wishful_simple_controller

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
wishful_controller_simple.py: First implementation of WiSHFUL controller

Usage:
   wishful_controller_simple.py [options] [-q | -v]

Options:
   --logfile name      Name of the logfile
   --config configFile Config file path

Example:
   ./wishful_simple_local_controller -v --config ./config.yaml 

Other options:
   -h, --help          show this help message and exit
   -q, --quiet         print less text
   -v, --verbose       print more text
   --version           show version and exit
"""

import sys
import datetime
import logging
import wishful_controller
import gevent
import yaml
from scapy.all import *
import wishful_upis as upis
from wishful_framework import TimeEvent, PktEvent, MovAvgFilter, PeakDetector, Match, Action, Permanance, PktMatch, FieldSelector


__author__ = "Piotr Gawlowicz, Anatolij Zubow"
__copyright__ = "Copyright (c) 2015, Technische Universitat Berlin"
__version__ = "0.1.0"
__email__ = "{gawlowicz, zubow}@tkn.tu-berlin.de"


log = logging.getLogger('wishful_agent.main')
controller = wishful_controller.Controller()
nodes = []


@controller.new_node_callback()
def new_node(node):
    nodes.append(node)
    print("New node appeared:")
    print(node)


@controller.node_exit_callback()
def node_exit(node, reason):
    if node in nodes:
        nodes.remove(node);
    print("NodeExit : NodeName : {} Reason : {}".format(node.name, reason))


@controller.set_default_callback()
def default_callback(group, node, cmd, data):
    print("DEFAULT CALLBACK : Group: {}, NodeName: {}, Cmd: {}, Returns: {}".format(group, node.name, cmd, data))


def upi_based_rule_matched(group, node, ruleId, data):
    print("{} UPI based Rule in node: {} with id: {} was matched; Data: {}".format(datetime.datetime.now(), node.name, ruleId, data)) 

def pkt_based_rule_matched(group, node, ruleId, data):
    print("{} Pkt Rule in node: {} with id: {} was matched; Data: {}".format(datetime.datetime.now(), node.name, ruleId, data)) 

def pkt_based_rule(group, node, ruleId, pkt):
    print("{} Pkt Rule in node: {} with id: {} was matched;".format(datetime.datetime.now(), node.name, ruleId)) 
    try:
      print(pkt.show())
    except:
      print(pkt)
    print("")


def main(args):
    log.debug(args)

    config_file_path = args['--config']
    config = None
    with open(config_file_path, 'r') as f:
        config = yaml.load(f)

    controller.load_config(config)
    controller.start()
   
    #wait for at leas one node
    while len(nodes) < 1:
        gevent.sleep(1)

    #control loop
    while True:
        print("Connected nodes", [str(node.name) for node in nodes])
        if nodes:
            #######################################################################################
            print("{} Get new generator".format(datetime.datetime.now()))
            myGenerator = controller.nodes(nodes[0]).generator.start(
                                        event=TimeEvent(func=upis.radio.get_noise, kwargs={}, interval=0.2, iface="wlan0"))

            print("{} New generator with ID: {}".format(datetime.datetime.now(), myGenerator.id))

            sampleNum = 5
            i = 0
            for sample in myGenerator:
                print(i, sample)
                i = i + 1
                if i == sampleNum:
                  break

            retVal = myGenerator.stop()
            print("{} Generator with ID: {} was {}".format(datetime.datetime.now(), myGenerator.id, retVal))
            gevent.sleep(2)

            #######################################################################################
            print("{} Get new generator".format(datetime.datetime.now()))
            myGenerator = controller.nodes(nodes[0]).generator.start(
                                        event=TimeEvent(func=upis.radio.get_rssi, kwargs={}, iface="wlan0"),
                                        filters=[MovAvgFilter(10)])

            print("{} New generator with ID: {}".format(datetime.datetime.now(), myGenerator.id))

            sampleNum = 5
            i = 0
            for sample in myGenerator:
                print(i, sample)
                i = i + 1
                if i == sampleNum:
                  break

            retVal = myGenerator.stop()
            print("{} Generator with ID: {} was {}".format(datetime.datetime.now(), myGenerator.id, retVal))
            gevent.sleep(2)

            #######################################################################################
            print("{} Get new generator".format(datetime.datetime.now()))
            myGenerator = controller.nodes(nodes[0]).generator.start(
                                        event=PktEvent(iface="eth0"),
                                        pktMatch=PktMatch("icmp"),
                                        selector=FieldSelector("IP.ttl"))

            print("{} New generator with ID: {}".format(datetime.datetime.now(), myGenerator.id))

            sampleNum = 5
            i = 0
            for sample in myGenerator:
                print(i, sample)
                i = i + 1
                if i == sampleNum:
                  break

            retVal = myGenerator.stop()
            print("{} Generator with ID: {} was {}".format(datetime.datetime.now(), myGenerator.id, retVal))
            gevent.sleep(2)

            #######################################################################################
            print("{} Get new generator".format(datetime.datetime.now()))
            myGenerator = controller.nodes(nodes[0]).generator.start(
                                        event=PktEvent(iface="eth0"),
                                        selector=FieldSelector("IP.dst"))

            print("{} New generator with ID: {}".format(datetime.datetime.now(), myGenerator.id))

            sampleNum = 5
            i = 0
            for sample in myGenerator:
                print(i, sample)
                i = i + 1
                if i == sampleNum:
                  break

            retVal = myGenerator.stop()
            print("{} Generator with ID: {} was {}".format(datetime.datetime.now(), myGenerator.id, retVal))
            gevent.sleep(2)

            #######################################################################################
            print("{} Get new generator".format(datetime.datetime.now()))
            myGenerator = controller.nodes(nodes[0]).generator.start(
                                        event=PktEvent(iface="eth0"))

            print("{} New generator with ID: {}".format(datetime.datetime.now(), myGenerator.id))

            sampleNum = 5
            i = 0
            for pkt in myGenerator:
                print(i, pkt.show())
                i = i + 1
                if i == sampleNum:
                  break

            retVal = myGenerator.stop()
            print("{} Generator with ID: {} was {}".format(datetime.datetime.now(), myGenerator.id, retVal))
            gevent.sleep(2)

            #######################################################################################



if __name__ == "__main__":
    try:
        from docopt import docopt
    except:
        print("""
        Please install docopt using:
            pip install docopt==0.6.1
        For more refer to:
        https://github.com/docopt/docopt
        """)
        raise

    args = docopt(__doc__, version=__version__)

    log_level = logging.INFO  # default
    if args['--verbose']:
        log_level = logging.DEBUG
    elif args['--quiet']:
        log_level = logging.ERROR

    logfile = None
    if args['--logfile']:
        logfile = args['--logfile']

    logging.basicConfig(filename=logfile, level=log_level,
        format='%(asctime)s - %(name)s.%(funcName)s() - %(levelname)s - %(message)s')

    try:
        main(args)
    except KeyboardInterrupt:
        log.debug("Controller exits")
    finally:
        log.debug("Exit")
        controller.stop()