Added amplification plugin

plugins/amplification/plugin.py

@@ -0,0 +1,207 @@
+"""
+Amplification Plugin driver
+
+Overrides the default evaluator plugin handling so we can optimize testing many strategies at once,
+since we will not use the engine.
+"""
+
+import argparse
+import calendar
+import copy
+import logging
+import os
+import random
+import socket
+import sys
+import tempfile
+import time
+import traceback
+import tqdm
+import urllib.request
+
+import requests
+from scapy.all import *
+
+import actions.utils
+from plugins.plugin import Plugin
+
+BASEPATH = os.path.dirname(os.path.abspath(__file__))
+PROJECT_ROOT = os.path.dirname(os.path.dirname(BASEPATH))
+
+
+def get_open_sport(strategy_ports, logger):
+    """
+    Returns a source port that is not currently being used.
+    """
+    while True:
+        # Pick a port somewhere between 10,000 and 60,000
+        sport = random.randint(10000, 60000)
+        # If the source port has already been used, try to find a different one
+        if sport in strategy_ports:
+            continue
+
+        # Bind TCP socket
+        try:
+            with socket.socket() as sock:
+                # If we can bind, nothing is listening
+                sock.bind(('', sport))
+                break
+        except OSError:
+            logger.debug("Port %d is in use, picking another" % sport)
+            continue
+    logger.debug("Using source port %d" % sport)
+    return sport
+
+
+class AmplificationPluginRunner(Plugin):
+    """
+    Defines the amplification plugin runner.
+    """
+    name = "amplification"
+    override_evaluation = True
+
+    def __init__(self, args):
+        """
+        Marks this plugin as enabled
+        """
+        self.enabled = True
+        self.logger = None
+        self.strategy_ports = {}
+        self.responses = {}
+        self.sent_sizes = {}
+        self.disregard_empty = False
+
+    def handle_packet(self, packet):
+        """
+        Called by scapy when a matching inbound packet is seen.
+        """
+        strategy_port = packet["TCP"].dport
+        # If not to any strategy, not from us
+        if strategy_port not in self.strategy_ports:
+            return
+
+        if not packet.haslayer("TCP"):
+            return
+
+        if self.disregard_empty and not packet["TCP"].payload:
+            return
+
+        self.logger.debug("[%s] Received packet (%d): %s / %s", self.strategy_ports[strategy_port].environment_id, len(bytes(packet)), packet.summary(), packet["TCP"].payload)
+
+        if strategy_port not in self.responses:
+            self.responses[strategy_port] = []
+
+        self.responses[strategy_port].append(packet)
+
+    def evaluate(self, args, evaluator, population, logger):
+        """
+        Runs the plugins
+        """
+        self.logger = logger
+        self.disregard_empty = args["disregard_empty"]
+        # Clear the responses for the start of this generation
+        self.responses.clear()
+
+        dport = int(args.get("port", 7))
+        logger.debug("Using port %d" % dport)
+        site = args["site"]
+        dst = args["dst"]
+
+        payload = 'GET / HTTP/1.1\r\nHost: %s\r\n\r\n' % site
+        payload = payload.encode()
+
+        # Create a sniffer
+        logger.debug("Starting sniffer")
+        sniffer = AsyncSniffer(filter="tcp and src port %d" % dport, prn=self.handle_packet, store=False)
+        sniffer.start()
+
+        # Maps source ports to strategies
+        self.strategy_ports = {}
+        self.sent_sizes = {}
+        for ind in tqdm.tqdm(population, leave=False, disable=(actions.utils.CONSOLE_LOG_LEVEL == "debug")):
+            sport = get_open_sport(self.strategy_ports, logger)
+            # Reserve this source port for this strategy
+            self.strategy_ports[sport] = ind
+            seq = int(RandInt())
+            ack = int(RandInt())
+            packets = [
+                IP(dst=dst)/TCP(sport=sport, dport=dport, flags="S", ack=0, seq=seq),
+                IP(dst=dst)/TCP(sport=sport, dport=dport, flags="A", ack=ack, seq=seq+1),
+                IP(dst=dst)/TCP(sport=sport, dport=dport, flags="PA", ack=ack, seq=seq+1)/Raw(payload)
+            ]
+            packets = [actions.packet.Packet(packet) for packet in packets]
+
+            packets_to_send = []
+            try:
+                for packet in packets:
+                    # Run the strategy on the packet
+                    packets_to_send += ind.act_on_packet(packet, logger)
+            except Exception:
+                logger.exception("Error running strategy")
+                ind.fitness = -1000
+                continue
+
+            # If the strategy sends no packets, punish and continue
+            if not packets_to_send:
+                ind.fitness = -1000
+                continue
+
+            for packet in packets_to_send:
+                # Record the size we're about to send
+                if sport not in self.sent_sizes:
+                    self.sent_sizes[sport] = 0
+                self.sent_sizes[sport] += len(bytes(packet.packet))
+            logger.debug("About to send %d bytes" % self.sent_sizes[sport])
+
+            for packet in packets_to_send:
+                if packet.sleep:
+                    time.sleep(packet.sleep)
+
+                self.logger.debug("Sending packet (%d) %s", len(bytes(packet)), str(packet))
+                # Send the packet
+                send(packet.packet, verbose=False)
+
+            # Sleep the requested milliseconds between generations
+            time.sleep(args["sleep"]/1000)
+
+        logger.info("Sleeping %d cooldown seconds to wait for packets to come in" % args["cooldown"])
+        time.sleep(args["cooldown"])
+
+        logger.debug("Stopping sniffer")
+        sniffer.stop()
+
+        # Zero out the fitnesses for strategies that do not get responses
+        for port in self.strategy_ports:
+            ind = self.strategy_ports[port]
+            if ind.fitness != -1000:
+                ind.fitness = 0
+                if port in self.responses:
+                    for response in self.responses[port]:
+                        ind.fitness += len(bytes(response))
+
+                    ind.fitness = round(ind.fitness / self.sent_sizes[port], 3)
+                    self.logger.debug("[%s] Fitness %s: %s" % (ind.environment_id, ind.fitness, str(ind)))
+
+                ind.fitness = actions.utils.punish_unused(ind.fitness, logger, ind)
+            logger.debug("[%s] Fitness: %s: %s" % (ind.environment_id, ind.fitness, str(ind)))
+
+        self.strategy_ports.clear()
+        self.responses.clear()
+        return population
+
+    @staticmethod
+    def get_args(command):
+        """
+        Defines required args for this plugin
+        """
+        parser = argparse.ArgumentParser(description='Amplification plugin runner', allow_abbrev=False)
+        parser.add_argument('--output-directory', action='store', help="Where to output results")
+        parser.add_argument('--sleep', action='store', type=int, default=500, help='milliseconds to sleep between each strategy')
+        parser.add_argument('--port', action='store', type=int, default=7, help='port to use')
+        parser.add_argument('--dst', action='store', help='IP to use')
+        parser.add_argument('--runs', action='store', help='Runs per strategy')
+        parser.add_argument('--site', action='store', default="pornhub.com", help='Site to include in the HTTP GET request')
+        parser.add_argument('--disregard-empty', action='store_true', help='Disregard packets without payloads (RSTs)')
+        parser.add_argument('--cooldown', action='store', type=int, default=8, help='amount of time after the last packet is sent to collect packets')
+        args, _ = parser.parse_known_args(command)
+        return vars(args)