LocationDB: add a name -> offset method

2 files changed, 17 insertions, 0 deletions

diff --git a/miasm2/core/locationdb.py b/miasm2/core/locationdb.py
index 39c1c99a..21e2d713 100644
--- a/miasm2/core/locationdb.py
+++ b/miasm2/core/locationdb.py
@@ -39,6 +39,8 @@ class LocationDB(object):
     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)
@@ -112,6 +114,16 @@ class LocationDB(object):
             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
+        """
+        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
diff --git a/test/core/locationdb.py b/test/core/locationdb.py
index b9a5f707..61bc4563 100644
--- a/test/core/locationdb.py
+++ b/test/core/locationdb.py
@@ -92,6 +92,11 @@ 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()
 
+# Name and offset manipulation
+assert loc_db.get_name_offset(name2) is None
+assert loc_db.get_name_offset("unk_name") is None
+assert loc_db.get_name_offset("first_name") == 0x5678
+
 # Merge
 loc_db2 = LocationDB()
 loc_db2.add_location(offset=0x3344)