From 676ff7448469b677892de4967440abd86d25e99b Mon Sep 17 00:00:00 2001 From: Aaron Adams Date: Wed, 22 May 2024 19:36:30 +0800 Subject: [PATCH] Add strip_colors as lib/strings API (#2164) * Add strip_colors as lib/strings API * Update strings.py --------- Co-authored-by: Disconnect3d --- pwndbg/commands/ai.py | 10 +++------- pwndbg/lib/strings.py | 8 ++++++++ 2 files changed, 11 insertions(+), 7 deletions(-) create mode 100644 pwndbg/lib/strings.py diff --git a/pwndbg/commands/ai.py b/pwndbg/commands/ai.py index c4b49c8d6..671b6d59e 100644 --- a/pwndbg/commands/ai.py +++ b/pwndbg/commands/ai.py @@ -18,6 +18,7 @@ import gdb import pwndbg import pwndbg.color.message as M import pwndbg.commands +import pwndbg.lib.strings from pwndbg.commands import CommandCategory from pwndbg.commands import context from pwndbg.gdblib import config @@ -250,7 +251,7 @@ def build_context_prompt_body(): {source} ``` """ - return strip_colors(prompt) + return pwndbg.lib.strings.strip_colors(prompt) def build_command_prompt_body(command): @@ -260,12 +261,7 @@ def build_command_prompt_body(command): output = gdb.execute(command, to_string=True) print(output) prompt += f"""\n```\n{output}\n```\n\n""" - return strip_colors(prompt) - - -def strip_colors(text): - ## Now remove all ANSI color codes from the prompt - return re.sub(r"\x1b[^m]*m", "", text) + return pwndbg.lib.strings.strip_colors(prompt) def query_openai_chat(prompt, model="gpt-3.5-turbo", max_tokens=100, temperature=0.0): diff --git a/pwndbg/lib/strings.py b/pwndbg/lib/strings.py new file mode 100644 index 000000000..6b201e1ac --- /dev/null +++ b/pwndbg/lib/strings.py @@ -0,0 +1,8 @@ +from __future__ import annotations + +import re + + +def strip_colors(text): + """Remove all ANSI color codes from the text""" + return re.sub(r"\x1b[^m]*m", "", text)