miasm/core/locationdb.py
import warnings
from builtins import int as int_types
from functools import reduce
from future.utils import viewitems, viewvalues
from miasm.core.utils import printable
from miasm.expression.expression import LocKey, ExprLoc
class LocationDB(object):
"""
LocationDB is a "database" of information associated to location.
An entry in a LocationDB is uniquely identified with a LocKey.
Additional 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 additional 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
>>> loc_db.get_location_offset("first_name")
0x5678
# 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
"""
assert isinstance(name, str)
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
"""
assert isinstance(name, str)
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.
@offset: 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 get_name_offset(self, name):
"""
Return the offset of @name if any, None otherwise.
@name: target name
"""
assert isinstance(name, str)
loc_key = self.get_name_location(name)
if loc_key is None:
return None
return self.get_location_offset(loc_key)
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 isinstance(name, str)
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
assert isinstance(name, str)
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 viewitems(self._offset_to_loc_key)}
assert reduce(
lambda x, y:x.union(y),
viewvalues(self._loc_key_to_names),
set(),
) == set(self._name_to_loc_key)
for name, loc_key in viewitems(self._name_to_loc_key):
assert name in self._loc_key_to_names[loc_key]
def find_free_name(self, name):
"""
If @name is not known in DB, return it
Else append an index to it corresponding to the next unknown name
@name: string
"""
assert isinstance(name, str)
if self.get_name_location(name) is None:
return name
i = 0
while True:
new_name = "%s_%d" % (name, i)
if self.get_name_location(new_name) is None:
return new_name
i += 1
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 may be updated and will be returned.
"""
# Deprecation handling
if isinstance(name, int_types):
assert offset is None or offset == name
warnings.warn("Deprecated API: use 'add_location(offset=)' instead."
" An additional '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:
assert isinstance(name, str)
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 is None:
if offset is not None:
self.set_location_offset(name_loc_key, offset)
elif 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:
# Check for already known name are checked above
return self.add_location_name(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)
new_names = set()
for name in names:
try:
name = name.decode()
except AttributeError:
pass
new_names.add(name)
names = new_names
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 list(self._name_to_loc_key)
@property
def offsets(self):
"""Return all known offsets"""
return list(self._offset_to_loc_key)
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(printable(name) for name in 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)