uidf refactor

5 files changed, 182 insertions, 20 deletions

diff --git a/python/function.py b/python/function.py
index d4c7c92c..83666ab9 100644
--- a/python/function.py
+++ b/python/function.py
@@ -2835,7 +2835,21 @@ class Function:
 			arch = self.arch
 		return core.BNIsCallInstruction(self.handle, arch.handle, addr)
 
-	def set_user_var_value(self, var: 'variable.Variable', def_addr: int, value: 'variable.PossibleValueSet') -> None:
+	def create_forced_var_version(self, var: 'variable.Variable', def_addr: int) -> None:
+		def_site = core.BNArchitectureAndAddress()
+		def_site.arch = self.arch.handle
+		def_site.address = def_addr
+
+		core.BNCreateForcedVariableVersion(self.handle, var.to_BNVariable(), def_site)
+
+	def clear_forced_var_version(self, var: 'variable.Variable', def_addr: int) -> None:
+		def_site = core.BNArchitectureAndAddress()
+		def_site.arch = self.arch.handle
+		def_site.address = def_addr
+
+		core.BNClearForcedVariableVersion(self.handle, var.to_BNVariable(), def_site)
+
+	def set_user_var_value(self, var: 'variable.Variable', def_addr: int, value: 'variable.PossibleValueSet', after: bool = True) -> None:
 		"""
 		`set_user_var_value` allows the user to specify a PossibleValueSet value for an MLIL variable at its \
 		definition site.
@@ -2846,8 +2860,9 @@ class Function:
 		the HLIL layer.
 
 		:param Variable var: Variable for which the value is to be set
-		:param int def_addr: Address of the definition site of the variable
+		:param int def_addr: Address where the variable is set
 		:param PossibleValueSet value: Informed value of the variable
+		:param bool after: Whether the value happens before or after the instruction
 		:rtype: None
 
 		:Example:
@@ -2857,27 +2872,27 @@ class Function:
 			>>> var_value = PossibleValueSet.constant(5)
 			>>> current_function.set_user_var_value(mlil_var, def_address, var_value)
 		"""
-		if var.index == 0:
-			# Special case: function parameters have index 0 and are defined at the start of the function
-			def_addr = self.start
-		else:
-			var_defs = self.mlil.get_var_definitions(var)
-			if var_defs is None:
-				raise ValueError("Could not get definition for Variable")
-			found = False
-			for site in var_defs:
-				if site.address == def_addr:
-					found = True
-					break
-			if not found:
-				raise ValueError("No definition for Variable found at given address")
+		#if var.index == 0:
+		#	# Special case: function parameters have index 0 and are defined at the start of the function
+		#	def_addr = self.start
+		#else:
+		#	var_defs = self.mlil.get_var_definitions(var)
+		#	if var_defs is None:
+		#		raise ValueError("Could not get definition for Variable")
+		#	found = False
+		#	for site in var_defs:
+		#		if site.address == def_addr:
+		#			found = True
+		#			break
+		#	if not found:
+		#		raise ValueError("No definition for Variable found at given address")
 		def_site = core.BNArchitectureAndAddress()
 		def_site.arch = self.arch.handle
 		def_site.address = def_addr
 
-		core.BNSetUserVariableValue(self.handle, var.to_BNVariable(), def_site, value._to_core_struct())
+		core.BNSetUserVariableValue(self.handle, var.to_BNVariable(), def_site, after, value._to_core_struct())
 
