Show context of pseudocodes instead whole of them (#630)

* Show context of pseudocodes instead whole of them

* Use generator expression instead of list

ida_script.py

@@ -115,6 +115,21 @@ def decompile(addr):
         return None
 
 
+def decompile_context(addr, context_lines):
+    cfunc = decompile(addr)
+    if cfunc is None:
+        return None
+    item = cfunc.body.find_closest_addr(addr)
+    y_holder = idaapi.int_pointer()
+    if not cfunc.find_item_coords(item, None, y_holder):
+        return cfunc
+    y = y_holder.value()
+    lines = cfunc.get_pseudocode()
+    retlines = (idaapi.tag_remove(lines[lnnum].line) for lnnum
+                    in range(max(0, y - context_lines),min(len(lines), y + context_lines)))
+    return '\n'.join(retlines)
+
+
 def versions():
     """Returns IDA & Python versions"""
     import sys
@@ -131,6 +146,7 @@ register_module(idautils)
 register_module(idaapi)
 server.register_function(lambda a: eval(a, globals(), locals()), 'eval')
 server.register_function(wrap(decompile)) # overwrites idaapi/ida_hexrays.decompile
+server.register_function(wrap(decompile_context), 'decompile_context')  # support context decompile
 server.register_function(versions)
 server.register_introspection_functions()
 

pwndbg/commands/context.py

@@ -299,10 +299,9 @@ def context_code():
     if not pwndbg.ida.available():
         return []
 
-    name = pwndbg.ida.GetFunctionName(pwndbg.regs.pc)
-    addr = pwndbg.ida.LocByName(name)
     # May be None when decompilation failed or user loaded wrong binary in IDA
-    code = pwndbg.ida.decompile(addr)
+    n = int(int(int(source_code_lines) / 2)) # int twice to make it a real int instead of inthook
+    code = pwndbg.ida.decompile_context(pwndbg.regs.pc, n)
     
     if code:
         return [pwndbg.ui.banner("Hexrays pseudocode")] + code.splitlines()

pwndbg/ida.py

@@ -392,6 +392,13 @@ def decompile(addr):
     return _ida.decompile(addr)
 
 
+@withHexrays
+@takes_address
+@pwndbg.memoize.reset_on_stop
+def decompile_context(pc, context_lines):
+    return _ida.decompile_context(pc, context_lines)
+
+
 @withIDA
 @pwndbg.memoize.forever
 def get_ida_versions():