Fix the color and alignment of the `checksec` output (#2197)

pwndbg/commands/checksec.py

@@ -14,21 +14,8 @@ parser = argparse.ArgumentParser(
 parser.add_argument("-f", "--file", type=str, help="Specify the file to run `checksec` on.")
 
 
-def color_line(line: str) -> str:
-    return pwndbg.color.normal(
-        line.replace("*", pwndbg.color.green("*"))
-        .replace(":", f":{pwndbg.color.GREEN}")
-        .replace("No", f"{pwndbg.color.RED}No")
-    )
-
-
-def color_lines(output: str) -> str:
-    return "\n".join(map(color_line, output.split("\n")))
-
-
 @pwndbg.commands.ArgparsedCommand(parser, command_name="checksec")
 @pwndbg.commands.OnlyWithFile
 def checksec(file: str) -> None:
     local_path = file or pwndbg.gdblib.file.get_proc_exe_file()
-    output = pwndbg.wrappers.checksec.get_raw_out(local_path)
-    print(color_lines(output))
+    print(pwndbg.wrappers.checksec.get_raw_out(local_path))

pwndbg/wrappers/checksec.py

@@ -1,13 +1,38 @@
 from __future__ import annotations
 
+import contextlib
+from typing import Iterator
+
+import pwnlib.term.text
 from pwnlib.elf import ELF
 
+import pwndbg.color.message as M
+
+
+@contextlib.contextmanager
+def monkeypatch_pwnlib_term_text() -> Iterator[None]:
+    # Note: It's kinda hacky to monkeypatch pwnlib.term.text like this, but I didn't find a better way to do it.
+    # The patch is for here:
+    # https://github.com/Gallopsled/pwntools/blob/f046fdd93e154bd892332f38cfbb518de130f1f2/pwnlib/elf/elf.py#L1999-L2001
+    # This might break in the future, so need to update this patch when the implementation of pwnlib changes.
+    pwnlib.term.text.red = M.error
+    pwnlib.term.text.green = M.success
+    pwnlib.term.text.yellow = M.warn
+    yield
+    del pwnlib.term.text.red
+    del pwnlib.term.text.green
+    del pwnlib.term.text.yellow
+
 
 def get_raw_out(local_path: str) -> str:
     elf = ELF(local_path)
-    output = f"File: {elf.path}\n"
-    output += f"Arch: {elf.arch}\n"
-    output += elf.checksec()
+    # 10 is the magic number used in elf.checksec() to align the output.
+    # https://github.com/Gallopsled/pwntools/blob/f046fdd93e154bd892332f38cfbb518de130f1f2/pwnlib/elf/elf.py#L2012
+    # We might need to update this number if the implementation of elf.checksec() changes in the future.
+    output = "File:".ljust(10) + elf.path + "\n"
+    output += "Arch:".ljust(10) + elf.arch + "\n"
+    with monkeypatch_pwnlib_term_text():
+        output += elf.checksec()
     return output