From a5adc94598d1810205a84189224c4cdc2d6d5ed0 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Wed, 21 Oct 2020 13:13:43 +0000 Subject: [PATCH] Fixed plumbing for plugins to specify interface to the engine --- plugins/plugin_client.py | 4 +++- plugins/plugin_server.py | 4 +++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/plugins/plugin_client.py b/plugins/plugin_client.py index 9c1ff46..4a6aaca 100644 --- a/plugins/plugin_client.py +++ b/plugins/plugin_client.py @@ -41,6 +41,7 @@ class ClientPlugin(Plugin): choices=("debug", "info", "warning", "critical", "error"), help="Sets the log level") parser.add_argument('--port', action='store', type=int, help='port to run this server on') + parser.add_argument('--interface', action='store', help='interface to run the engine on for client-side training') parser.add_argument('--wait-for-censor', action='store_true', help='send control packets to the censor to get startup confirmation') @@ -60,6 +61,7 @@ class ClientPlugin(Plugin): eid = args.get("environment_id") use_engine = not args.get("no_engine") port = args.get("port") + interface = args.get("interface") server_side = args.get("server_side") assert port, "Need to specify a port in order to launch a sniffer" @@ -69,7 +71,7 @@ class ClientPlugin(Plugin): with actions.sniffer.Sniffer(pcap_filename, port, logger) as sniff: # Conditionally initialize the engine - with engine.Engine(port, args.get("strategy"), server_side=False, environment_id=eid, output_directory=output_path, log_level=args.get("log", "info"), enabled=use_engine) as eng: + with engine.Engine(port, args.get("strategy"), server_side=False, environment_id=eid, output_directory=output_path, log_level=args.get("log", "info"), interface=interface, enabled=use_engine) as eng: # Wait for the censor to start up, if one is running if args.get("wait_for_censor"): self.wait_for_censor(args.get("server"), port, eid, output_path) diff --git a/plugins/plugin_server.py b/plugins/plugin_server.py index 0cd2e86..b451b89 100644 --- a/plugins/plugin_server.py +++ b/plugins/plugin_server.py @@ -45,6 +45,7 @@ class ServerPlugin(Plugin): choices=("debug", "info", "warning", "critical", "error"), help="Sets the log level") parser.add_argument('--port', action='store', type=int, help='port to run this server on') + parser.add_argument('--interface', action='store_true', help='interface to run the engine on for server-side training') parser.add_argument('--external-server', action='store_true', help="use an external server for testing.") @@ -70,6 +71,7 @@ class ServerPlugin(Plugin): eid = args["environment_id"] use_engine = not args.get("no_engine", False) port = args["port"] + interface = args.get("interface") server_side = args["server_side"] log_level = args["log"] strategy = args.get("strategy", "") @@ -94,7 +96,7 @@ class ServerPlugin(Plugin): self.sniffer = actions.sniffer.Sniffer(pcap_filename, int(port), logger).__enter__() # Conditionally initialize the engine - self.engine = engine.Engine(port, strategy, server_side=True, environment_id=eid, output_directory=output_path, log_level=args.get("log", "info"), enabled=use_engine, forwarder=forwarder).__enter__() + self.engine = engine.Engine(port, strategy, server_side=True, environment_id=eid, output_directory=output_path, log_level=args.get("log", "info"), enabled=use_engine, interface=interface, forwarder=forwarder).__enter__() # Run the plugin self.server_proc = multiprocessing.Process(target=self.start_thread, args=(args, logger))