From 18a90adcbb5623030368847e75fff4e11612de71 Mon Sep 17 00:00:00 2001
From: Jim Carciofini <jim.carciofini@galois.com>
Date: Wed, 21 Feb 2024 11:43:25 -0600
Subject: [PATCH] Code reorg/cleanup

---
 pate_binja/pate.py | 24 +++++++++++++-----------
 pate_binja/view.py | 23 +++++++++++++++++------
 2 files changed, 30 insertions(+), 17 deletions(-)

diff --git a/pate_binja/pate.py b/pate_binja/pate.py
index e7fa2d1f..9c7f72d1 100644
--- a/pate_binja/pate.py
+++ b/pate_binja/pate.py
@@ -428,13 +428,6 @@ def show_message(self, rec: Any):
 
 
 class CFARNode:
-
-    id: str
-    desc: str
-    data = None
-    predomain = None
-    postdomain = None
-    external_postdomain = None
     exits: list[CFARNode]
 
     def __init__(self, id: str, desc: str, data: dict):
@@ -444,6 +437,10 @@ def __init__(self, id: str, desc: str, data: dict):
         self.update_node(desc, data)
         self.desc = desc
         self.data = data
+        self.predomain = None
+        self.postdomain = None
+        self.external_postdomain = None
+        self.addr = None
 
     def update_node(self, desc: str, data: dict):
         self.desc = desc
@@ -528,7 +525,6 @@ def pprint_node_event_trace_patched(self, trace, label: str, pre: str = '', out:
             pprint_event_trace(f'{label} Patched', trace['traces']['patched'], pre, out)
 
 class CFARGraph:
-
     nodes: dict[str, CFARNode]
 
     def __init__(self):
@@ -1265,10 +1261,15 @@ def run_replay(file: str) -> Popen:
         )
 
 
-def run_pate_config(file):
+def get_run_config(file: os.PathLike) -> dict:
     with open(file, 'r') as f:
         config = json.load(f)
-    cwd = os.path.dirname(file)
+    config['cwd'] = os.path.dirname(file)
+    return config
+
+
+def run_config(config: dict):
+    cwd = config.get('cwd')
     original = config.get('original')
     patched = config.get('patched')
     rawargs = config.get('args')
@@ -1292,7 +1293,8 @@ def run_pate(cwd: str, original: str, patched: str, args: list[str]) -> Popen:
 
 def run_pate_config_or_replay_file(f: str) -> Popen:
     if f.endswith(".run-config.json"):
-        test_live(lambda ignore: run_pate_config(f))
+        config = get_run_config(f)
+        test_live(lambda ignore: run_config(config))
     elif f.endswith(".replay"):
         test_replay(lambda ignore: run_replay(f))
 
diff --git a/pate_binja/view.py b/pate_binja/view.py
index 40541f0b..b45d5cfb 100644
--- a/pate_binja/view.py
+++ b/pate_binja/view.py
@@ -27,10 +27,11 @@
 
 
 class PateWidget(QWidget):
-    def __init__(self, parent: QWidget, filename: str) -> None:
+    def __init__(self, parent: QWidget, filename: str, config: dict) -> None:
         super().__init__(parent)
 
         self.filename = filename
+        self.config = config
         self.pate_thread: PateThread = None
 
         self.flow_graph_widget = MyFlowGraphWidget(self)
@@ -125,6 +126,10 @@ def show_cfar_graph(self, graph: pate.CFARGraph) -> None:
 
 
 class PateThread(Thread):
+    proc: Popen or None
+    oBv: BinaryView | None
+    pBv: BinaryView | None
+
     # TODO: Look at interaction.run_progress_dialog
     # handle cancel and restart
     def __init__(self, bv, run_fn, pate_widget: PateWidget, replay=False, show_ce_trace=True, trace_file=None):
@@ -135,7 +140,9 @@ def __init__(self, bv, run_fn, pate_widget: PateWidget, replay=False, show_ce_tr
         self.run_fn = run_fn
         self.trace_file = trace_file
         self.show_ce_trace = show_ce_trace
-        self.proc: Popen = None
+        self.proc = None
+        self.oBv = None
+        self.pBv = None
 
     def run(self):
 
@@ -352,16 +359,20 @@ def launch_pate(context: UIActionContext):
     if f.endswith(".run-config.json"):
         replay = False
         trace_file = os.path.join(os.path.dirname(f), 'lastrun.replay')
-        fn = lambda ignore: pate.run_pate_config(f)
+        config = pate.get_run_config(f)
     elif f.endswith(".replay"):
         replay = True
         trace_file = None
-        fn = lambda ignore: pate.run_replay(f)
+        config = {}  # TODO: need config for replay. Save it in replay file as first line?
 
-    pate_widget = PateWidget(context.widget, f)
+    pate_widget = PateWidget(context.widget, f, config)
     tab = context.context.createTabForWidget("PATE " + os.path.basename(f), pate_widget)
 
-    pt = PateThread(None, fn, pate_widget, replay=replay, trace_file=trace_file)
+    if replay:
+        pt = PateThread(None, lambda i: pate.run_replay(f), pate_widget, replay=replay, trace_file=trace_file)
+    else:
+        pt = PateThread(None, lambda i: pate.run_config(config), pate_widget, replay=replay, trace_file=trace_file)
+
     pate_widget.pate_thread = pt
     pt.start()