From dd13dc390352946bb2fc94938ce27895e1b9c6bf Mon Sep 17 00:00:00 2001
From: George Hughey <ghughey@umd.edu>
Date: Mon, 9 Mar 2020 21:16:22 -0700
Subject: [PATCH] Fixed nits

---
 actions/utils.py   |  2 +-
 engine.py          | 25 ++++++++++++-------------
 tests/test_tree.py |  2 +-
 3 files changed, 14 insertions(+), 15 deletions(-)

diff --git a/actions/utils.py b/actions/utils.py
index a795442..97fc447 100644
--- a/actions/utils.py
+++ b/actions/utils.py
@@ -183,7 +183,7 @@ def get_interface():
     """
     if os.name == 'nt':
         # Windows code
-        return # TODO: Fix this 
+        return
     else:
         ifaces = netifaces.interfaces()
         for iface in ifaces:
diff --git a/engine.py b/engine.py
index 9a8e697..cdfd17f 100644
--- a/engine.py
+++ b/engine.py
@@ -28,10 +28,7 @@ import actions.utils
 
 BASEPATH = os.path.dirname(os.path.abspath(__file__))
 
-if os.name == 'nt':
-    WINDOWS = True
-else:
-    WINDOWS = False
+WINDOWS = (os.name == 'nt')
 
 if WINDOWS:
     import pydivert
@@ -41,21 +38,21 @@ else:
 
 from abc import ABC, abstractmethod
 
-def Engine(server_port, string_strategy, environment_id=None, output_directory="trials", log_level="info"):
+def Engine(server_port, string_strategy, **kwargs):
     # Factory function to dynamically choose which engine to use.
     # Users should initialize an Engine using this.
     if WINDOWS:
         eng = WindowsEngine(server_port, 
                     string_strategy,
-                    environment_id=environment_id, 
-                    output_directory=output_directory, 
-                    log_level=log_level)
+                    environment_id=kwargs.get("environment_id", None), 
+                    output_directory=kwargs.get("output_directory", "trials"), 
+                    log_level=kwargs.get("log_level", "info"))
     else:
         eng = LinuxEngine(server_port, 
                     string_strategy,
-                    environment_id=environment_id, 
-                    output_directory=output_directory, 
-                    log_level=log_level)
+                    environment_id=kwargs.get("environment_id", None), 
+                    output_directory=kwargs.get("output_directory", "trials"), 
+                    log_level=kwargs.get("log_level", "info"))
 
     return eng
 
@@ -109,6 +106,7 @@ class GenericEngine(ABC):
         """
         self.shutdown()
 
+
 class WindowsEngine(GenericEngine):
     def __init__(self, server_port, string_strategy, environment_id=None, output_directory="trials", log_level="info"):
         super().__init__(server_port, string_strategy, environment_id=environment_id, output_directory=output_directory, log_level=log_level)
@@ -129,7 +127,8 @@ class WindowsEngine(GenericEngine):
         self.logger.debug("Initializing Divert")
 
         self.divert = pydivert.WinDivert("tcp.DstPort == %d || tcp.SrcPort == %d || udp.DstPort == %d || udp.SrcPort == %d"  \
-            % (int(self.server_port), int(self.server_port), int(self.server_port), int(self.server_port)))self.divert.open()
+            % (int(self.server_port), int(self.server_port), int(self.server_port), int(self.server_port)))
+        self.divert.open()
         self.divert_thread = threading.Thread(target=self.run_divert)
         self.divert_thread.start()
 
@@ -499,7 +498,7 @@ def main(args):
             strategy = LIBRARY[int(args["strategy_index"])][0]
         else:
             # Default to first strategy
-            strategy = LIBRARY[0][0]
+            strategy = LIBRARY[6][0]
         eng = Engine(args["server_port"],
                         strategy,
                         environment_id=args.get("environment_id"),
diff --git a/tests/test_tree.py b/tests/test_tree.py
index c69abbd..ca80ae6 100644
--- a/tests/test_tree.py
+++ b/tests/test_tree.py
@@ -131,7 +131,7 @@ def test_pretty_print():
     """
     Print complex tree, although difficult to test
     """
-    if platform.system() == "Windows":
+    if os.name == 'nt':
         return
     t = actions.trigger.Trigger("field", "flags", "TCP")
     a = actions.tree.ActionTree("out", trigger=t)