From 95dd553ab7ee28866d97f919c822b2778d7c9b97 Mon Sep 17 00:00:00 2001 From: Gulshan Singh Date: Fri, 28 Jun 2024 13:37:53 -0700 Subject: [PATCH] Start migrating to Python logging and add log_level command (#2230) * Start migrating to Python logging and add log_level command * Add debug log messages to gdbinit.py --- docs/env_vars.md | 1 + gdbinit.py | 30 +++++++++++++++++++++++++++++- pwndbg/color/message.py | 10 ++++++++++ pwndbg/commands/__init__.py | 35 ++++++++++++++++++----------------- pwndbg/commands/dev.py | 18 ++++++++++++++++++ pwndbg/log.py | 20 ++++++++++++++++++++ 6 files changed, 96 insertions(+), 18 deletions(-) create mode 100644 pwndbg/log.py diff --git a/docs/env_vars.md b/docs/env_vars.md index 2a03476b..a404581c 100644 --- a/docs/env_vars.md +++ b/docs/env_vars.md @@ -7,6 +7,7 @@ Pwndbg relies on several environment variables to customize its behavior. Below - `HOME`, `XDG_CACHE_HOME`: Used by `lib.tempfile` to determine temporary file locations. - `PWNDBG_VENV_PATH`: Specifies the virtual environment path for Pwndbg. - `PWNDBG_DISABLE_COLORS`: Disables colored output in Pwndbg. +- `PWNDBG_LOGLEVEL`: Initial log level to use for log messages. - `OPENAI_API_KEY`, `ANTHROPIC_API_KEY`: Used by the `ai` command for accessing respective AI APIs. - `GITHUB_ACTIONS`, `RUN_FLAKY`: Used by `tests_commands.py` to determine the test environment. - `PWNDBG_PROFILE`: Enables profiling for benchmarking. diff --git a/gdbinit.py b/gdbinit.py index 61e8f389..7ce179af 100644 --- a/gdbinit.py +++ b/gdbinit.py @@ -2,6 +2,7 @@ from __future__ import annotations import cProfile import hashlib +import logging import os import shutil import site @@ -29,9 +30,10 @@ def hash_file(file_path: str | Path) -> str: def run_poetry_install(poetry_path: os.PathLike[str], dev: bool = False) -> Tuple[str, str, int]: - command: List[str | os.PathLike[str]] = [poetry_path, "install"] + command: List[str] = [str(poetry_path), "install"] if dev: command.extend(("--with", "dev")) + logging.debug(f"Updating deps with command: {' '.join(command)}") result = subprocess.run(command, capture_output=True, text=True) return result.stdout.strip(), result.stderr.strip(), result.returncode @@ -60,9 +62,14 @@ def update_deps(src_root: Path, venv_path: Path) -> None: poetry_lock_hash_path = venv_path / "poetry.lock.hash" current_hash = hash_file(src_root / "poetry.lock") + logging.debug(f"Current poetry.lock hash: {current_hash}") + stored_hash = None if poetry_lock_hash_path.exists(): stored_hash = poetry_lock_hash_path.read_text().strip() + logging.debug(f"Stored poetry.lock hash: {stored_hash}") + else: + logging.debug("No stored hash found") # If the hashes don't match, update the dependencies if current_hash != stored_hash: @@ -128,6 +135,21 @@ def skip_venv(src_root) -> bool: ) +def init_logger(): + log_level_env = os.environ.get("PWNDBG_LOGLEVEL", "WARNING") + log_level = getattr(logging, log_level_env.upper()) + + root_logger = logging.getLogger() + root_logger.setLevel(log_level) + + # Add a custom StreamHandler we will use to customize log message formatting. We + # configure the handler later, after pwndbg has been imported. + handler = logging.StreamHandler() + root_logger.addHandler(handler) + + return handler + + def main() -> None: profiler = cProfile.Profile() @@ -136,6 +158,8 @@ def main() -> None: start_time = time.time() profiler.enable() + handler = init_logger() + src_root = Path(__file__).parent.resolve() if not skip_venv(src_root): venv_path = get_venv_path(src_root) @@ -156,8 +180,12 @@ def main() -> None: pwndbg.dbg = pwndbg.dbg_mod.gdb.GDB() pwndbg.dbg.setup() + import pwndbg.log import pwndbg.profiling + # ColorFormatter relies on pwndbg being loaded, so we can't set it up until now + handler.setFormatter(pwndbg.log.ColorFormatter()) + pwndbg.profiling.init(profiler, start_time) if os.environ.get("PWNDBG_PROFILE") == "1": pwndbg.profiling.profiler.stop("pwndbg-load.pstats") diff --git a/pwndbg/color/message.py b/pwndbg/color/message.py index 61ec6fc0..156d8378 100644 --- a/pwndbg/color/message.py +++ b/pwndbg/color/message.py @@ -21,6 +21,8 @@ config_hint_color = theme.add_color_param( config_success_color = theme.add_color_param( "message-success-color", "green", "color of success messages" ) +config_debug_color = theme.add_color_param("message-debug-color", "blue", "color of debug messages") +config_info_color = theme.add_color_param("message-info-color", "white", "color of info messages") config_warning_color = theme.add_color_param( "message-warning-color", "yellow", "color of warning messages" ) @@ -62,6 +64,14 @@ def success(msg: object) -> str: return generateColorFunction(config.message_success_color)(msg) +def debug(msg: object) -> str: + return generateColorFunction(config.message_warning_color)(msg) + + +def info(msg: object) -> str: + return generateColorFunction(config.message_warning_color)(msg) + + def warn(msg: object) -> str: return generateColorFunction(config.message_warning_color)(msg) diff --git a/pwndbg/commands/__init__.py b/pwndbg/commands/__init__.py index 996195a7..e6cdab6b 100644 --- a/pwndbg/commands/__init__.py +++ b/pwndbg/commands/__init__.py @@ -3,6 +3,7 @@ from __future__ import annotations import argparse import functools import io +import logging from enum import Enum from typing import Any from typing import Callable @@ -22,12 +23,13 @@ import pwndbg.gdblib.kernel import pwndbg.gdblib.proc import pwndbg.gdblib.qemu import pwndbg.gdblib.regs -from pwndbg.color import message from pwndbg.gdblib.heap.ptmalloc import DebugSymsHeap from pwndbg.gdblib.heap.ptmalloc import GlibcMemoryAllocator from pwndbg.gdblib.heap.ptmalloc import HeuristicHeap from pwndbg.gdblib.heap.ptmalloc import SymbolUnresolvableError +log = logging.getLogger(__name__) + T = TypeVar("T") P = ParamSpec("P") @@ -266,9 +268,9 @@ def OnlyWithFile(function: Callable[P, T]) -> Callable[P, Optional[T]]: return function(*a, **kw) else: if pwndbg.gdblib.qemu.is_qemu(): - print(message.error("Could not determine the target binary on QEMU.")) + log.error("Could not determine the target binary on QEMU.") else: - print(message.error(f"{function.__name__}: There is no file loaded.")) + log.error(f"{function.__name__}: There is no file loaded.") return None return _OnlyWithFile @@ -280,7 +282,7 @@ def OnlyWhenQemuKernel(function: Callable[P, T]) -> Callable[P, Optional[T]]: if pwndbg.gdblib.qemu.is_qemu_kernel(): return function(*a, **kw) else: - print( + log.error( f"{function.__name__}: This command may only be run when debugging the Linux kernel in QEMU." ) return None @@ -294,7 +296,7 @@ def OnlyWhenUserspace(function: Callable[P, T]) -> Callable[P, Optional[T]]: if not pwndbg.gdblib.qemu.is_qemu_kernel(): return function(*a, **kw) else: - print( + log.error( f"{function.__name__}: This command may only be run when not debugging a QEMU kernel target." ) return None @@ -317,9 +319,8 @@ def OnlyWithArch(arch_names: List[str]) -> Callable[[Callable[P, T]], Callable[P return function(*a, **kw) else: arches_str = ", ".join(arch_names) - print( - f"%s: This command may only be run on the {arches_str} architecture(s)" - % function.__name__ + log.error( + f"{function.__name__}: This command may only be run on the {arches_str} architecture(s)" ) return None @@ -334,7 +335,7 @@ def OnlyWithKernelDebugSyms(function: Callable[P, T]) -> Callable[P, Optional[T] if pwndbg.gdblib.kernel.has_debug_syms(): return function(*a, **kw) else: - print( + log.error( f"{function.__name__}: This command may only be run when debugging a Linux kernel with debug symbols." ) return None @@ -348,7 +349,7 @@ def OnlyWhenPagingEnabled(function: Callable[P, T]) -> Callable[P, Optional[T]]: if pwndbg.gdblib.kernel.paging_enabled(): return function(*a, **kw) else: - print(f"{function.__name__}: This command may only be run when paging is enabled.") + log.error(f"{function.__name__}: This command may only be run when paging is enabled.") return None return _OnlyWhenPagingEnabled @@ -360,7 +361,7 @@ def OnlyWhenRunning(function: Callable[P, T]) -> Callable[P, Optional[T]]: if pwndbg.gdblib.proc.alive: return function(*a, **kw) else: - print(f"{function.__name__}: The program is not being run.") + log.error(f"{function.__name__}: The program is not being run.") return None return _OnlyWhenRunning @@ -373,7 +374,7 @@ def OnlyWithTcache(function: Callable[P, T]) -> Callable[P, Optional[T]]: if pwndbg.gdblib.heap.current.has_tcache(): return function(*a, **kw) else: - print( + log.error( f"{function.__name__}: This version of GLIBC was not compiled with tcache support." ) return None @@ -387,7 +388,7 @@ def OnlyWhenHeapIsInitialized(function: Callable[P, T]) -> Callable[P, Optional[ if pwndbg.gdblib.heap.current is not None and pwndbg.gdblib.heap.current.is_initialized(): return function(*a, **kw) else: - print(f"{function.__name__}: Heap is not initialized yet.") + log.error(f"{function.__name__}: Heap is not initialized yet.") return None return _OnlyWhenHeapIsInitialized @@ -400,8 +401,8 @@ def _is_statically_linked() -> bool: def _try2run_heap_command(function: Callable[P, T], *a: P.args, **kw: P.kwargs) -> T | None: - e = lambda s: print(message.error(s)) - w = lambda s: print(message.warn(s)) + e = log.error + w = log.warning # Note: We will still raise the error for developers when exception-* is set to "on" try: return function(*a, **kw) @@ -438,8 +439,8 @@ def _try2run_heap_command(function: Callable[P, T], *a: P.args, **kw: P.kwargs) def OnlyWithResolvedHeapSyms(function: Callable[P, T]) -> Callable[P, T | None]: @functools.wraps(function) def _OnlyWithResolvedHeapSyms(*a: P.args, **kw: P.kwargs) -> T | None: - e = lambda s: print(message.error(s)) - w = lambda s: print(message.warn(s)) + e = log.error + w = log.warn if ( isinstance(pwndbg.gdblib.heap.current, HeuristicHeap) and pwndbg.config.resolve_heap_via_heuristic == "auto" diff --git a/pwndbg/commands/dev.py b/pwndbg/commands/dev.py index 492f37a6..1a21f88b 100644 --- a/pwndbg/commands/dev.py +++ b/pwndbg/commands/dev.py @@ -1,6 +1,7 @@ from __future__ import annotations import argparse +import logging import pwndbg.color.message as MessageColor import pwndbg.commands @@ -63,3 +64,20 @@ def dev_dump_instruction(address=None, force_emulate=False, no_emulate=False) -> if instructions: insn = instructions[0] print(repr(insn)) + + +parser = argparse.ArgumentParser(description="Set the log level.") +parser.add_argument( + "level", + type=str, + nargs="?", + choices=["debug", "info", "warning", "error", "critical"], + default="warning", + help="The log level to set.", +) + + +@pwndbg.commands.ArgparsedCommand(parser, category=CommandCategory.DEV) +def log_level(level: str) -> None: + logging.getLogger().setLevel(getattr(logging, level.upper())) + print(f"Log level set to {level}") diff --git a/pwndbg/log.py b/pwndbg/log.py new file mode 100644 index 00000000..eccdded1 --- /dev/null +++ b/pwndbg/log.py @@ -0,0 +1,20 @@ +from __future__ import annotations + +import logging + +import pwndbg.color + + +class ColorFormatter(logging.Formatter): + log_funcs = { + logging.DEBUG: pwndbg.color.message.debug, + logging.INFO: pwndbg.color.message.info, + logging.WARNING: pwndbg.color.message.warn, + logging.ERROR: pwndbg.color.message.error, + logging.CRITICAL: pwndbg.color.message.error, + } + + def format(self, record): + log_func = self.log_funcs.get(record.levelno) + formatter = logging.Formatter(log_func("%(message)s")) + return formatter.format(record)