Speed up disasm with color (#2334)

* Speed up disasm with color * Update CHANGELOG --------- Co-authored-by: Peace-Maker <[email protected]>
Gallopsled · Jan 17, 2024 · 5805f5e · 5805f5e
1 parent 000ca1f
commit 5805f5e
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 10 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -80,6 +80,7 @@ The table below shows which release corresponds to each branch, and what date th
 - [#2308][2308] Fix WinExec shellcraft to make sure it's 16 byte aligned
 - [#2279][2279] Make `pwn template` always set context.binary
 - [#2310][2310] Add support to start a process on Windows
+- [#2334][2334] Speed up disasm commandline tool with colored output
 
 [2242]: https://github.com/Gallopsled/pwntools/pull/2242
 [2277]: https://github.com/Gallopsled/pwntools/pull/2277
@@ -91,6 +92,7 @@ The table below shows which release corresponds to each branch, and what date th
 [2308]: https://github.com/Gallopsled/pwntools/pull/2308
 [2279]: https://github.com/Gallopsled/pwntools/pull/2279
 [2310]: https://github.com/Gallopsled/pwntools/pull/2310
+[2334]: https://github.com/Gallopsled/pwntools/pull/2334
 
 ## 4.12.0 (`beta`)
 

diff --git a/pwnlib/asm.py b/pwnlib/asm.py
@@ -860,6 +860,8 @@ def disasm(data, vma = 0, byte = True, offset = True, instructions = True):
 
 
     lines = []
+
+    # Note: those patterns are also used in pwnlib/commandline/disasm.py
     pattern = '^( *[0-9a-f]+: *)', '((?:[0-9a-f]+ )+ *)', '(.*)'
     if not byte:
         pattern = pattern[::2]

diff --git a/pwnlib/commandline/disasm.py b/pwnlib/commandline/disasm.py
@@ -3,6 +3,7 @@
 from __future__ import print_function
 
 import argparse
+import re
 import string
 import sys
 
@@ -76,18 +77,34 @@ def main(args):
         from pygments.formatters import TerminalFormatter
         from pwnlib.lexer import PwntoolsLexer
 
-        offsets = disasm(dat, vma=safeeval.const(args.address), instructions=False, byte=False)
-        bytes   = disasm(dat, vma=safeeval.const(args.address), instructions=False, offset=False)
-        instrs  = disasm(dat, vma=safeeval.const(args.address), byte=False, offset=False)
-        # instrs  = highlight(instrs, PwntoolsLexer(), TerminalFormatter())
+        dis = disasm(dat, vma=safeeval.const(args.address))
 
-        highlight_bytes = lambda t: ''.join(map(lambda x: x.replace('00', text.red('00')).replace('0a', text.red('0a')), group(2, t)))
-        for o,b,i in zip(*map(str.splitlines, (offsets, bytes, instrs))):
-            b = ' '.join(highlight_bytes(bb) for bb in b.split(' '))
-            i = highlight(i.strip(), PwntoolsLexer(), TerminalFormatter()).strip()
-            i = i.replace(',',', ')
+        # Note: those patterns are copied from disasm function
+        pattern = '^( *[0-9a-f]+: *)((?:[0-9a-f]+ )+ *)(.*)'
+        lines = []
+        for line in dis.splitlines():
+            match = re.search(pattern, line)
+            if not match:
+                # Append as one element tuple
+                lines.append((line,))
+                continue
+
+            groups = match.groups()
+            o, b, i = groups
+
+            lines.append((o, b, i))
 
-            print(o,b,i)
+
+        highlight_bytes = lambda t: ''.join(map(lambda x: x.replace('00', text.red('00')).replace('0a', text.red('0a')), group(2, t)))
+        for line in lines:
+            if len(line) == 3:
+                o, b, i = line
+                b = ' '.join(highlight_bytes(bb) for bb in b.split(' '))
+                i = highlight(i.strip(), PwntoolsLexer(), TerminalFormatter()).strip()
+                i = i.replace(',',', ')
+                print(o,b,i)
+            else:
+                print(line[0])
         return
 
     print(disasm(dat, vma=safeeval.const(args.address)))