Fixes #18 but need some TU

haystack/abc/interfaces.py

@@ -432,12 +432,36 @@ class IModuleConstraints(object):
     x['struct_1']['field1'] contains a list of contraints.
     """
 
-    def get_records(self):
-        """get the list of record names."""
+    def get_constraints(self):
+        """
+        get the list of record_type_name,IConstraint for all fields of
+        :return dict
+        """
+        raise NotImplementedError('Please implement all methods')
+
+    def set_constraints(self, record_type_name, record_constraints):
+        """
+        Add constraints for that record_type name
+        :param record_type_name:
+        :param record_constraints:
+        :return
+        """
         raise NotImplementedError('Please implement all methods')
 
-    def get_constraints_for_record(self, record_name):
-        """get the list of IConstraint for all fields of
+    def get_dynamic_constraints(self):
+        """
+        get the record_type_name,IRecordTypeDynamicConstraintsValidator
+
+        :return dict
+        """
+        raise NotImplementedError('Please implement all methods')
+
+    def set_dynamic_constraints(self, record_type_name, record_constraints):
+        """
+        Add dynamic constraints validator for that record_type name
+        :param record_type_name: str
+        :param record_constraints: IRecordTypeDynamicConstraintsValidator
+        :return:
         """
         raise NotImplementedError('Please implement all methods')
 
@@ -508,3 +532,20 @@ def load_members(self, record, max_depth):
         :return:
         """
         raise NotImplementedError('Please implement all methods')
+
+
+class IRecordTypeDynamicConstraintsValidator(object):
+    """
+    A record-type-based constraints validation class
+    """
+    def get_record_type_name(self):
+        """Return the name of the record_type for which these advanced checks can occur"""
+        raise NotImplementedError('Please implement all methods')
+
+    def is_valid(self, record):
+        """
+        Advanced checks that cannot be expressed in the constraints files
+        """
+        raise NotImplementedError('Please implement all methods')
+
+    # TODO get_list_tuples

haystack/basicmodel.py

@@ -55,9 +55,10 @@ def get_record_type_fields(record_type):
         yield (f[0], f[1])
     raise StopIteration
 
+
 class CTypesRecordConstraintValidator(interfaces.IRecordConstraintsValidator):
     """
