From 2c97b4f8d48bcee33a6e8c08849b752ea99a7271 Mon Sep 17 00:00:00 2001
From: "joxeankoret@yahoo.es" <joxeankoret@yahoo.es>
Date: Fri, 8 Mar 2024 13:21:49 +0100
Subject: [PATCH] Various bug fixes

BUG: Messages from threads were being ignored by `log()`.
BUG: Timeouts in heuristics were not properly handled. Fix for issue #296.
HEUR: Increase the threads join time (`THREADS_WAIT_TIME`) to 1 second.
GUI: Display the progress when exporting a large number of compilation units every 1% of CUs exported.
---
 diaphora.py        | 19 ++++++++++++-------
 diaphora_config.py |  2 +-
 diaphora_ida.py    |  4 +++-
 3 files changed, 16 insertions(+), 9 deletions(-)

diff --git a/diaphora.py b/diaphora.py
index 59c253b..b3d89d6 100755
--- a/diaphora.py
+++ b/diaphora.py
@@ -174,11 +174,10 @@ def log(message):
   Print a message
   """
   # pylint: disable=protected-access
-  if isinstance(threading.current_thread(), threading._MainThread):
-    if IS_IDA or os.getenv("DIAPHORA_LOG_PRINT") is not None:
-      print(f"[Diaphora: {time.asctime()}] {message}")
-    else:
-      logging.info(message)
+  if IS_IDA or os.getenv("DIAPHORA_LOG_PRINT") is not None:
+    print(f"[Diaphora: {time.asctime()}] {message}")
+  else:
+    logging.info(message)
   # pylint: enable=protected-access
 
 
@@ -2041,8 +2040,8 @@ def add_matches_internal(
     t = time.monotonic()
     while self.continue_getting_sql_rows(i):
       if time.monotonic() - t > self.timeout or cur_thread.timeout:
-        log("Timeout")
-        break
+        log_refresh(f"Timeout with heuristic '{cur_thread.name}'")
+        raise SystemExit()
 
       i += 1
       if i % 50000 == 0:
@@ -2113,6 +2112,8 @@ def add_matches_from_query_ratio(
       self.add_matches_internal(
         cur, best=best, partial=partial, unreliable=unreliable, debug=debug
       )
+    except SystemExit:
+      pass
     except:
       log(f"Error: {str(sys.exc_info()[1])}")
       print("*" * 80)
@@ -2136,6 +2137,8 @@ def add_matches_from_query_ratio_max(self, sql, best, partial, val):
       self.add_matches_internal(
         cur, best=best, partial=partial, val=val, unreliable="unreliable"
       )
+    except SystemExit:
+      pass
     except:
       log(f"Error: {str(sys.exc_info()[1])}")
       print("*" * 80)
@@ -2160,6 +2163,8 @@ def add_matches_from_query_ratio_max_trusted(self, sql, val):
       self.add_matches_internal(
         cur, best="best", partial="partial", val=val, unreliable="partial"
       )
+    except SystemExit:
+      pass
     except:
       log(f"Error: {str(sys.exc_info()[1])}")
       print("*" * 80)
diff --git a/diaphora_config.py b/diaphora_config.py
index 691eb6b..a748885 100644
--- a/diaphora_config.py
+++ b/diaphora_config.py
@@ -190,7 +190,7 @@
 # joined for a fraction of seconds to wait for them to finish after each 
 # iteration in a thread. This value indicates the number of seconds to join. It
 # basically translates to this: thread.join(THREADS_WAIT_TIME).
-THREADS_WAIT_TIME = 0.5
+THREADS_WAIT_TIME = 1
 
 
 #-------------------------------------------------------------------------------
diff --git a/diaphora_ida.py b/diaphora_ida.py
index 9e1c155..5d0673b 100644
--- a/diaphora_ida.py
+++ b/diaphora_ida.py
@@ -3240,8 +3240,10 @@ def get_modules_using_lfa(self):
 
   def save_compilation_units(self):
     log_refresh("Finding compilation units...")
+    msg("Finding compilation units...")
     lfa_modules = self.get_modules_using_lfa()
     log_refresh("Saving compilation units...")
+    msg("Saving compilation units...")
 
     sql1 = """insert into compilation_units (name, start_ea, end_ea)
                   values (?, ?, ?)"""
@@ -3258,7 +3260,7 @@ def save_compilation_units(self):
     try:
       dones = set()
       total = len(lfa_modules)
-      checkpoint = int(total / 10)
+      checkpoint = int(total / 100)
       for i, module in enumerate(lfa_modules):
         if i > 0 and checkpoint > 0 and i % checkpoint == 0:
           log_refresh(f"Processing compilation unit {i} out of {total}...")