mirror of https://github.com/Kkevsterrr/geneva
208 lines
7.4 KiB
Python
208 lines
7.4 KiB
Python
"""
|
|
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)
|