From 0e3ab0e2a93979c756bc5f2b29033aaa8714fa56 Mon Sep 17 00:00:00 2001
From: chkp-eyalit <eyalit@checkpoint.com>
Date: Wed, 27 May 2020 14:24:38 +0300
Subject: [PATCH] [Bug Fix] Enforce logic applies only to active code types

Fixes the main bug in issue #41. When there was only 1 active code
type (THUMB) and we were scanning a line from the other code type (ARM)
we triggered an exception.

Scanned the code to make sure no prediction will be made on code
types which are not active at the moment of the decision.
---
 src/thumbs_up/analyzer_utils.py     |  5 +++--
 src/thumbs_up/analyzers/analyzer.py | 15 +++++++++++++--
 src/thumbs_up/utils/fptr.py         |  2 +-
 3 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/src/thumbs_up/analyzer_utils.py b/src/thumbs_up/analyzer_utils.py
index b844253..ca5c4b0 100644
--- a/src/thumbs_up/analyzer_utils.py
+++ b/src/thumbs_up/analyzer_utils.py
@@ -153,8 +153,10 @@ def functionScan(analyzer, scs):
             if search_func or analyzer.switch_identifier.isSwitchCase(line.start_ea):
                 line = line.next
                 continue
+            original_code_type = analyzer.codeType(line.start_ea)
             # If this is code, check that it matches the start of a function, and make it a function
-            if line.is_code and analyzer.func_classifier.predictFunctionStartMixed(line.start_ea):
+            if line.is_code and analyzer.supportedCodeType(original_code_type) and \
+                        analyzer.func_classifier.predictFunctionStartMixed(line.start_ea):
                 if not ida_funcs.add_func(line.start_ea):
                     line = line.next
                 else:
@@ -168,7 +170,6 @@ def functionScan(analyzer, scs):
             # If unknown, check if a function and don't try to keep the same code type
             if line.is_unknown:
                 guess_code_type = analyzer.func_classifier.predictFunctionStartType(line.start_ea)
-                original_code_type = analyzer.codeType(line.start_ea)
                 if analyzer.func_classifier.predictFunctionStart(line.start_ea, guess_code_type):
                     if original_code_type != guess_code_type:
                         analyzer.setCodeType(line.start_ea, line.start_ea + 1, guess_code_type)
diff --git a/src/thumbs_up/analyzers/analyzer.py b/src/thumbs_up/analyzers/analyzer.py
index 935b78e..cd48994 100644
--- a/src/thumbs_up/analyzers/analyzer.py
+++ b/src/thumbs_up/analyzers/analyzer.py
@@ -247,7 +247,7 @@ def isValidCodePtr(self, ptr_ea):
             True iff the code pointer is valid
         """
         ptr_type = self.ptrCodeType(ptr_ea)
-        return self.isCodeAligned(self.cleanPtr(ptr_ea), ptr_type) and ptr_type in self.activeCodeTypes()
+        return self.isCodeAligned(self.cleanPtr(ptr_ea), ptr_type) and self.supportedCodeType(ptr_type)
 
     def hasCodeTypes(self):
         """Check if the given CPU has multiple code types.
@@ -289,9 +289,20 @@ def disableCodeType(self, code_type):
         Args:
             code_type (int): code type to be disabled
         """
-        if code_type in self._active_code_types:
+        if self.supportedCodeType(code_type):
             self._active_code_types.remove(code_type)
 
+    def supportedCodeType(self, code_type):
+        """Check if a given code_type is actively supported.
+
+        Args:
+            code_type (int): code type to be checked
+
+        Return Value:
+            The code type of the annotated pointer
+        """
+        return code_type in self._active_code_types
+
     def ptrCodeType(self, ptr_ea):
         """Extract the code type of the annotated pointer.
 
diff --git a/src/thumbs_up/utils/fptr.py b/src/thumbs_up/utils/fptr.py
index a74bacb..bfbbd7f 100644
--- a/src/thumbs_up/utils/fptr.py
+++ b/src/thumbs_up/utils/fptr.py
@@ -200,7 +200,7 @@ def locateDataPtrs(self, scs, sds):
                     continue
                 # check for a function ptr
                 value = self._analyzer.parseAdderss(cur_ea)
-                # make sure it is valid
+                # make sure it is valid (enforces that the code_type is active)
                 if self.isValidCodePtr(value, scs):
                     func_value = self._analyzer.cleanPtr(value)
                     code_type  = self._analyzer.ptrCodeType(value)