-    This is the main class, to be inherited by all ctypes structure.
+    This is the main class, to be inherited by all ctypes record validators.
     It adds a generic validation framework, based on simple assertion,
     and on more complex constraint on members values.
     """
@@ -80,14 +81,24 @@ def __init__(self, memory_handler, my_constraints):
         self._target = self._memory_handler.get_target_platform()
         self._ctypes = self._target.get_target_ctypes()
         self._utils = self._target.get_target_ctypes_utils()
-        self._constraints = my_constraints
+        self._constraints_base = None
+        self._constraints_dynamic = None
+        if my_constraints is not None:
+            self._constraints_base = my_constraints.get_constraints()
+            self._constraints_dynamic = my_constraints.get_dynamic_constraints()
 
     def _get_constraints_for(self, record):
         n = record.__class__.__name__
-        if self._constraints and n in self._constraints:
-            return self._constraints[n]
+        if self._constraints_base and n in self._constraints_base:
+            return self._constraints_base[n]
         return dict()
 
+    def _get_dynamic_constraints_for(self, record):
+        n = record.__class__.__name__
+        if self._constraints_dynamic and n in self._constraints_dynamic:
+            return self._constraints_dynamic[n]
+        return None
+
     def get_orig_addr(self, record):
         """ returns the vaddr of this instance."""
         haystack_addr = self._ctypes.addressof(record)
@@ -124,35 +135,24 @@ def is_valid(self, record):
             raise TypeError('Feed me a record')
         valid = self._is_valid(record, self._get_constraints_for(record))
         log.debug('-- <%s> isValid = %s', record.__class__.__name__, valid)
-        return valid
+        # check dynamic constraints last.
+        if valid:
+            return self._is_valid_dynamic_constraints(record)
+        return False
+
+    def _is_valid_dynamic_constraints(self, record):
+        dynamic_constraints = self._get_dynamic_constraints_for(record)
+        if dynamic_constraints:
+            return dynamic_constraints.is_valid(record)
+        return True
 
     def _is_valid(self, record, record_constraints):
         """ real implementation.    check expectedValues first, then the other fields """
         log.debug(' -- <%s> isValid --', record.__class__.__name__)
         done = []
-        log.debug("constraints are on %s", record_constraints)
-        # check all fields inline
-        #for attrname, attrtype in get_fields(record):
-        #    attr = getattr(record, attrname)
-        #    ignore = False
-        #    # shorcut ignores
-        #    if attrname in record_constraints:
-        #        for constraint in record_constraints[attrname]:
-        #            if constraint is constraints.IgnoreMember:
-        #                log.debug('ignoring %s as requested', attrname)
-        #                ignore = True
-        #                break
-        #            elif isinstance(constraint, constraints.ListLimitDepthValidation):
-        #                max_depth = constraint.max_depth
-        #                log.debug('setting max_depth %d as requested', max_depth)
-        #        continue
-        #    if ignore:
-        #        continue
-        #    # validate the attr
-        #    if not self._is_valid_attr(attr, attrname, attrtype, record_constraints):
-        #        return False
-
         # we check constrained field first to stop early if possible
+        # then we test the other fields
+        log.debug("constraints are on %s", record_constraints)
         _fieldsTuple = get_fields(record)
         myfields = dict(_fieldsTuple)
         for attrname, _constraints in record_constraints.iteritems():
@@ -174,9 +174,8 @@ def _is_valid(self, record, record_constraints):
                 continue
             if not self._is_valid_attr(attr, attrname, attrtype, record_constraints):
                 return False
-        # check the rest for validation
-        todo = [(name, typ)
-                for name, typ in get_fields(record) if name not in done]
+        # check the other fields for validation
+        todo = [(name, typ) for name, typ in get_fields(record) if name not in done]
         for attrname, attrtype, in todo:
             attr = getattr(record, attrname)
             if not self._is_valid_attr(attr, attrname, attrtype, record_constraints):

haystack/constraints.py

@@ -19,6 +19,7 @@
 
 log = logging.getLogger('constraints')
 
+
 from haystack.abc import interfaces
 
 
@@ -57,8 +58,7 @@ def read(self, filename):
         # each section anem is the name of the target structure
         for struct_name in parser.sections():
             log.debug('handling structure %s', struct_name)
-            if struct_name not in _constraints:
-                _constraints[struct_name] = RecordConstraints()
+            record_constraints = RecordConstraints()
             # each config entry is a field and its IConstraint
             for field, value in parser.items(struct_name):
                 log.debug('%s: field %s ::= %s', struct_name, field, value)
@@ -68,12 +68,14 @@ def read(self, filename):
                     raise ValueError("%s: struct_name: %s Field: %s constraint: %s" % (
                                      e.message, struct_name, field, value))
                 # each field can only have one IConstraint (which can be a list of)
-                if field not in _constraints[struct_name]:
-                    _constraints[struct_name][field] = []
+                if field not in record_constraints:
+                    record_constraints[field] = []
                 if isinstance(value, list):
-                    _constraints[struct_name][field].extend(value)
+                    record_constraints[field].extend(value)
                 else:
-                    _constraints[struct_name][field].append(value)
+                    record_constraints[field].append(value)
+            # we set it
+            _constraints.set_constraints(struct_name, record_constraints)
         return _constraints
 
     def _parse(self, value):
@@ -167,18 +169,49 @@ def _try_numbers(self, _arg):
         return ret
 
 
-class ModuleConstraints(interfaces.IModuleConstraints, dict):
+class ModuleConstraints(interfaces.IModuleConstraints):
     """
     Holds the constraints for all record types of a module.
     """
-    def get_records(self):
-        """get the list of record names."""
-        return self.keys()
+    def __init__(self):
+        self.__constraints = {}
+        self.__dynamics = {}
+
+    def get_constraints(self):
+        """
+        get the list of IConstraint for all fields of record_name
+
+        :return the list of IConstraint for that record
+        """
+        return self.__constraints
+
+    def set_constraints(self, record_type_name, record_constraints):
+        """
+        Add constraints for that record_type name
+        :param record_type_name:
+        :param record_constraints:
+        :return:
+        """
+        self.__constraints[record_type_name] = record_constraints
 
-    def get_constraints_for_record(self, record_name):
-        """get the list of IConstraint for all fields of
+    def get_dynamic_constraints(self):
         """
-        return self[record_name]
+        get the IRecordTypeDynamicConstraintsValidator for record_type_name
+
+        :return the list of IRecordTypeDynamicConstraintsValidator
+        """
+        return self.__dynamics
+
+    def set_dynamic_constraints(self, record_type_name, record_constraints):
+        """
+        Add dynamic constraints validator for that record_type name
+        :param record_type_name: str
+        :param record_constraints: IRecordTypeDynamicConstraintsValidator
+        :return:
+        """
+        assert isinstance(record_constraints, interfaces.IRecordTypeDynamicConstraintsValidator)
+        self.__dynamics[record_type_name] = record_constraints
+
 
 class RecordConstraints(interfaces.IRecordConstraints, dict):
     """
@@ -193,6 +226,7 @@ def get_constraints_for_field(self, field_name):
         """
         return self[field_name]
 
+
 class IgnoreMember(interfaces.IConstraint):
 
     """
@@ -204,6 +238,7 @@ class IgnoreMember(interfaces.IConstraint):
     def __contains__(self, obj):
         return True
 
+
 class ListLimitDepthValidation(interfaces.IConstraint):
 
     """

haystack/reverse/heuristics/reversers.py

@@ -113,7 +113,6 @@ def __init__(self, _memory_handler, record_names, record_constraints):
         self.__record_names = record_names
         self.__constraints = record_constraints
         self.__search_results = {}
-        self.__constraints = []
 
     def _iterate_records(self, _context):
         for x in self.__record_names:

haystack/search/searcher.py