-	def clear_user_var_value(self, var: 'variable.Variable', def_addr: int) -> None:
+	def clear_user_var_value(self, var: 'variable.Variable', def_addr: int, after: bool = True) -> None:
 		"""
 		Clears a previously defined user variable value.
 
@@ -2904,7 +2919,7 @@ class Function:
 		def_site.arch = self.arch.handle
 		def_site.address = def_addr
 
-		core.BNClearUserVariableValue(self.handle, var.to_BNVariable(), def_site)
+		core.BNClearUserVariableValue(self.handle, var.to_BNVariable(), def_site, after)
 
 	def get_all_user_var_values(
 	    self
diff --git a/python/highlevelil.py b/python/highlevelil.py
index 85be4794..154af642 100644
--- a/python/highlevelil.py
+++ b/python/highlevelil.py
@@ -766,6 +766,12 @@ class HighLevelILInstruction(BaseILInstruction):
 	def get_label(self, operand_index: int) -> GotoLabel:
 		return GotoLabel(self.function, self.core_instr.operands[operand_index])
 
+	def get_constraint(self, operand_index: int) -> variable.PossibleValueSet:
+		value = core.BNGetCachedHighLevelILPossibleValueSet(self.function.handle, self.core_instr.operands[operand_index])
+		result = variable.PossibleValueSet(self.function.arch, value)
+		core.BNFreePossibleValueSet(value)
+		return result
+
 	@property
 	def raw_operands(self) -> OperandsType:
 		"""Raw operand expression indices as specified by the core structure (read-only)"""
@@ -2285,6 +2291,46 @@ class HighLevelILFcmpO(HighLevelILComparisonBase, FloatingPoint):
 class HighLevelILFcmpUo(HighLevelILComparisonBase, FloatingPoint):
 	pass
 
+@dataclass(frozen=True, repr=False, eq=False)
+class HighLevelILAssert(HighLevelILInstruction):
+    @property
+    def src(self) -> variable.Variable:
+        return self.get_var(0)
+
+    @property
+    def constraint(self) -> variable.PossibleValueSet:
+        return self.get_constraint(1)
+
+@dataclass(frozen=True, repr=False, eq=False)
+class HighLevelILAssertSsa(HighLevelILInstruction, SSA):
+    @property
+    def src(self) -> 'mediumlevelil.SSAVariable':
+        return self.get_var_ssa(0, 1)
+
+    @property
+    def constraint(self) -> variable.PossibleValueSet:
+        return self.get_constraint(2)
+
+@dataclass(frozen=True, repr=False, eq=False)
+class HighLevelILForceVer(HighLevelILInstruction):
+    @property
+    def dest(self) -> variable.Variable:
+        return self.get_var(0)
+
+    @property
+    def src(self) -> variable.Variable:
+        return self.get_var(1)
+
+@dataclass(frozen=True, repr=False, eq=False)
+class HighLevelILForceVerSsa(HighLevelILInstruction, SSA):
+    @property
+    def dest(self) -> 'mediumlevelil.SSAVariable':
+        return self.get_var_ssa(0, 1)
+
+    @property
+    def src(self) -> 'mediumlevelil.SSAVariable':
+        return self.get_var_ssa(2, 3)
+
 
 ILInstruction = {
     HighLevelILOperation.HLIL_NOP: HighLevelILNop,  #  ,
@@ -2423,6 +2469,10 @@ ILInstruction = {
     HighLevelILOperation.HLIL_FCMP_GT: HighLevelILFcmpGt,  #  ("left", "expr"), ("right", "expr"),
     HighLevelILOperation.HLIL_FCMP_O: HighLevelILFcmpO,  #  ("left", "expr"), ("right", "expr"),
     HighLevelILOperation.HLIL_FCMP_UO: HighLevelILFcmpUo,  #  ("left", "expr"), ("right", "expr"),
+    HighLevelILOperation.HLIL_ASSERT: HighLevelILAssert,
+    HighLevelILOperation.HLIL_ASSERT_SSA: HighLevelILAssertSsa,
+    HighLevelILOperation.HLIL_FORCE_VER: HighLevelILForceVer,
+    HighLevelILOperation.HLIL_FORCE_VER_SSA: HighLevelILForceVerSsa,
 }
 
 
diff --git a/python/lowlevelil.py b/python/lowlevelil.py
index 2d40c22c..4c2f0de9 100644
--- a/python/lowlevelil.py
+++ b/python/lowlevelil.py
@@ -1171,6 +1171,12 @@ class LowLevelILInstruction(BaseILInstruction):
 		finally:
 			core.BNLowLevelILFreeOperandList(operand_list)
 
+	def _get_constraint(self, operand_index: int) -> variable.PossibleValueSet:
+		value = core.BNGetCachedLowLevelILPossibleValueSet(self.function.handle, self.instr.operands[operand_index])
+		result = variable.PossibleValueSet(self.function.arch, value)
+		core.BNFreePossibleValueSet(value)
+		return result
+
 
 @dataclass(frozen=True, repr=False, eq=False)
 class LowLevelILBinaryBase(LowLevelILInstruction, BinaryOperation):
@@ -3011,6 +3017,42 @@ class LowLevelILStoreSsa(LowLevelILInstruction, Store, SSA):
 			("src", self.src, "LowLevelILInstruction"),
 		]
 
+@dataclass(frozen=True, repr=False, eq=False)
+class LowLevelILAssert(LowLevelILInstruction):
+    @property
+    def src(self) -> ILRegister:
+        return self._get_reg(0)
+
+    @property
+    def constraint(self) -> variable.PossibleValueSet:
+        return self._get_constraint(1)
+
+@dataclass(frozen=True, repr=False, eq=False)
+class LowLevelILAssertSsa(LowLevelILInstruction, SSA):
+    @property
+    def src(self) -> SSARegister:
+        return self._get_reg_ssa(0, 1)
+
+    @property
+    def constraint(self) -> variable.PossibleValueSet:
+        return self._get_constraint(2)
+
+@dataclass(frozen=True, repr=False, eq=False)
+class LowLevelILForceVer(LowLevelILInstruction):
+    @property
+    def dest(self) -> ILRegister:
+        return self._get_reg(0)
+
+@dataclass(frozen=True, repr=False, eq=False)
+class LowLevelILForceVerSsa(LowLevelILInstruction, SSA):
+    @property
+    def dest(self) -> SSARegister:
+        return self._get_reg_ssa(0, 1)
+
+    @property
+    def src(self) -> SSARegister:
+        return self._get_reg_ssa(2, 3)
+
 
 ILInstruction:Dict[LowLevelILOperation, LowLevelILInstruction] = {  # type: ignore
     LowLevelILOperation.LLIL_NOP: LowLevelILNop,                                    #  [],
@@ -3152,6 +3194,10 @@ ILInstruction:Dict[LowLevelILOperation, LowLevelILInstruction] = {  # type: igno
     LowLevelILOperation.LLIL_REG_STACK_PHI: LowLevelILRegStackPhi,                  #  [("dest", "reg_stack_ssa"), ("src", "reg_stack_ssa_list")],
     LowLevelILOperation.LLIL_FLAG_PHI: LowLevelILFlagPhi,                           #  [("dest", "flag_ssa"), ("src", "flag_ssa_list")],
     LowLevelILOperation.LLIL_MEM_PHI: LowLevelILMemPhi,                             #  [("dest_memory", "int"), ("src_memory", "int_list")]
+    LowLevelILOperation.LLIL_ASSERT: LowLevelILAssert,                              #  [("src", "reg"), ("constraint", "constraint")]
+    LowLevelILOperation.LLIL_ASSERT_SSA: LowLevelILAssertSsa,                       #  [("src", "reg_ssa"), ("constraint", "constraint")]
+    LowLevelILOperation.LLIL_FORCE_VER: LowLevelILForceVer,                         #  [("dest", "reg"), ("src", "reg")]
+    LowLevelILOperation.LLIL_FORCE_VER_SSA: LowLevelILForceVerSsa,                  #  [("dest", "reg_ssa"), ("src", "reg_ssa")]
 }
 
 
diff --git a/python/mediumlevelil.py b/python/mediumlevelil.py
index 976c1a85..4ec79d8d 100644
--- a/python/mediumlevelil.py
+++ b/python/mediumlevelil.py
@@ -1101,6 +1101,12 @@ class MediumLevelILInstruction(BaseILInstruction):
 		finally:
 			core.BNMediumLevelILFreeOperandList(operand_list)
 
+	def _get_constraint(self, operand_index: int) -> variable.PossibleValueSet:
+		value = core.BNGetCachedMediumLevelILPossibleValueSet(self.function.handle, self.instr.operands[operand_index])
+		result = variable.PossibleValueSet(self.function.arch, value)
+		core.BNFreePossibleValueSet(value)
+		return result
+
 
 @dataclass(frozen=True, repr=False, eq=False)
 class MediumLevelILConstBase(MediumLevelILInstruction, Constant):
@@ -3001,6 +3007,47 @@ class MediumLevelILStoreStructSsa(MediumLevelILInstruction, Store, SSA):
 			('src', self.src, 'MediumLevelILInstruction'),
 		]
 
+@dataclass(frozen=True, repr=False, eq=False)
+class MediumLevelILAssert(MediumLevelILInstruction):
+    @property
+    def src(self) -> variable.Variable:
+        return self._get_var(0)
+
+    @property
+    def constraint(self) -> variable.PossibleValueSet:
+        return self._get_constraint(1)
+
+@dataclass(frozen=True, repr=False, eq=False)
+class MediumLevelILAssertSsa(MediumLevelILInstruction, SSA):
+    @property
+    def src(self) -> SSAVariable:
+        return self._get_var_ssa(0, 1)
+
+    @property
+    def constraint(self) -> variable.PossibleValueSet:
+        return self._get_constraint(2)
+
+@dataclass(frozen=True, repr=False, eq=False)
+class MediumLevelILForceVer(MediumLevelILInstruction):
+    @property
+    def dest(self) -> variable.Variable:
+        return self._get_var(0)
+
+    @property
+    def src(self) -> variable.Variable:
+        return self._get_var(1)
+
+@dataclass(frozen=True, repr=False, eq=False)
+class MediumLevelILForceVerSsa(MediumLevelILInstruction, SSA):
+    @property
+    def dest(self) -> SSAVariable:
+        return self._get_var_ssa(0, 1)
+
+    @property
+    def src(self) -> SSAVariable:
+        return self._get_var_ssa(2, 3)
+
+
 
 ILInstruction = {
     MediumLevelILOperation.MLIL_NOP: MediumLevelILNop,  # [],
@@ -3164,6 +3211,10 @@ ILInstruction = {
         MediumLevelILCallUntyped,  # [("output", "expr"), ("dest", "expr"), ("params", "expr"), ("stack", "expr")],
     MediumLevelILOperation.MLIL_STORE_STRUCT_SSA:
         MediumLevelILStoreStructSsa,  # [("dest", "expr"), ("offset", "int"), ("dest_memory", "int"), ("src_memory", "int"), ("src", "expr")],
+    MediumLevelILOperation.MLIL_ASSERT: MediumLevelILAssert,
+    MediumLevelILOperation.MLIL_ASSERT_SSA: MediumLevelILAssertSsa,
+    MediumLevelILOperation.MLIL_FORCE_VER: MediumLevelILForceVer,
+    MediumLevelILOperation.MLIL_FORCE_VER_SSA: MediumLevelILForceVerSsa,
 }
 
 
diff --git a/python/variable.py b/python/variable.py
index 2a8ee026..e33f9af6 100644
--- a/python/variable.py
+++ b/python/variable.py
@@ -380,7 +380,7 @@ class PossibleValueSet:
 			return self.ranges == other.ranges
 		elif self.type in [RegisterValueType.InSetOfValues, RegisterValueType.NotInSetOfValues]:
 			return self.values == other.values
-		elif self.type == RegisterValueType.UndeterminedValue and hasattr(other, 'type'):
+		elif self.type == RegisterValueType.UndeterminedValue and hasattr(other, '_type'):
 			return self.type == other.type
 		else:
 			return self == other