diff options
| -rw-r--r-- | miasm2/core/asmblock.py | 231 | ||||
| -rw-r--r-- | miasm2/core/locationdb.py | 453 | ||||
| -rw-r--r-- | test/core/locationdb.py | 108 | ||||
| -rwxr-xr-x | test/test_all.py | 1 |
4 files changed, 569 insertions, 224 deletions
diff --git a/miasm2/core/asmblock.py b/miasm2/core/asmblock.py index 08ff25e9..2829737e 100644 --- a/miasm2/core/asmblock.py +++ b/miasm2/core/asmblock.py @@ -12,6 +12,7 @@ from miasm2.expression.modint import moduint, modint from miasm2.core.utils import Disasm_Exception, pck from miasm2.core.graph import DiGraph, DiGraphSimplifier, MatchGraphJoker from miasm2.core.interval import interval +from miasm2.core.locationdb import LocationDB log_asmblock = logging.getLogger("asmblock") @@ -344,235 +345,17 @@ class asm_block_bad(AsmBlockBad): warnings.warn('DEPRECATION WARNING: use "AsmBlockBad" instead of "asm_block_bad"') super(asm_block_bad, self).__init__(loc_key, alignment, *args, **kwargs) +class AsmSymbolPool(LocationDB): + """[DEPRECATED API] use 'LocationDB' instead""" -class AsmSymbolPool(object): - """ - Store symbols. - - A symbol links a name to an (optional) offset - - Rules and limitations: - - two different symbols cannot have the same offset - - two different symbols cannot have the same name - - symbols manipulation (comparison, creation ...) can only be done on - symbols generated by the same symbol pool - """ - - def __init__(self): - self._loc_keys = set() - - self._loc_key_to_offset = {} - self._loc_key_to_name = {} - - self._name_to_loc_key = {} - self._offset_to_loc_key = {} - - self._loc_key_num = 0 - - def loc_key_to_offset(self, loc_key): - """ - Return offset of @loc_key, None otherwise. - @loc_key: LocKey instance - """ - assert isinstance(loc_key, LocKey) - return self._loc_key_to_offset.get(loc_key) - - def loc_key_to_name(self, loc_key): - """ - Return name of @loc_key. - @loc_key: LocKey instance - """ - assert isinstance(loc_key, LocKey) - return self._loc_key_to_name[loc_key] - - def add_location(self, name, offset=None): - """ - Create and add a location to the symbol_pool - @name: loc_key's name (never empty). If offset is None and name is int, - generate loc_key with generic name and name as offset - @offset: (optional) loc_key's offset - """ - - if is_int(name): - assert offset is None or offset == name - offset = name - name = "loc_%.16X" % (int(name) & 0xFFFFFFFFFFFFFFFF) - if offset is not None: - offset = int(offset) - - assert name - - # Test for collisions - known_loc_key = self.getby_name(name) - if known_loc_key is not None: - known_offset = self.loc_key_to_offset(known_loc_key) - if known_offset != offset: - raise ValueError( - 'symbol %s with different offset %s %s' % ( - name, offset, known_offset - ) - ) - return known_loc_key - - elif self.getby_offset(offset) is not None: - raise ValueError( - 'offset %s with different names %s' % ( - offset, - name - ) - ) - - 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 - - self._name_to_loc_key[name] = loc_key - self._loc_key_to_name[loc_key] = name - return loc_key - - def remove_loc_key(self, loc_key): - """ - Delete a @loc_key - """ - name = self._loc_key_to_name.pop(loc_key, None) - self._name_to_loc_key.pop(name, None) - - 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 del_loc_key_offset(self, loc_key): - """Unpin the @loc_key from its offset""" - offset = self._loc_keys_to_offset.pop(loc_key) - self._offset_to_loc_key.pop(offset, None) - - def getby_offset(self, offset): - """ - Retrieve loc_key using its @offset, None otherwise. - @offset: int - """ - return self._offset_to_loc_key.get(offset) - - def getby_name(self, name): - """ - Retrieve loc_key using its @name, None otherwise. - @name: str - """ - return self._name_to_loc_key.get(name) - - def getby_name_create(self, name): - """Get a loc_key from its @name, create it if it doesn't exist""" - loc_key = self.getby_name(name) - if loc_key is None: - loc_key = self.add_location(name) - return loc_key - - def getby_offset_create(self, offset): - """Get a loc_key from its @offset, create it if it doesn't exist""" - loc_key = self.getby_offset(offset) - if loc_key is None: - loc_key = self.add_location(offset) - return loc_key - - def rename_location(self, loc_key, newname): - """Rename the @loc_key name to @newname""" - if newname in self._name_to_loc_key: - raise ValueError('Symbol already known') - name = self._loc_key_to_name[loc_key] - assert name is not None - self._name_to_loc_key.pop(name) - self._loc_key_to_name[loc_key] = newname - - def set_offset(self, loc_key, offset): - """Pin the @loc_key to an @offset - Note that there is a special case when the offset is a list - it happens when offsets are recomputed in resolve_symbol* - """ - assert isinstance(loc_key, LocKey) - assert offset not in self._offset_to_loc_key - if loc_key not in self._loc_keys: - raise ValueError('Foreign loc_key %s' % loc_key) - - old_offset = self._loc_key_to_offset.pop(loc_key, None) - self._offset_to_loc_key.pop(old_offset, None) - - self._loc_key_to_offset[loc_key] = offset - self._offset_to_loc_key[offset] = loc_key - - @property - def loc_keys(self): - """Return all loc_keys""" - return self._loc_keys - - @property - def items(self): - """Return all loc_keys""" - warnings.warn('DEPRECATION WARNING: use "loc_keys" instead of "items"') - return list(self._loc_keys) - - def __str__(self): - return "".join("%s\n" % loc_key for loc_key in self._loc_keys) - - def __getitem__(self, item): - warnings.warn('DEPRECATION WARNING: use "getby_name" or "getby_offset"') - 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 "getby_name" or "getby_offset"') - return item in self._name_to_loc_key or item in self._offset_to_loc_key - - def merge(self, symbol_pool): - """Merge with another @symbol_pool""" - self._loc_keys.update(symbol_pool.loc_keys) - self._name_to_loc_key.update(symbol_pool._name_to_loc_key) - self._offset_to_loc_key.update(symbol_pool._offset_to_loc_key) - - 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.getby_offset_create(int(expr)) - ret = ExprLoc(loc_key, expr.size) - return ret - return expr - - def gen_loc_key(self): - """Generate a new unpinned loc_key""" - loc_key = self.add_location("lbl_gen_%.8X" % (self._loc_key_num)) - return loc_key - - def str_loc_key(self, loc_key): - name = self.loc_key_to_name(loc_key) - offset = self.loc_key_to_offset(loc_key) - if name is None: - name = str(loc_key) - if offset is not None: - offset = hex(offset) - out = name - if offset is not None: - out = "%s:%s" % (out, offset) - return out - + def __init__(self, *args, **kwargs): + warnings.warn("Deprecated API, use 'LocationDB' instead") + super(AsmSymbolPool, self).__init__(*args, **kwargs) class asm_symbol_pool(AsmSymbolPool): def __init__(self): - warnings.warn('DEPRECATION WARNING: use "AsmSymbolPool" instead of "asm_symbol_pool"') + warnings.warn('DEPRECATION WARNING: use "LocationDB" instead of "asm_symbol_pool"') super(asm_symbol_pool, self).__init__() 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) diff --git a/test/core/locationdb.py b/test/core/locationdb.py new file mode 100644 index 00000000..b9a5f707 --- /dev/null +++ b/test/core/locationdb.py @@ -0,0 +1,108 @@ +from miasm2.core.locationdb import LocationDB + + +# Basic tests (LocationDB description) +loc_db = LocationDB() +loc_key1 = loc_db.add_location() +loc_key2 = loc_db.add_location(offset=0x1234) +loc_key3 = loc_db.add_location(name="first_name") +loc_db.add_location_name(loc_key3, "second_name") +loc_db.set_location_offset(loc_key3, 0x5678) +loc_db.remove_location_name(loc_key3, "second_name") + +assert loc_db.get_location_offset(loc_key1) is None +assert loc_db.get_location_offset(loc_key2) == 0x1234 + +assert loc_db.pretty_str(loc_key1) == str(loc_key1) +assert loc_db.pretty_str(loc_key2) == "loc_1234" +assert loc_db.pretty_str(loc_key3) == "first_name" +loc_db.consistency_check() + +# Offset manipulation +loc_key4 = loc_db.add_location() +assert loc_db.get_location_offset(loc_key4) is None +loc_db.set_location_offset(loc_key4, 0x1122) +assert loc_db.get_location_offset(loc_key4) == 0x1122 +loc_db.unset_location_offset(loc_key4) +assert loc_db.get_location_offset(loc_key4) is None +try: + loc_db.set_location_offset(loc_key4, 0x1234) + has_raised = False +except KeyError: + has_raised = True +assert has_raised +assert loc_db.get_location_offset(loc_key4) is None +loc_db.set_location_offset(loc_key4, 0x1122) +try: + loc_db.set_location_offset(loc_key4, 0x1123) + has_raised = False +except ValueError: + has_raised = True +assert has_raised +assert loc_db.get_location_offset(loc_key4) == 0x1122 +loc_db.set_location_offset(loc_key4, 0x1123, force=True) +assert loc_db.get_location_offset(loc_key4) == 0x1123 +assert 0x1123 in loc_db.offsets +try: + loc_db.add_location(offset=0x1123) + has_raised = False +except ValueError: + has_raised = True +assert loc_db.add_location(offset=0x1123, strict=False) == loc_key4 +assert loc_db.get_offset_location(0x1123) == loc_key4 +assert loc_db.get_or_create_offset_location(0x1123) == loc_key4 +loc_key4_bis = loc_db.get_or_create_offset_location(0x1144) +assert loc_db.get_offset_location(0x1144) == loc_key4_bis +loc_db.consistency_check() + +# Names manipulation +loc_key5 = loc_db.add_location() +name1 = "name1" +name2 = "name2" +name3 = "name3" +assert len(loc_db.get_location_names(loc_key5)) == 0 +loc_db.add_location_name(loc_key5, name1) +loc_db.add_location_name(loc_key5, name2) +assert name1 in loc_db.names +assert name2 in loc_db.names +assert name1 in loc_db.get_location_names(loc_key5) +assert name2 in loc_db.get_location_names(loc_key5) +assert loc_db.get_name_location(name1) == loc_key5 +loc_db.remove_location_name(loc_key5, name1) +assert name1 not in loc_db.names +assert name1 not in loc_db.get_location_names(loc_key5) +try: + loc_db.remove_location_name(loc_key5, name1) + has_raised = False +except KeyError: + has_raised = True +try: + loc_db.add_location_name(loc_key1, name2) + has_raised = False +except KeyError: + has_raised = True +try: + loc_db.add_location(name=name2) + has_raised = False +except ValueError: + has_raised = True +assert loc_db.add_location(name=name2, strict=False) == loc_key5 +assert loc_db.get_or_create_name_location(name2) == loc_key5 +loc_key5_bis = loc_db.get_or_create_name_location(name3) +assert loc_db.get_name_location(name3) == loc_key5_bis +loc_db.consistency_check() + +# Merge +loc_db2 = LocationDB() +loc_db2.add_location(offset=0x3344) +loc_db2.add_location(name=name2) +loc_db.merge(loc_db2) +assert 0x3344 in loc_db.offsets +assert name2 in loc_db.names +loc_db.consistency_check() +assert loc_db.get_name_location(name2) == loc_key5 + +# Delete +loc_db.remove_location(loc_key5) +assert loc_db.get_name_location(name2) is None +loc_db.consistency_check() diff --git a/test/test_all.py b/test/test_all.py index b1e36573..40df315c 100755 --- a/test/test_all.py +++ b/test/test_all.py @@ -232,6 +232,7 @@ for script in ["interval.py", "parse_asm.py", "utils.py", "sembuilder.py", + "locationdb.py", "test_types.py", ]: testset += RegressionTest([script], base_dir="core") |