add getby_offset in AsmCFG

2 files changed, 10 insertions, 0 deletions

diff --git a/miasm2/core/asmblock.py b/miasm2/core/asmblock.py
index 07b8ceeb..8740aeb7 100644
--- a/miasm2/core/asmblock.py
+++ b/miasm2/core/asmblock.py
@@ -736,6 +736,14 @@ class AsmCFG(DiGraph):
                     yield predecessor
                     done.add(predecessor)
 
+    def getby_offset(self, offset):
+        """Return block containing @offset"""
+        for block in self:
+            if block.lines[0].offset <= offset < \
+                    (block.lines[-1].offset + block.lines[-1].l):
+                return block
+        return None
+
     def sanity_check(self):
         """Do sanity checks on blocks' constraints:
         * no pendings
diff --git a/test/core/asmblock.py b/test/core/asmblock.py
index eb7b54b2..7f0dbc5f 100644
--- a/test/core/asmblock.py
+++ b/test/core/asmblock.py
@@ -41,6 +41,8 @@ assert len(blocks.pendings) == 0
 assert len(blocks.nodes()) == 17
 assert len(blocks.edges2constraint) == len(blocks.edges())
 assert len(blocks.edges()) == 24
+assert blocks.getby_offset(0x63).lines[0].offset == 0x5f
+assert blocks.getby_offset(0x69).lines[0].offset == 0x69
 
 ## Convert to dot
 open("graph.dot", "w").write(blocks.dot())