Uses Context Manager in main method

This commit is contained in:
VwCSXg 2022-07-27 18:54:53 +02:00
parent 2fc5b2f2f4
commit e495652a4a
1 changed files with 18 additions and 21 deletions

View File

@ -443,29 +443,26 @@ def main(args):
"""
Kicks off the engine with the given arguments.
"""
try:
nat_config = {}
if args.get("sender_ip") and args.get("routing_ip") and args.get("forward_ip"):
nat_config = {"sender_ip" : args["sender_ip"],
"routing_ip" : args["routing_ip"],
"forward_ip" : args["forward_ip"]}
nat_config = {}
if args.get("sender_ip") and args.get("routing_ip") and args.get("forward_ip"):
nat_config = {"sender_ip": args["sender_ip"],
"routing_ip": args["routing_ip"],
"forward_ip": args["forward_ip"]}
eng = Engine(args["server_port"],
args["strategy"],
environment_id=args["environment_id"],
server_side=args["server_side"],
output_directory=args["output_directory"],
forwarder=nat_config,
log_level=args["log"],
in_queue_num=args["in_queue_num"],
out_queue_num=args["out_queue_num"],
save_seen_packets=args["no_save_packets"],
demo_mode=args["demo_mode"])
with Engine(args["server_port"],
args["strategy"],
environment_id=args["environment_id"],
server_side=args["server_side"],
output_directory=args["output_directory"],
forwarder=nat_config,
log_level=args["log"],
in_queue_num=args["in_queue_num"],
out_queue_num=args["out_queue_num"],
save_seen_packets=args["no_save_packets"],
demo_mode=args["demo_mode"]) as eng:
eng.initialize_nfqueue()
while True:
time.sleep(0.5)
finally:
eng.shutdown_nfqueue()
threading.Event().wait() # Wait forever
if __name__ == "__main__":