@@ -58,12 +58,7 @@ def search(self, struct_type, max_res=10, max_depth=10):
         """
         outputs = []
         for m in self._target_mappings:
-            outputs.extend(
-                self._search_in(
-                    m,
-                    struct_type,
-                    nb=max_res-len(outputs),
-                    depth=max_depth))
+            outputs.extend(self._search_in(m, struct_type, nb=max_res-len(outputs), depth=max_depth))
             # check out
             if len(outputs) >= max_res:
                 break
@@ -97,7 +92,7 @@ def _search_in(self, mem_map, struct_type, nb=10, depth=99):
         for addr, size in walker.get_user_allocations():
             # FIXME, heap walker should give a hint
             # minimum chunk size varies...
-            #if size != struct_size:
+            # if size != struct_size:
             if size < struct_size:
                 log.debug("size %d < struct_size %d", size, struct_size)
                 continue
@@ -113,10 +108,9 @@ def _search_in(self, mem_map, struct_type, nb=10, depth=99):
             for offset in utils.xrange(start, end, plen):
                 # a - load and validate the record
                 # DEBUG
-                #offset = addr
+                # offset = addr
                 log.debug('load_at(%d) ', offset)
-                instance, validated = self._load_at(
-                    mem_map, offset, struct_type, depth)
+                instance, validated = self._load_at(mem_map, offset, struct_type, depth)
                 if validated:
                     log.debug("found instance @ 0x%lx", offset)
                     # do stuff with it.
@@ -180,23 +174,10 @@ def __init__(self, memory_handler, my_constraints=None, target_mappings=None, up
         :param update_cb: callback function to call for each valid result
         :return:
         """
-        if not isinstance(memory_handler, interfaces.IMemoryHandler):
-            raise TypeError("Feed me a IMemoryHandler")
-        if my_constraints and not isinstance(my_constraints, interfaces.IModuleConstraints):
-            raise TypeError("Feed me a IModuleConstraints")
-        if target_mappings is not None and not isinstance(target_mappings, list):
-            raise TypeError("Feed me a list of IMemoryMapping")
-        elif target_mappings is None:
+        if target_mappings is None:
             # default to all heaps
             target_mappings = memory_handler.get_mappings()
-        self._memory_handler = memory_handler
-        self._my_constraints = my_constraints
-        self._target_mappings = target_mappings
-        self._update_cb = update_cb
-        log.debug(
-            'StructFinder created for %s. Search Perimeter on %d mappings.',
-            self._memory_handler.get_name(),
-            len(self._target_mappings))
+        super(AnyOffsetRecordSearcher, self).__init__(memory_handler, my_constraints, target_mappings, update_cb)
         return
 
     def _search_in(self, mem_map, struct_type, nb=10, depth=99, align=None):
@@ -236,13 +217,11 @@ def _search_in(self, mem_map, struct_type, nb=10, depth=99, align=None):
             # print a debug message every now and then
             if offset % (1024 << 6) == 0:
                 p2 = offset - start
-                log.debug('processed %d bytes    - %02.02f test/sec',
-                          p2, (p2 - p) / (plen * (time.time() - t0)))
+                log.debug('processed %d bytes    - %02.02f test/sec', p2, (p2 - p) / (plen * (time.time() - t0)))
                 t0 = time.time()
                 p = p2
             # a - load and validate the record
-            instance, validated = self._load_at(
-                mem_map, offset, struct_type, depth)
+            instance, validated = self._load_at(mem_map, offset, struct_type, depth)
             if validated:
                 log.debug("found instance @ 0x%lx", offset)
                 # do stuff with it.

test/haystack/test_basicmodel.py

@@ -18,7 +18,7 @@
 __status__ = "Production"
 
 from test.haystack import SrcTests
-
+from haystack import constraints
 
 class TestLoadMembers(SrcTests):
 
@@ -31,7 +31,8 @@ def setUp(self):
         self.my_ctypes = self.my_target.get_target_ctypes()
         self.my_utils = self.my_target.get_target_ctypes_utils()
         self.my_model = self.memory_handler.get_model()
-        self.validator = basicmodel.CTypesRecordConstraintValidator(self.memory_handler, dict())
+        _constraints = constraints.ModuleConstraints() # empty
+        self.validator = basicmodel.CTypesRecordConstraintValidator(self.memory_handler, _constraints)
         self.ctypes_gen32 = self.my_model.import_module("test.src.ctypes5_gen32")
 
     def tearDown(self):
@@ -139,7 +140,8 @@ def setUp(self):
         self.my_ctypes = self.my_target.get_target_ctypes()
         self.my_utils = self.my_target.get_target_ctypes_utils()
         self.my_model = self.memory_handler.get_model()
-        self.validator = basicmodel.CTypesRecordConstraintValidator(self.memory_handler,dict())
+        _constraints = constraints.ModuleConstraints() # empty
+        self.validator = basicmodel.CTypesRecordConstraintValidator(self.memory_handler, _constraints)
         try:
             self.sslsnoop = self.my_model.import_module("sslsnoop")
         except ImportError: