Browse Source

Added amplification plugin

pull/30/head
Kkevsterrr 10 months ago
parent
commit
e788720bf3
  1. 207
      plugins/amplification/plugin.py

207
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)
Loading…
Cancel
Save