Add more IDA stuff

pwndbg/disasm.py

@@ -31,7 +31,7 @@ def get(address, instructions=1):
             try:
                 target = split[1]
                 name   = pwndbg.symbol.get(int(target, 0))
-                asm    = asm.replace(target, name)
+                asm    = asm + ' <%s>' % name
             except ValueError:
                 pass
 

pwndbg/enhance.py

@@ -62,6 +62,15 @@ def enhance(value):
     exe    = page and page.execute
     rwx    = page and page.rwx
 
+    # For the purpose of following pointers, don't display
+    # anything on the stack or heap as 'code'
+    if 'stack' in page.objfile or 'heap' in page.objfile:
+        rwx = exe = False
+
+    # If IDA doesn't think it's in a function, don't display it as code.
+    if pwndbg.ida.available() and not pwndbg.ida.GetFunctionName(value):
+        rwx = exe = False
+
     if exe:
         instr = pwndbg.disasm.get(value, 1)[0].asm
 
@@ -89,6 +98,8 @@ def enhance(value):
     if instr and 'stack' in page.objfile:
         retval = [intval, szval]
 
+
+
     # If it's RWX but a small value, don't display it as code in a chain.
     elif instr and rwx and intval0 < 0x1000:
         retval = [intval, szval]

pwndbg/ida.py

@@ -117,7 +117,6 @@ def GetBptQty():
 @returns_address
 def GetBptEA(i):
     return _ida.GetBptEA(i)
-    return request_nocache('map(GetBptEA, range(GetBptQty()))')
 
 _breakpoints=[]
 
@@ -177,17 +176,42 @@ def Auto_Jump():
 
 @withIDA
 @returns_address
+@pwndbg.memoize.reset_on_objfile
 def LocByName(name):
     return _ida.LocByName(str(name))
 
 @withIDA
 @takes_address
 @returns_address
+@pwndbg.memoize.reset_on_objfile
 def PrevHead(addr):
     return _ida.PrevHead(addr)
 
 @withIDA
 @takes_address
 @returns_address
+@pwndbg.memoize.reset_on_objfile
 def NextHead(addr):
     return _ida.NextHead(addr)
+
+@withIDA
+@takes_address
+@pwndbg.memoize.reset_on_objfile
+def GetFunctionName(addr):
+    return _ida.GetFunctionName(addr)
+
+@withIDA
+@takes_address
+@pwndbg.memoize.reset_on_objfile
+def GetFlags(addr):
+    return _ida.GetFlags(addr)
+
+@withIDA
+@pwndbg.memoize.reset_on_objfile
+def isASCII(flags):
+    return _ida.isASCII(flags)
+
+@withIDA
+@pwndbg.memoize.reset_on_objfile
+def isFunc(flags):
+    return _ida.isASCII(flags)