/client.py

import websockets
import logging
import json
import asyncio

logging.SPAM = 5


def domainFromName(name):
    for part in name.split('.'):
        if part.startswith('as'):
            return part[2:]


class SarnetUIConnection():
    def __init__(self, uri, ssl):
        self.ws = None
        self.uri = uri
        self.ctx = ssl
        self.identifier = {'version': 'sc17-0.1',
                           'agent': 'python3/websockets',
                           'screen': '0x0@1',
                           'url': 'cli'}

    async def __aenter__(self):
        self._conn = websockets.connect(self.uri, ssl=self.ctx)
        self.ws = await self._conn.__aenter__()
        await self.oldsend('identify', data=self.identifier)
        return self

    async def __aexit__(self, *args, **kwargs):
        await self._conn.__aexit__(*args, **kwargs)

    async def recieve(self):
        try:
            msg = await asyncio.wait_for(self.ws.recv(), timeout=20)
        except asyncio.TimeoutError:
            # No data in 20 seconds, check the connection.
            try:
                pong_waiter = await self.ws.ping()
                await asyncio.wait_for(pong_waiter, timeout=10)
            except asyncio.TimeoutError:
                # No response to ping in 10 seconds, disconnect.
                raise asyncio.TimeoutError
        else:
            logging.log(logging.SPAM, 'recv: %s', msg)
            parsed = json.loads(msg)
        return parsed

    async def _send(self, msg):
        logging.debug('sent: %s', json.dumps(msg))
        await self.ws.send(json.dumps(msg))

    async def recv_handler(self):
        while True:
            parsed = await self.recieve()
            self.handle_msg(parsed)

    async def newsend(self, kind, **kwargs):
        msg = {'@': kind}
        if msg:
            msg.update(kwargs)
        await self._send(msg)

    async def oldsend(self, kind, data=None, **kwargs):
        msg = {'kind': kind}
        if data is not None or kwargs:
            if data:
                data.update(kwargs)
            else:
                data = kwargs
            msg['data'] = data
        await self._send(msg)

    def handle_msg(self, msg):
        if 'm' in msg and 'cmd/host-netctl/containers/va' in msg['m']:
            logging.info('got: %s', msg)


