diff --git a/pwndbg/commands/checksec.py b/pwndbg/commands/checksec.py index 645ab94383b..746e05f8521 100755 --- a/pwndbg/commands/checksec.py +++ b/pwndbg/commands/checksec.py @@ -14,21 +14,8 @@ 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)) diff --git a/pwndbg/wrappers/checksec.py b/pwndbg/wrappers/checksec.py index dd1793d5f33..66e9788717d 100644 --- a/pwndbg/wrappers/checksec.py +++ b/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