Create ColorConfig class and use it in nearpc (#1317)

* Create ColorConfig class and use it in nearpc

* Address comments

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",

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=" "):

pwndbg/color/nearpc.py

@@ -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)

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. '<main+8>'
         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: