mirror of https://github.com/Kkevsterrr/geneva
Added amplification plugin
This commit is contained in:
Kkevsterrr
parent
9c256c790b
commit
e788720bf3
|
|
@ -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)
|
||||
Loading…
Reference in New Issue