LocationDB: introduced to replace AsmSymbolPool

1 files changed, 453 insertions, 0 deletions

diff --git a/miasm2/core/locationdb.py b/miasm2/core/locationdb.py
new file mode 100644
index 00000000..39c1c99a
--- /dev/null
+++ b/miasm2/core/locationdb.py
@@ -0,0 +1,453 @@
+import warnings
+
+from miasm2.expression.expression import LocKey, ExprLoc
+from miasm2.expression.modint import moduint, modint
+
+
+def is_int(a):
+    return isinstance(a, (int, long, moduint, modint))
+
+
+class LocationDB(object):
+    """
+    LocationDB is a "database" of information associated to location.
+
+    An entry in a LocationDB is uniquely identified with a LocKey.
+    Additionnal information which can be associated with a LocKey are:
+    - an offset (uniq per LocationDB)
+    - several names (each are uniqs per LocationDB)
+
+    As a schema:
+    loc_key  1 <-> 0..1  offset
+             1 <-> 0..n  name
+
+    >>> loc_db = LocationDB()
+    # Add a location with no additionnal information
+    >>> loc_key1 = loc_db.add_location()
+    # Add a location with an offset
+    >>> loc_key2 = loc_db.add_location(offset=0x1234)
+    # Add a location with several names
+    >>> loc_key3 = loc_db.add_location(name="first_name")
+    >>> loc_db.add_location_name(loc_key3, "second_name")
+    # Associate an offset to an existing location
+    >>> loc_db.set_location_offset(loc_key3, 0x5678)
+    # Remove a name from an existing location
+    >>> loc_db.remove_location_name(loc_key3, "second_name")
+
+    # Get back offset
+    >>> loc_db.get_location_offset(loc_key1)
+    None
+    >>> loc_db.get_location_offset(loc_key2)
+    0x1234
+
+    # Display a location
+    >>> loc_db.pretty_str(loc_key1)
+    loc_key_1
+    >>> loc_db.pretty_str(loc_key2)
+    loc_1234
+    >>> loc_db.pretty_str(loc_key3)
+    first_name
+    """
+
+    def __init__(self):
+        # Known LocKeys
+        self._loc_keys = set()
+
+        # Association tables
+        self._loc_key_to_offset = {}
+        self._loc_key_to_names = {}
+        self._name_to_loc_key = {}
+        self._offset_to_loc_key = {}
+
+        # Counter for new LocKey generation
+        self._loc_key_num = 0
+
+    def get_location_offset(self, loc_key):
+        """
+        Return the offset of @loc_key if any, None otherwise.
+        @loc_key: LocKey instance
+        """
+        assert isinstance(loc_key, LocKey)
+        return self._loc_key_to_offset.get(loc_key)
+
+    def get_location_names(self, loc_key):
+        """
+        Return the frozenset of names associated to @loc_key
+        @loc_key: LocKey instance
+        """
+        assert isinstance(loc_key, LocKey)
+        return frozenset(self._loc_key_to_names.get(loc_key, set()))
+
+    def get_name_location(self, name):
+        """
+        Return the LocKey of @name if any, None otherwise.
+        @name: target name
+        """
+        return self._name_to_loc_key.get(name)
+
+    def get_or_create_name_location(self, name):
+        """
+        Return the LocKey of @name if any, create one otherwise.
+        @name: target name
+        """
+        loc_key = self._name_to_loc_key.get(name)
+        if loc_key is not None:
+            return loc_key
+        return self.add_location(name=name)
+
+    def get_offset_location(self, offset):
+        """
+        Return the LocKey of @offset if any, None otherwise.
+        @name: target offset
+        """
+        return self._offset_to_loc_key.get(offset)
+
+    def get_or_create_offset_location(self, offset):
+        """
+        Return the LocKey of @offset if any, create one otherwise.
+        @offset: target offset
+        """
+        loc_key = self._offset_to_loc_key.get(offset)
+        if loc_key is not None:
+            return loc_key
+        return self.add_location(offset=offset)
+
+    def add_location_name(self, loc_key, name):
+        """Associate a name @name to a given @loc_key
+        @name: str instance
+        @loc_key: LocKey instance
+        """
+        assert loc_key in self._loc_keys
+        already_existing_loc = self._name_to_loc_key.get(name)
+        if already_existing_loc is not None and already_existing_loc != loc_key:
+            raise KeyError("%r is already associated to a different loc_key "
+                           "(%r)" % (name, already_existing_loc))
+        self._loc_key_to_names.setdefault(loc_key, set()).add(name)
+        self._name_to_loc_key[name] = loc_key
+
+    def remove_location_name(self, loc_key, name):
+        """Disassociate a name @name from a given @loc_key
+        Fail if @name is not already associated to @loc_key
+        @name: str instance
+        @loc_key: LocKey instance
+        """
+        assert loc_key in self._loc_keys
+        already_existing_loc = self._name_to_loc_key.get(name)
+        if already_existing_loc is None:
+            raise KeyError("%r is not already associated" % name)
+        if already_existing_loc != loc_key:
+            raise KeyError("%r is already associated to a different loc_key "
+                           "(%r)" % (name, already_existing_loc))
+        del self._name_to_loc_key[name]
+        self._loc_key_to_names[loc_key].remove(name)
+
+    def set_location_offset(self, loc_key, offset, force=False):
+        """Associate the offset @offset to an LocKey @loc_key
+
+        If @force is set, override silently. Otherwise, if an offset is already
+        associated to @loc_key, an error will be raised
+        """
+        assert loc_key in self._loc_keys
+        already_existing_loc = self.get_offset_location(offset)
+        if already_existing_loc is not None and already_existing_loc != loc_key:
+            raise KeyError("%r is already associated to a different loc_key "
+                           "(%r)" % (offset, already_existing_loc))
+        already_existing_off = self._loc_key_to_offset.get(loc_key)
+        if (already_existing_off is not None and
+            already_existing_off != offset):
+            if not force:
+                raise ValueError(
+                    "%r already has an offset (0x%x). Use 'force=True'"
+                    " for silent overriding" % (
+                        loc_key, already_existing_off
+                    ))
+            else:
+                self.unset_location_offset(loc_key)
+        self._offset_to_loc_key[offset] = loc_key
+        self._loc_key_to_offset[loc_key] = offset
+
+    def unset_location_offset(self, loc_key):
+        """Disassociate LocKey @loc_key's offset
+
+        Fail if there is already no offset associate with it
+        @loc_key: LocKey
+        """
+        assert loc_key in self._loc_keys
+        already_existing_off = self._loc_key_to_offset.get(loc_key)
+        if already_existing_off is None:
+            raise ValueError("%r already has no offset" % (loc_key))
+        del self._offset_to_loc_key[already_existing_off]
+        del self._loc_key_to_offset[loc_key]
+
+    def consistency_check(self):
+        """Ensure internal structures are consistent with each others"""
+        assert set(self._loc_key_to_names).issubset(self._loc_keys)
+        assert set(self._loc_key_to_offset).issubset(self._loc_keys)
+        assert self._loc_key_to_offset == {v: k for k, v in self._offset_to_loc_key.iteritems()}
+        assert reduce(
+            lambda x, y:x.union(y),
+            self._loc_key_to_names.itervalues(),
+            set(),
+        ) == set(self._name_to_loc_key)
+        for name, loc_key in self._name_to_loc_key.iteritems():
+            assert name in self._loc_key_to_names[loc_key]
+
+    def add_location(self, name=None, offset=None, strict=True):
+        """Add a new location in the locationDB. Returns the corresponding LocKey.
+        If @name is set, also associate a name to this new location.
+        If @offset is set, also associate an offset to this new location.
+
+        Strict mode (set by @strict, default):
+          If a location with @offset or @name already exists, an error will be
+        raised.
+        Otherwise:
+          If a location with @offset or @name already exists, the corresponding
+        LocKey will be returned.
+        """
+
+        # Deprecation handling
+        if is_int(name):
+            assert offset is None or offset == name
+            warnings.warn("Deprecated API: use 'add_location(offset=)' instead."
+                          " An additionnal 'name=' can be provided to also "
+                          "associate a name (there is no more default name)")
+            offset = name
+            name = None
+
+        # Argument cleaning
+        offset_loc_key = None
+        if offset is not None:
+            offset = int(offset)
+            offset_loc_key = self.get_offset_location(offset)
+
+        # Test for collisions
+        name_loc_key = None
+        if name is not None:
+            name_loc_key = self.get_name_location(name)
+
+        if strict:
+            if name_loc_key is not None:
+                raise ValueError("An entry for %r already exists (%r), and "
+                                 "strict mode is enabled" % (
+                                     name, name_loc_key
+                                 ))
+            if offset_loc_key is not None:
+                raise ValueError("An entry for 0x%x already exists (%r), and "
+                                 "strict mode is enabled" % (
+                                     offset, offset_loc_key
+                                 ))
+        else:
+            # Non-strict mode
+            if name_loc_key is not None:
+                known_offset = self.get_offset_location(name_loc_key)
+                if known_offset != offset:
+                    raise ValueError(
+                        "Location with name '%s' already have an offset: 0x%x "
+                        "(!= 0x%x)" % (name, offset, known_offset)
+                        )
+                # Name already known, same offset -> nothing to do
+                return name_loc_key
+
+            elif offset_loc_key is not None:
+                if name is not None:
+                    # This is an error. Check for already known name are checked above
+                    raise ValueError(
+                        "Location with offset 0x%x already exists."
+                        "To add a name to this location, use the dedicated API"
+                        "'add_location_name(%r, %r)'" % (
+                            offset_loc_key,
+                            name
+                        ))
+                # Offset already known, no name specified
+                return offset_loc_key
+
+        # No collision, this is a brand new location
+        loc_key = LocKey(self._loc_key_num)
+        self._loc_key_num += 1
+        self._loc_keys.add(loc_key)
+
+        if offset is not None:
+            assert offset not in self._offset_to_loc_key
+            self._offset_to_loc_key[offset] = loc_key
+            self._loc_key_to_offset[loc_key] = offset
+
+        if name is not None:
+            self._name_to_loc_key[name] = loc_key
+            self._loc_key_to_names[loc_key] = set([name])
+
+        return loc_key
+
+    def remove_location(self, loc_key):
+        """
+        Delete the location corresponding to @loc_key
+        @loc_key: LocKey instance
+        """
+        assert isinstance(loc_key, LocKey)
+        if loc_key not in self._loc_keys:
+            raise KeyError("Unknown loc_key %r" % loc_key)
+        names = self._loc_key_to_names.pop(loc_key, [])
+        for name in names:
+            del self._name_to_loc_key[name]
+        offset = self._loc_key_to_offset.pop(loc_key, None)
+        self._offset_to_loc_key.pop(offset, None)
+        self._loc_keys.remove(loc_key)
+
+    def pretty_str(self, loc_key):
+        """Return a human readable version of @loc_key, according to information
+        available in this LocationDB instance"""
+        names = self.get_location_names(loc_key)
+        if names:
+            return ",".join(names)
+        offset = self.get_location_offset(loc_key)
+        if offset is not None:
+            return "loc_%x" % offset
+        return str(loc_key)
+
+    @property
+    def loc_keys(self):
+        """Return all loc_keys"""
+        return self._loc_keys
+
+    @property
+    def names(self):
+        """Return all known names"""
+        return self._name_to_loc_key.keys()
+
+    @property
+    def offsets(self):
+        """Return all known offsets"""
+        return self._offset_to_loc_key.keys()
+
+    def __str__(self):
+        out = []
+        for loc_key in self._loc_keys:
+            names = self.get_location_names(loc_key)
+            offset = self.get_location_offset(loc_key)
+            out.append("%s: %s - %s" % (
+                loc_key,
+                "0x%x" % offset if offset is not None else None,
+                ",".join(names)
+            ))
+        return "\n".join(out)
+
+    def merge(self, location_db):
+        """Merge with another LocationDB @location_db
+
+        WARNING: old reference to @location_db information (such as LocKeys)
+        must be retrieved from the updated version of this instance. The
+        dedicated "get_*" APIs may be used for this task
+        """
+        # A simple merge is not doable here, because LocKey will certainly
+        # collides
+
+        for foreign_loc_key in location_db.loc_keys:
+            foreign_names = location_db.get_location_names(foreign_loc_key)
+            foreign_offset = location_db.get_location_offset(foreign_loc_key)
+            if foreign_names:
+                init_name = list(foreign_names)[0]
+            else:
+                init_name = None
+            loc_key = self.add_location(offset=foreign_offset, name=init_name,
+                                        strict=False)
+            cur_names = self.get_location_names(loc_key)
+            for name in foreign_names:
+                if name not in cur_names and name != init_name:
+                    self.add_location_name(loc_key, name=name)
+
+    def canonize_to_exprloc(self, expr):
+        """
+        If expr is ExprInt, return ExprLoc with corresponding loc_key
+        Else, return expr
+
+        @expr: Expr instance
+        """
+        if expr.is_int():
+            loc_key = self.get_or_create_offset_location(int(expr))
+            ret = ExprLoc(loc_key, expr.size)
+            return ret
+        return expr
+
+    # Deprecated APIs
+    @property
+    def items(self):
+        """Return all loc_keys"""
+        warnings.warn('DEPRECATION WARNING: use "loc_keys" instead of "items"')
+        return list(self._loc_keys)
+
+    def __getitem__(self, item):
+        warnings.warn('DEPRECATION WARNING: use "get_name_location" or '
+                      '"get_offset_location"')
+        if item in self._name_to_loc_key:
+            return self._name_to_loc_key[item]
+        if item in self._offset_to_loc_key:
+            return self._offset_to_loc_key[item]
+        raise KeyError('unknown symbol %r' % item)
+
+    def __contains__(self, item):
+        warnings.warn('DEPRECATION WARNING: use "get_name_location" or '
+                      '"get_offset_location", or ".offsets" or ".names"')
+        return item in self._name_to_loc_key or item in self._offset_to_loc_key
+
+    def loc_key_to_name(self, loc_key):
+        """[DEPRECATED API], see 'get_location_names'"""
+        warnings.warn("Deprecated API: use 'get_location_names'")
+        return sorted(self.get_location_names(loc_key))[0]
+
+    def loc_key_to_offset(self, loc_key):
+        """[DEPRECATED API], see 'get_location_offset'"""
+        warnings.warn("Deprecated API: use 'get_location_offset'")
+        return self.get_location_offset(loc_key)
+
+    def remove_loc_key(self, loc_key):
+        """[DEPRECATED API], see 'remove_location'"""
+        warnings.warn("Deprecated API: use 'remove_location'")
+        self.remove_location(loc_key)
+
+    def del_loc_key_offset(self, loc_key):
+        """[DEPRECATED API], see 'unset_location_offset'"""
+        warnings.warn("Deprecated API: use 'unset_location_offset'")
+        self.unset_location_offset(loc_key)
+
+    def getby_offset(self, offset):
+        """[DEPRECATED API], see 'get_offset_location'"""
+        warnings.warn("Deprecated API: use 'get_offset_location'")
+        return self.get_offset_location(offset)
+
+    def getby_name(self, name):
+        """[DEPRECATED API], see 'get_name_location'"""
+        warnings.warn("Deprecated API: use 'get_name_location'")
+        return self.get_name_location(name)
+
+    def getby_offset_create(self, offset):
+        """[DEPRECATED API], see 'get_or_create_offset_location'"""
+        warnings.warn("Deprecated API: use 'get_or_create_offset_location'")
+        return self.get_or_create_offset_location(offset)
+
+    def getby_name_create(self, name):
+        """[DEPRECATED API], see 'get_or_create_name_location'"""
+        warnings.warn("Deprecated API: use 'get_or_create_name_location'")
+        return self.get_or_create_name_location(name)
+
+    def rename_location(self, loc_key, newname):
+        """[DEPRECATED API], see 'add_name_location' and 'remove_location_name'
+        """
+        warnings.warn("Deprecated API: use 'add_location_name' and "
+                      "'remove_location_name'")
+        for name in self.get_location_names(loc_key):
+            self.remove_location_name(loc_key, name)
+        self.add_location_name(loc_key, name)
+
+    def set_offset(self, loc_key, offset):
+        """[DEPRECATED API], see 'set_location_offset'"""
+        warnings.warn("Deprecated API: use 'set_location_offset'")
+        self.set_location_offset(loc_key, offset, force=True)
+
+    def gen_loc_key(self):
+        """[DEPRECATED API], see 'add_location'"""
+        warnings.warn("Deprecated API: use 'add_location'")
+        return self.add_location()
+
+    def str_loc_key(self, loc_key):
+        """[DEPRECATED API], see 'pretty_str'"""
+        warnings.warn("Deprecated API: use 'pretty_str'")
+        return self.pretty_str(loc_key)