class SarnetUIClient(SarnetUIConnection):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.topology = None
        self.metadata = {}
        self.nodes = []
        self.possible_attackers = []
        self.possible_victims = []
        self.scenario = None
        self.msg_callback = {}
        self.learning = set()
        self.thresholds = {}

    def handle_msg(self, msg):
        if msg['@'] == 't':
            self.setTopology(msg['t'])
        elif msg['@'] == 'm':  # metadata updates
            self.setNodeMetadata(msg['i'], msg['m'])
        elif msg['@'] == 'M':  # node metadata
            for node_name in msg['M'].keys():
                self.nodes.append(node_name)
                self.setNodeMetadata(node_name, msg['M'][node_name], True)
        elif msg['@'] == 'f':
            # flows
            pass
        elif msg['@'] == 'B':
            if msg['B']['kind'] == 'reload-topo':
                logging.warning('reload-topo requested')
                # reload topo
                pass
        else:
            logging.debug('client: got unknown msg: %', msg)

    async def set_collab(self, level):
        if level in ['local', 'upstream', 'alliance'] or level.isdigit():
            if level == '0':
                level = 'local'
            await self._send({"@": "pub",
                              "k": "cmd/sarnet/collaboration",
                              "v": level,
                              "r": True})
            logging.debug('client: set level to %s', level)
            return level
        return None

    async def set_defense(self, level):
        if level in ['filter', 'rateup', 'honeypot']:
            await self._send({"@": "pub",
                              "k": "cmd/sarnet/defense",
                              "v": level,
                              "r": True})
            logging.debug('client: set defense to %s', level)
            return level
        return None

    def setTopology(self, topology):
        self.topology = topology
        self.topology['flow_owner'] = {}
        self.topology['stat_owner'] = {}
        flowMacs = {}

        for k, v in topology['nodes'].items():
            if '/service/' in k:
                self.possible_victims.append(k)
            if '/client/' in k:
                self.possible_attackers.append(k)

            a = self.metadata.setdefault(v['name'], {})
            a['iflink'] = {}
            a['ifmac'] = {}

            for iface in v['ifaces']:
                if 'links' not in iface or len(iface['links']) == 0:
                    continue
                linkName = iface['links'][0]
                try:
                    link = topology['links'][linkName]
                except KeyError:
                    link = None

                if not link:
                    continue

                if 'ifname' in iface:
                    a.iflink[linkName] = iface['ifname']

                if iface['mac']:
                    mac = iface['mac'].lower()

                    a['ifmac'][mac] = linkName

                    flowMacs[mac.replace(':', '')] = {
                        'left': link['source'],
                        'right': link['target'],
                        'owner': v['name'],
                        'link': linkName}

    def setNodeNetstats(self, name, kv):
        # don't seem to need it
        pass

    def trackValue(self, kv, k, v):
        if v:
            kv[k] = v
        else:
            kv.pop(k, None)

    def setNodeMetadata(self, name, kv, reset=False):
        meta = {}

        if not kv:
            self.metadata.pop(name, None)
            meta = {}
            kv = {}
            reset = True
        else:
            meta = self.metadata.setdefault(name, {})

        if reset:
            meta['status'] = {}
            meta['linkStatus'] = {}

        # 'undefined' may be Null
        if 'status' not in meta.keys() or meta['status'] is None:
            meta['status'] = {}
        if 'linkStatus' not in meta.keys() or meta['linkStatus'] is None:
            meta['linkStatus'] = {}
        if 'vm/ns' in kv.keys():
            self.setNodeNetstats(name, kv['vm/ns'])

        for k, v in kv.items():
            if k.startswith('if/'):
                iface = json.loads(v)
                # idgi
                if iface['mac'] and iface['iface']:
                    mac = iface['mac'].lower()
                    meta.setdefault('iflink', {})[iface['iface']] = mac
                    if 'ifmac' in meta:
                        meta.setdefault('iflink', {})[meta['ifmac'][mac]] \
                            = iface['iface']
            elif k.startswith('sarnet/a/'):
                self.trackValue(meta['status'], k, v or False)
            elif k.startswith('sarnet/o/'):
                self.trackValue(meta['status'], k, v != '{"state": true}')
            elif k == 'sarnet/log':
                # node log
                pass
            elif k == 'sarnet/learn_mode':
                self.trackValue(meta['status'], k, v is True)
                pass
            elif k == 'cmd/host-netctl/iface-filter':
                pass
            elif k == 'cmd/host-netctl/sdn-redirect':
                pass
            elif k == 'cmd/host-netctl/nfv':
                self.trackValue(meta['status'], k, v != '[]')
            elif k == 'vm/containers':
                atk = False
                cname = v
                for cname in json.loads(v):
                    if cname.startswith('va-'):
                        atk = True
                self.trackValue(meta['status'], '$attacker', atk)
            elif k == 'vm/egress-policer':
                meta['linkStatus']['policer'] = json.loads(v)
                pass

            for pattern, callback in self.msg_callback.items():
                if callback and k.startswith(pattern):
                    callback(name, k, v)

    def set_callback(self, pattern, func):
        self.msg_callback[pattern] = func

    async def stop(self):
        self.scenario = None
        await self._send({'@': 'attack', 'attack': 'stop'})

    async def set_behaviour(self, data):
        logging.info('client: broadcasting domain behaviour parameters')
        await self._send({'@': 'behaviour', 'data': data})

    def N(self, name):
        def url_to_ident(ident):
            if ident.startswith('http://geni-orca.renci.org/owl/'):
                return ident[31:]
            return ident

        n = self.topology['nodes'][name]
        if n and 'unit_url' in n.keys():
            return url_to_ident(n['unit_url'])
        return name

    async def attack(self, att_type, attackers, victim, rate=.8):

        def ddos_rate(attackers, rate):
            max_link_capacity = 100
            max_rate = max_link_capacity * len(attackers)
            if rate == 'high':
                rate = .9
            elif rate == 'med':
                rate = .6
            elif rate == 'low':
                rate = .3
            return max_rate * rate

        def pwd_rate(rate):
            if rate == 'high':
                rate = .9
            elif rate == 'med':
                rate = .6
            elif rate == 'low':
                rate = .3
            return int(10 * rate)

        async def attack_victim(att_type, attackers, victim, rate=.8):
            victimDomain = domainFromName(victim)

            nrate = ddos_rate(attackers, rate)
            nworkers = pwd_rate(rate)
            logging.info('ddos rate = %s, pwd workers = %s, victim=%s', str(nrate),
                         str(int(nworkers)), victim)
            d = {'@': 'attack', 'attack': 'cs', 'cs': {}}
            for a in attackers:
                localDomain = domainFromName(a)
                env = {}
                docker_img = None
                if att_type == 'pw':
                    docker_img = 'vnet-client'
                    env['CONTAINER_ARGS'] = '-mode,pw,-workers,' + str(int(nworkers)) + ',172.30.' + victimDomain + '.64'
                elif att_type == 'reflect':
                    docker_img = 'vnet-ddos-udp'
                    logging.debug('client: reflect select:  %s %s',
                                  localDomain, int(localDomain) < 13)
                    env['SRC'] = '172.30.' + victimDomain + '.250'
                    # TODO: need a reflector select mechanism
                    env['DST'] = '172.30.' + ('15' if int(localDomain) < 13
                                              else '11') + '.63'
                    env['DST'] = '172.30.' + '11' + '.63'
                    env['RATE'] = str((nrate / (3 * len(attackers)) or 0))
                    env['SIZE'] = '500'
                else:
                    docker_img = 'vnet-ddos-udp'
                    env['SRC'] = '172.30.' + localDomain + '.250'
                    env['DST'] = '172.30.' + victimDomain + '.250'
                    env['RATE'] = str(min(95, nrate / len(attackers) or 0))
                    env['SIZE'] = '1460'
                d['cs'][self.N(a)] = [{'name': 'va-attack',
                                       'img': docker_img,
                                       'env': env}]

            for a in self.possible_attackers:
                if a not in attackers:
                    d['cs'][self.N(a)] = []

            await self._send(d)

        if isinstance(victim, list):
            n_victims = len(victim)
            n_attackers = len(attackers)
            n_att_vic = int(n_attackers / n_victims)
            if n_attackers < n_victims:
                raise ValueError("num_attackers less than num_victims")
                # can improve it with someting modulo etc.
                # though better to error because it's probably
                # unintentional if this occurs


            for i, v in enumerate(victim):
                idx = n_att_vic * i
                att = attackers[idx:idx + n_att_vic]
                await attack_victim(att_type, att, v, rate)
        else:
            await attack_victim(att_type, attackers, victim, rate)

    async def start(self, scenario):
        if self.scenario:
            raise RuntimeError('already running scenario')
        self.scenario = scenario
        if scenario.attack not in ['pw', 'reflect', 'ddos']:
            return
        if scenario.level:
            await self.set_collab(scenario.level)
        if scenario.defense:
            await self.set_defense(scenario.defense)
        if scenario.behaviour:
            await self.set_behaviour(scenario.behaviour)
        if scenario.rate:
            rate = scenario.rate
        else:
            rate = .3
        await self.attack(scenario.attack, scenario.attackers,
                          scenario.victims, rate)

    async def clean(self, ignore=[]):
        logging.info('cleanup: start')
        clean = False
        while not clean:
            need_cleaning = []
            await asyncio.sleep(1)
            logging.debug('cleanup: cleaning...')
            for name in self.topology['nodes']:
                m = self.metadata[name]
                if not m:
                    continue
                ignore_obs = any([i for i in ignore if i in name])
                if 'status' not in m.keys():
                    continue
                for k in m['status'].keys():
                    if ignore_obs and k.startswith('sarnet/o'):
                        logging.debug('cleanup: ignoring %s', name)
                        continue
                    need_cleaning.append('{}/{}'.format(name, k))
            if len(need_cleaning) == 0:
                logging.info('cleanup: done')
                clean = True
            else:
                logging.debug('cleanup: need to clean up: %s', need_cleaning)
        return clean

    async def learn(self):
        """ Relearns sarnet's baselines (exposes thresholds) """
        await self._send({"@": "pub",
                          "k": "cmd/sarnet/learn",
                          "v": True,
                          "r": True})
        logging.info('learning: start')
        self.set_callback('sarnet/thresh', self.handle_thresh)
        self.set_callback('sarnet/learn_mode', self.handle_learn)
        await asyncio.sleep(.5)
        while len(self.learning) > 0:
            await asyncio.sleep(1)
        logging.info('learning: ended')

    def handle_thresh(self, name, key, value):
        key = key.split('/')[2]
        self.thresholds.setdefault(name, {})[key.split('@')[0]] = float(value)

    def handle_learn(self, name, key, value):
        if value == 'true':
            self.learning.add(name)
        elif value == 'false':
            if name in self.learning:
                self.learning.remove(name)
        else:
            RuntimeError('expected true or false')