From 4bd4bda36dc73e0492775dbb80c3ddcc63f59873 Mon Sep 17 00:00:00 2001 From: Gulshan Singh Date: Sat, 22 Oct 2022 21:27:46 -0700 Subject: [PATCH] Create ColorConfig class and use it in nearpc (#1317) * Create ColorConfig class and use it in nearpc * Address comments --- pwndbg/arguments.py | 2 +- pwndbg/color/__init__.py | 27 ++++++++++++++++++-- pwndbg/color/nearpc.py | 53 --------------------------------------- pwndbg/commands/nearpc.py | 28 +++++++++++++++------ 4 files changed, 47 insertions(+), 63 deletions(-) delete mode 100644 pwndbg/color/nearpc.py diff --git a/pwndbg/arguments.py b/pwndbg/arguments.py index b78afe911..198dfea40 100644 --- a/pwndbg/arguments.py +++ b/pwndbg/arguments.py @@ -7,7 +7,6 @@ from capstone import CS_GRP_CALL from capstone import CS_GRP_INT import pwndbg.chain -import pwndbg.color.nearpc as N import pwndbg.constants import pwndbg.disasm import pwndbg.gdblib.arch @@ -19,6 +18,7 @@ import pwndbg.ida import pwndbg.lib.abi import pwndbg.lib.funcparser import pwndbg.lib.functions +from pwndbg.commands.nearpc import c as N ida_replacements = { "__int64": "signed long long int", diff --git a/pwndbg/color/__init__.py b/pwndbg/color/__init__.py index 0a9bdb0f0..a65b2d720 100644 --- a/pwndbg/color/__init__.py +++ b/pwndbg/color/__init__.py @@ -1,5 +1,8 @@ import os import re +from collections import namedtuple +from typing import Any +from typing import List import pwndbg.lib.memoize @@ -130,7 +133,27 @@ def generateColorFunctionInner(old, new): return wrapper -def generateColorFunction(config): +ColorParamSpec = namedtuple("ColorParamSpec", ["name", "default", "doc"]) + + +class ColorConfig: + def __init__(self, namespace: str, params: List[ColorParamSpec]): + self._namespace = namespace + self._params = {} + for param in params: + self._params[param.name] = theme.add_color_param( + f"{self._namespace}-{param.name}-color", param.default, param.doc + ) + + def __getattr__(self, attr): + param_name = attr.replace("_", "-") + if param_name in self._params: + return generateColorFunction(self._params[param_name]) + + raise AttributeError(f"ColorConfig object for {self._namespace} has no attribute '{attr}'") + + +def generateColorFunction(config: str): # the `x` here may be a config Parameter object # and if we run with disable_colors or if the config value # is empty, we need to ensure we cast it to string @@ -152,7 +175,7 @@ def strip(x): def terminateWith(x, color): - return re.sub("\x1b\\[0m", NORMAL + color, x) + return x.replace("\x1b[0m", NORMAL + color) def ljust_colored(x, length, char=" "): diff --git a/pwndbg/color/nearpc.py b/pwndbg/color/nearpc.py deleted file mode 100644 index f60549d52..000000000 --- a/pwndbg/color/nearpc.py +++ /dev/null @@ -1,53 +0,0 @@ -import pwndbg.color.theme as theme -from pwndbg.color import generateColorFunction -from pwndbg.gdblib import config - -config_symbol = theme.add_color_param( - "nearpc-symbol-color", "normal", "color for nearpc command (symbol)" -) -config_address = theme.add_color_param( - "nearpc-address-color", "normal", "color for nearpc command (address)" -) -config_prefix = theme.add_color_param( - "nearpc-prefix-color", "none", "color for nearpc command (prefix marker)" -) -config_syscall_name = theme.add_color_param( - "nearpc-syscall-name-color", "red", "color for nearpc command (resolved syscall name)" -) -config_argument = theme.add_color_param( - "nearpc-argument-color", "bold", "color for nearpc command (target argument)" -) -config_ida_anterior = theme.add_color_param( - "nearpc-ida-anterior-color", "bold", "color for nearpc command (IDA anterior)" -) -config_branch_marker = theme.add_color_param( - "nearpc-branch-marker-color", "normal", "color for nearpc command (branch marker line)" -) - - -def symbol(x): - return generateColorFunction(config.nearpc_symbol_color)(x) - - -def address(x): - return generateColorFunction(config.nearpc_address_color)(x) - - -def prefix(x): - return generateColorFunction(config.nearpc_prefix_color)(x) - - -def syscall_name(x): - return generateColorFunction(config.nearpc_syscall_name_color)(x) - - -def argument(x): - return generateColorFunction(config.nearpc_argument_color)(x) - - -def ida_anterior(x): - return generateColorFunction(config.nearpc_ida_anterior_color)(x) - - -def branch_marker(x): - return generateColorFunction(config.nearpc_branch_marker_color)(x) diff --git a/pwndbg/commands/nearpc.py b/pwndbg/commands/nearpc.py index ac9ba44d0..34fe2dada 100644 --- a/pwndbg/commands/nearpc.py +++ b/pwndbg/commands/nearpc.py @@ -7,7 +7,6 @@ import pwndbg.arguments import pwndbg.color import pwndbg.color.context as C import pwndbg.color.disasm as D -import pwndbg.color.nearpc as N import pwndbg.color.theme import pwndbg.commands.comments import pwndbg.disasm @@ -19,6 +18,8 @@ import pwndbg.gdblib.vmmap import pwndbg.ida import pwndbg.lib.functions import pwndbg.ui +from pwndbg.color import ColorConfig +from pwndbg.color import ColorParamSpec from pwndbg.color import message @@ -27,6 +28,19 @@ def ljust_padding(lst): return [s.ljust(longest_len) for s in lst] +c = ColorConfig( + "nearpc", + [ + ColorParamSpec("symbol", "normal", "color for nearpc command (symbol)"), + ColorParamSpec("address", "normal", "color for nearpc command (address)"), + ColorParamSpec("prefix", "none", "color for nearpc command (prefix marker)"), + ColorParamSpec("syscall-name", "red", "color for nearpc command (resolved syscall name)"), + ColorParamSpec("argument", "bold", "color for nearpc command (target argument)"), + ColorParamSpec("ida-anterior", "bold", "color for nearpc command (IDA anterior)"), + ColorParamSpec("branch-marker", "normal", "color for nearpc command (branch marker line)"), + ], +) + nearpc_branch_marker = pwndbg.color.theme.add_param( "nearpc-branch-marker", " ↓", "branch marker line for nearpc command" ) @@ -149,17 +163,17 @@ def nearpc(pc=None, lines=None, to_string=False, emulate=False): # or when showing current instruction for the second time show_prefix = instr.address == pc and not nearpc.repeat and first_pc prefix = " %s" % (prefix_sign if show_prefix else " " * len(prefix_sign)) - prefix = N.prefix(prefix) + prefix = c.prefix(prefix) pre = pwndbg.ida.Anterior(instr.address) if pre: - result.append(N.ida_anterior(pre)) + result.append(c.ida_anterior(pre)) # Colorize address and symbol if not highlighted # symbol is fetched from gdb and it can be e.g. '' if instr.address != pc or not pwndbg.gdblib.config.highlight_pc or nearpc.repeat: - address_str = N.address(address_str) - symbol = N.symbol(symbol) + address_str = c.address(address_str) + symbol = c.symbol(symbol) elif pwndbg.gdblib.config.highlight_pc and first_pc: prefix = C.highlight(prefix) address_str = C.highlight(address_str) @@ -171,7 +185,7 @@ def nearpc(pc=None, lines=None, to_string=False, emulate=False): # If there was a branch before this instruction which was not # contiguous, put in some ellipses. if prev and prev.address + prev.size != instr.address: - result.append(N.branch_marker("%s" % nearpc_branch_marker)) + result.append(c.branch_marker("%s" % nearpc_branch_marker)) # Otherwise if it's a branch and it *is* contiguous, just put # and empty line. @@ -183,7 +197,7 @@ def nearpc(pc=None, lines=None, to_string=False, emulate=False): if instr.address == pc: syscall_name = pwndbg.arguments.get_syscall_name(instr) if syscall_name: - line += " <%s>" % N.syscall_name("SYS_" + syscall_name) + line += " <%s>" % c.syscall_name("SYS_" + syscall_name) # For Comment Function try: