STRLCPY/geneva

Added amplification plugin
Kkevsterrr committed 3 years ago

e788720b

1 parent 9c256c79

Total 1 files

■ ■ ■ ■ ■ ■

plugins/amplification/plugin.py

1	+	"""
2	+	Amplification Plugin driver
3	+
4	+	Overrides the default evaluator plugin handling so we can optimize testing many strategies at once,
5	+	since we will not use the engine.
6	+	"""
7	+
8	+	import argparse
9	+	import calendar
10	+	import copy
11	+	import logging
12	+	import os
13	+	import random
14	+	import socket
15	+	import sys
16	+	import tempfile
17	+	import time
18	+	import traceback
19	+	import tqdm
20	+	import urllib.request
21	+
22	+	import requests
23	+	from scapy.all import *
24	+
25	+	import actions.utils
26	+	from plugins.plugin import Plugin
27	+
28	+	BASEPATH = os.path.dirname(os.path.abspath(__file__))
29	+	PROJECT_ROOT = os.path.dirname(os.path.dirname(BASEPATH))
30	+
31	+
32	+	def get_open_sport(strategy_ports, logger):
33	+	"""
34	+	Returns a source port that is not currently being used.
35	+	"""
36	+	while True:
37	+	# Pick a port somewhere between 10,000 and 60,000
38	+	sport = random.randint(10000, 60000)
39	+	# If the source port has already been used, try to find a different one
40	+	if sport in strategy_ports:
41	+	continue
42	+
43	+	# Bind TCP socket
44	+	try:
45	+	with socket.socket() as sock:
46	+	# If we can bind, nothing is listening
47	+	sock.bind(('', sport))
48	+	break
49	+	except OSError:
50	+	logger.debug("Port %d is in use, picking another" % sport)
51	+	continue
52	+	logger.debug("Using source port %d" % sport)
53	+	return sport
54	+
55	+
56	+	class AmplificationPluginRunner(Plugin):
57	+	"""
58	+	Defines the amplification plugin runner.
59	+	"""
60	+	name = "amplification"
61	+	override_evaluation = True
62	+
63	+	def __init__(self, args):
64	+	"""
65	+	Marks this plugin as enabled
66	+	"""
67	+	self.enabled = True
68	+	self.logger = None
69	+	self.strategy_ports = {}
70	+	self.responses = {}
71	+	self.sent_sizes = {}
72	+	self.disregard_empty = False
73	+
74	+	def handle_packet(self, packet):
75	+	"""
76	+	Called by scapy when a matching inbound packet is seen.
77	+	"""
78	+	strategy_port = packet["TCP"].dport
79	+	# If not to any strategy, not from us
80	+	if strategy_port not in self.strategy_ports:
81	+	return
82	+
83	+	if not packet.haslayer("TCP"):
84	+	return
85	+
86	+	if self.disregard_empty and not packet["TCP"].payload:
87	+	return
88	+
89	+	self.logger.debug("[%s] Received packet (%d): %s / %s", self.strategy_ports[strategy_port].environment_id, len(bytes(packet)), packet.summary(), packet["TCP"].payload)
90	+
91	+	if strategy_port not in self.responses:
92	+	self.responses[strategy_port] = []
93	+
94	+	self.responses[strategy_port].append(packet)
95	+
96	+	def evaluate(self, args, evaluator, population, logger):
97	+	"""
98	+	Runs the plugins
99	+	"""
100	+	self.logger = logger
101	+	self.disregard_empty = args["disregard_empty"]
102	+	# Clear the responses for the start of this generation
103	+	self.responses.clear()
104	+
105	+	dport = int(args.get("port", 7))
106	+	logger.debug("Using port %d" % dport)
107	+	site = args["site"]
108	+	dst = args["dst"]
109	+
110	+	payload = 'GET / HTTP/1.1\r\nHost: %s\r\n\r\n' % site
111	+	payload = payload.encode()
112	+
113	+	# Create a sniffer
114	+	logger.debug("Starting sniffer")
115	+	sniffer = AsyncSniffer(filter="tcp and src port %d" % dport, prn=self.handle_packet, store=False)
116	+	sniffer.start()
117	+
118	+	# Maps source ports to strategies
119	+	self.strategy_ports = {}
120	+	self.sent_sizes = {}
121	+	for ind in tqdm.tqdm(population, leave=False, disable=(actions.utils.CONSOLE_LOG_LEVEL == "debug")):
122	+	sport = get_open_sport(self.strategy_ports, logger)
123	+	# Reserve this source port for this strategy
124	+	self.strategy_ports[sport] = ind
125	+	seq = int(RandInt())
126	+	ack = int(RandInt())
127	+	packets = [
128	+	IP(dst=dst)/TCP(sport=sport, dport=dport, flags="S", ack=0, seq=seq),
129	+	IP(dst=dst)/TCP(sport=sport, dport=dport, flags="A", ack=ack, seq=seq+1),
130	+	IP(dst=dst)/TCP(sport=sport, dport=dport, flags="PA", ack=ack, seq=seq+1)/Raw(payload)
131	+	]
132	+	packets = [actions.packet.Packet(packet) for packet in packets]
133	+
134	+	packets_to_send = []
135	+	try:
136	+	for packet in packets:
137	+	# Run the strategy on the packet
138	+	packets_to_send += ind.act_on_packet(packet, logger)
139	+	except Exception:
140	+	logger.exception("Error running strategy")
141	+	ind.fitness = -1000
142	+	continue
143	+
144	+	# If the strategy sends no packets, punish and continue
145	+	if not packets_to_send:
146	+	ind.fitness = -1000
147	+	continue
148	+
149	+	for packet in packets_to_send:
150	+	# Record the size we're about to send
151	+	if sport not in self.sent_sizes:
152	+	self.sent_sizes[sport] = 0
153	+	self.sent_sizes[sport] += len(bytes(packet.packet))
154	+	logger.debug("About to send %d bytes" % self.sent_sizes[sport])
155	+
156	+	for packet in packets_to_send:
157	+	if packet.sleep:
158	+	time.sleep(packet.sleep)
159	+
160	+	self.logger.debug("Sending packet (%d) %s", len(bytes(packet)), str(packet))
161	+	# Send the packet
162	+	send(packet.packet, verbose=False)
163	+
164	+	# Sleep the requested milliseconds between generations
165	+	time.sleep(args["sleep"]/1000)
166	+
167	+	logger.info("Sleeping %d cooldown seconds to wait for packets to come in" % args["cooldown"])
168	+	time.sleep(args["cooldown"])
169	+
170	+	logger.debug("Stopping sniffer")
171	+	sniffer.stop()
172	+
173	+	# Zero out the fitnesses for strategies that do not get responses
174	+	for port in self.strategy_ports:
175	+	ind = self.strategy_ports[port]
176	+	if ind.fitness != -1000:
177	+	ind.fitness = 0
178	+	if port in self.responses:
179	+	for response in self.responses[port]:
180	+	ind.fitness += len(bytes(response))
181	+
182	+	ind.fitness = round(ind.fitness / self.sent_sizes[port], 3)
183	+	self.logger.debug("[%s] Fitness %s: %s" % (ind.environment_id, ind.fitness, str(ind)))
184	+
185	+	ind.fitness = actions.utils.punish_unused(ind.fitness, logger, ind)
186	+	logger.debug("[%s] Fitness: %s: %s" % (ind.environment_id, ind.fitness, str(ind)))
187	+
188	+	self.strategy_ports.clear()
189	+	self.responses.clear()
190	+	return population
191	+
192	+	@staticmethod
193	+	def get_args(command):
194	+	"""
195	+	Defines required args for this plugin
196	+	"""
197	+	parser = argparse.ArgumentParser(description='Amplification plugin runner', allow_abbrev=False)
198	+	parser.add_argument('--output-directory', action='store', help="Where to output results")
199	+	parser.add_argument('--sleep', action='store', type=int, default=500, help='milliseconds to sleep between each strategy')
200	+	parser.add_argument('--port', action='store', type=int, default=7, help='port to use')
201	+	parser.add_argument('--dst', action='store', help='IP to use')
202	+	parser.add_argument('--runs', action='store', help='Runs per strategy')
203	+	parser.add_argument('--site', action='store', default="pornhub.com", help='Site to include in the HTTP GET request')
204	+	parser.add_argument('--disregard-empty', action='store_true', help='Disregard packets without payloads (RSTs)')
205	+	parser.add_argument('--cooldown', action='store', type=int, default=8, help='amount of time after the last packet is sent to collect packets')
206	+	args, _ = parser.parse_known_args(command)
207	+	return vars(args)
208	+

Added amplification plugin