Adds tests related to UIDF

3 files changed, 232 insertions, 7 deletions

diff --git a/binaryview.cpp b/binaryview.cpp
index 59f14eef..3953cee6 100644
--- a/binaryview.cpp
+++ b/binaryview.cpp
@@ -2446,7 +2446,6 @@ uint64_t BinaryView::GetPreviousDataVariableStartBeforeAddress(uint64_t addr)
 }
 
 
-<<<<<<< HEAD
 bool BinaryView::ParsePossibleValueSetString(const string& value, BNRegisterValueType state, PossibleValueSet& result, uint64_t here, string& errors)
 {
 	BNPossibleValueSet res;
diff --git a/python/function.py b/python/function.py
index bb4f817a..44c2f383 100644
--- a/python/function.py
+++ b/python/function.py
@@ -263,8 +263,11 @@ class ValueRange(object):
 			return "<range: %#x to %#x>" % (self.start, self.end)
 		return "<range: %#x to %#x, step %#x>" % (self.start, self.end, self.step)
 
-
-
+	def __eq__(self, other):
+		if not isinstance(other, self.__class__):
+			return NotImplemented
+		return self.start == other.start and self.end == other.end and self.step == other.step
+		
 	@property
 	def start(self):
 		""" """
@@ -361,7 +364,7 @@ class PossibleValueSet(object):
 		if self._type == RegisterValueType.ConstantPointerValue:
 			return "<const ptr %#x>" % self.value
 		if self._type == RegisterValueType.StackFrameOffset:
-			return "<stack frame offset %#x>" % self.offset
+			return "<stack frame offset %#x>" % self._offset
 		if self._type == RegisterValueType.SignedRangeValue:
 			return "<signed ranges: %s>" % repr(self.ranges)
 		if self._type == RegisterValueType.UnsignedRangeValue:
@@ -377,10 +380,18 @@ class PossibleValueSet(object):
 		return "<undetermined>"
 
 	def __eq__(self, other):
-		if self.type in [RegisterValueType.ConstantValue, RegisterValueType.ConstantValue] and isinstance(other, numbers.Integral):
+		if self.type in [RegisterValueType.ConstantValue, RegisterValueType.ConstantPointerValue] and isinstance(other, numbers.Integral):
 			return self.value == other
-		if self.type in [RegisterValueType.ConstantValue, RegisterValueType.ConstantValue] and hasattr(other, 'type') and other.type == self.type:
+		if self.type in [RegisterValueType.ConstantValue, RegisterValueType.ConstantPointerValue] and hasattr(other, 'type') and other.type == self.type:
 			return self.value == other.value
+		elif self.type == RegisterValueType.StackFrameOffset and hasattr(other, 'type') and other.type == self.type:
+			return self.offset == other.offset
+		elif self.type in [RegisterValueType.SignedRangeValue, RegisterValueType.UnsignedRangeValue] and hasattr(other, 'type') and other.type == self.type:
+			return self.ranges == other.ranges
+		elif self.type in [RegisterValueType.InSetOfValues, RegisterValueType.NotInSetOfValues] and hasattr(other, 'type') and other.type == self.type:
+			return self.values == other.values
+		elif self.type == RegisterValueType.UndeterminedValue and hasattr(other, 'type'):
+			return self.type == other.type
 		else:
 			return self == other
 
@@ -581,7 +592,7 @@ class PossibleValueSet(object):
 		"""
 		result = PossibleValueSet()
 		result.type = RegisterValueType.StackFrameOffset
-		result.value = value
+		result.offset = offset
 		return result
 
 	@classmethod
diff --git a/suite/testcommon.py b/suite/testcommon.py
index 33b3b800..a523fe3a 100644
--- a/suite/testcommon.py
+++ b/suite/testcommon.py
@@ -883,6 +883,56 @@ class VerifyBuilder(Builder):
     def get_comments(self, bv):
         return bv.functions[0].comments
 
+    def test_possiblevalueset_parse(self):
+        """ Failed to parse PossibleValueSet from string"""
+        file_name = self.unpackage_file("helloworld")
+        try:
+            bv = binja.BinaryViewType.get_view_of_file(file_name)
+            # ConstantValue
+            lhs = bv.parse_possiblevalueset_string("0", binja.RegisterValueType.ConstantValue) 
+            rhs = binja.PossibleValueSet.constant(0)
+            assert lhs == rhs
+            lhs = bv.parse_possiblevalueset_string("$here + 2", binja.RegisterValueType.ConstantValue, 0x2000)
+            rhs = binja.PossibleValueSet.constant(0x2000 + 2)
+            assert lhs == rhs
+            # ConstantPointerValue
+            lhs = bv.parse_possiblevalueset_string("0x8000", binja.RegisterValueType.ConstantPointerValue)
+            rhs = binja.PossibleValueSet.constant_ptr(0x8000)
+            assert lhs == rhs
+            # StackFrameOffset
+            lhs = bv.parse_possiblevalueset_string("16", binja.RegisterValueType.StackFrameOffset)
+            rhs = binja.PossibleValueSet.stack_frame_offset(0x16)
+            assert lhs == rhs
+            # SignedRangeValue
+            lhs = bv.parse_possiblevalueset_string("-10:0:2", binja.RegisterValueType.SignedRangeValue)
+            rhs = binja.PossibleValueSet.signed_range_value([binja.ValueRange(-0x10, 0, 2)])
+            assert lhs == rhs
+            lhs = bv.parse_possiblevalueset_string("-10:0:2,2:5:1", binja.RegisterValueType.SignedRangeValue)
+            rhs = binja.PossibleValueSet.signed_range_value([binja.ValueRange(-0x10, 0, 2), binja.ValueRange(2, 5, 1)])
+            assert lhs == rhs
+            # UnsignedRangeValue
+            lhs = bv.parse_possiblevalueset_string("1:10:1", binja.RegisterValueType.UnsignedRangeValue)
+            rhs = binja.PossibleValueSet.unsigned_range_value([binja.ValueRange(1, 0x10, 1)])
+            assert lhs == rhs
+            lhs = bv.parse_possiblevalueset_string("1:10:1, 2:20:2", binja.RegisterValueType.UnsignedRangeValue)
+            rhs = binja.PossibleValueSet.unsigned_range_value([binja.ValueRange(1, 0x10, 1), binja.ValueRange(2, 0x20, 2)])
+            assert lhs == rhs
+            # InSetOfValues
+            lhs = bv.parse_possiblevalueset_string("1,2,3,3,4", binja.RegisterValueType.InSetOfValues)
+            rhs = binja.PossibleValueSet.in_set_of_values([1,2,3,4])
+            assert lhs == rhs
+            # NotInSetOfValues
+            lhs = bv.parse_possiblevalueset_string("1,2,3,4,4", binja.RegisterValueType.NotInSetOfValues)
+            rhs = binja.PossibleValueSet.not_in_set_of_values([1,2,3,4])
+            assert lhs == rhs
+            # UndeterminedValue
+            lhs = bv.parse_possiblevalueset_string("", binja.RegisterValueType.UndeterminedValue)
+            rhs = binja.PossibleValueSet.undetermined()
+            assert lhs == rhs
+            return True
+        finally:
+            self.delete_package("helloworld")
+
     def test_expression_parse(self):
         file_name = self.unpackage_file("helloworld")
         try:
@@ -1205,3 +1255,168 @@ class VerifyBuilder(Builder):
         finally:
             binja.Settings().set_string_list("files.universal.architecturePreference", save_setting_value)
             self.delete_package("fat_macho_9arch")
+
+    def test_user_informed_dataflow(self):
+        """User-informed dataflow tests"""
+        file_name = self.unpackage_file("helloworld")
+        try:
+            with binja.BinaryViewType['ELF'].open(file_name) as bv:
+                bv.update_analysis_and_wait()
+                func = bv.get_function_at(0x00008440)
+
+                ins_idx = func.mlil.get_instruction_start(0x845c)
+                ins = func.mlil[ins_idx]
+                assert(ins.operation == binja.MediumLevelILOperation.MLIL_IF)
+                assert(len(ins.vars_read) == 1)
+                var = ins.vars_read[0]
+                defs = func.mlil.get_var_definitions(var)
+                assert(len(defs) == 1)
+                def_site = defs[0].address
+
+                # Set variable value to 0
+                bv.begin_undo_actions()
+                func.set_user_var_value(var, def_site, binja.PossibleValueSet.constant(0))
+                bv.commit_undo_actions()
+                bv.update_analysis_and_wait()
+
+                ins_idx = func.mlil.get_instruction_start(0x845c)
+                ins = func.mlil[ins_idx]
+                assert(ins.operation == binja.MediumLevelILOperation.MLIL_IF)
+                # test if condition value is updated to true
+                assert(ins.condition.value == True)
+                # test if register value is updated to 0
+                assert(ins.get_reg_value_after('r3') == 0)
+                # test if branch is eliminated in hlil
+                for hlil_ins in func.hlil.instructions:
+                    assert(hlil_ins.operation != binja.HighLevelILOperation.HLIL_IF)
+
+                # test undo action
+                bv.undo()
+                bv.update_analysis_and_wait()
+                ins_idx = func.mlil.get_instruction_start(0x845c)
+                ins = func.mlil[ins_idx]
+                assert(ins.operation == binja.MediumLevelILOperation.MLIL_IF)
+                # test if condition value is updated to undetermined
+                assert(ins.condition.value.type == binja.RegisterValueType.UndeterminedValue)
+                # test if register value is updated to undetermined
+                assert(ins.get_reg_value_after('r3').type == binja.RegisterValueType.EntryValue)
+                # test if branch is restored in hlil
+                found = False
+                for hlil_ins in func.hlil.instructions:
+                    if hlil_ins.operation == binja.HighLevelILOperation.HLIL_IF:
+                        found = True
+                assert(found)
+
+                # test redo action
+                bv.redo()
+                bv.update_analysis_and_wait()
+                ins_idx = func.mlil.get_instruction_start(0x845c)
+                ins = func.mlil[ins_idx]
+                assert(ins.operation == binja.MediumLevelILOperation.MLIL_IF)
+                # test if condition value is updated to true
+                assert(ins.condition.value == True)
+                # test if register value is updated to 0
+                assert(ins.get_reg_value_after('r3') == 0)
+                # test if branch is eliminated in hlil
+                for hlil_ins in func.hlil.instructions:
+                    assert(hlil_ins.operation != binja.HighLevelILOperation.HLIL_IF)
+
+                # test bndb round trip
+                temp_name = next(tempfile._get_candidate_names()) + ".bndb"
+                bv.create_database(temp_name)
+                bv.file.close()
+
+            with binja.FileMetadata(temp_name).open_existing_database(temp_name).get_view_of_type('ELF') as bv:
+                bv.update_analysis_and_wait()
+                func = bv.get_function_at(0x00008440)
+
+                ins_idx = func.mlil.get_instruction_start(0x845c)
+                ins = func.mlil[ins_idx]
+                assert(ins.operation == binja.MediumLevelILOperation.MLIL_IF)
+                # test if condition value is updated to true
+                assert(ins.condition.value == True)
+                # test if register value is updated to 0
+                assert(ins.get_reg_value_after('r3') == 0)
+                # test if branch is eliminated in hlil
+                for hlil_ins in func.hlil.instructions:
+                    assert(hlil_ins.operation != binja.HighLevelILOperation.HLIL_IF)
+
+                # test undo after round trip
+                bv.undo()
+                bv.update_analysis_and_wait()
+                ins_idx = func.mlil.get_instruction_start(0x845c)
+                ins = func.mlil[ins_idx]
+                assert(ins.operation == binja.MediumLevelILOperation.MLIL_IF)
+                # test if condition value is updated to undetermined
+                assert(ins.condition.value.type == binja.RegisterValueType.UndeterminedValue)
+                # test if register value is updated to undetermined
+                assert(ins.get_reg_value_after('r3').type == binja.RegisterValueType.EntryValue)
+                # test if branch is restored in hlil
+                found = False
+                for hlil_ins in func.hlil.instructions:
+                    if hlil_ins.operation == binja.HighLevelILOperation.HLIL_IF:
+                        found = True
+                assert(found)
+
+                bv.file.close()
+
+            os.unlink(temp_name)
+            # Check retrieval of user variable values
+            return True
+        finally:
+            self.delete_package("helloworld")
+
+    def test_possiblevalueset_ser_and_deser(self):
+        """PossibleValueSet serialization and deserialization"""
+        def test_helper(value):
+            file_name = self.unpackage_file("helloworld")
+            try:
+                with binja.BinaryViewType['ELF'].open(file_name) as bv:
+                    bv.update_analysis_and_wait()
+                    func = bv.get_function_at(0x00008440)
+
+                    ins_idx = func.mlil.get_instruction_start(0x845c)
+                    ins = func.mlil[ins_idx]
+
+                    var = ins.vars_read[0]
+                    defs = func.mlil.get_var_definitions(var)
+                    def_site = defs[0].address
+
+                    func.set_user_var_value(var, def_site, value)
+                    bv.update_analysis_and_wait()
+
+                    def_ins_idx = func.mlil.get_instruction_start(def_site)
+                    def_ins = func.mlil[def_ins_idx]
+
+                    assert(def_ins.get_possible_reg_values_after('r3') == value)
+
+                    temp_name = next(tempfile._get_candidate_names()) + ".bndb"
+                    bv.create_database(temp_name)
+                    bv.file.close()
+
+                with binja.FileMetadata(temp_name).open_existing_database(temp_name).get_view_of_type('ELF') as bv:
+                    bv.update_analysis_and_wait()
+                    func = bv.get_function_at(0x00008440)
+
+                    ins_idx = func.mlil.get_instruction_start(0x845c)
+                    ins = func.mlil[ins_idx]
+
+                    def_ins_idx = func.mlil.get_instruction_start(def_site)
+                    def_ins = func.mlil[def_ins_idx]
+
+                    assert(def_ins.get_possible_reg_values_after('r3') == value)
+                    bv.file.close()
+
+                os.unlink(temp_name)
+                return True
+
+            finally:
+                self.delete_package("helloworld")
+
+        assert(test_helper(binja.PossibleValueSet.constant(0)))
+        assert(test_helper(binja.PossibleValueSet.constant_ptr(0x8000)))
+        assert(test_helper(binja.PossibleValueSet.unsigned_range_value([binja.ValueRange(1, 10, 2)])))
+        # assert(test_helper(binja.PossibleValueSet.signed_range_value([binja.ValueRange(-10, 0, 2)])))
+        assert(test_helper(binja.PossibleValueSet.in_set_of_values([1,2,3,4])))
+        assert(test_helper(binja.PossibleValueSet.not_in_set_of_values([1,2,3,4])))
+        return True