From 58280266f990e0843f4d56feaa6a922173d48dee Mon Sep 17 00:00:00 2001 From: Ads Dawson <104169244+GangGreenTemperTatum@users.noreply.github.com> Date: Mon, 30 Dec 2024 14:04:26 -0500 Subject: [PATCH 1/6] feat: hf chat completion plugin support --- burpference/api_adapters.py | 165 ++++++++++++++++++------- configs/README.md | 27 ++++ configs/huggingface_bigstar_coder.json | 16 +++ 3 files changed, 163 insertions(+), 45 deletions(-) create mode 100644 configs/huggingface_bigstar_coder.json diff --git a/burpference/api_adapters.py b/burpference/api_adapters.py index 5be8ff5..806fc11 100644 --- a/burpference/api_adapters.py +++ b/burpference/api_adapters.py @@ -19,18 +19,18 @@ def prepare_request(self, user_content, system_content=None): def process_response(self, response_data): pass + # Ollama /generate API adapter class class OllamaGenerateAPIAdapter(BaseAPIAdapter): def prepare_request(self, system_content, user_content): - prompt = "{0}\n\nUser request:\n{1}".format( - system_content, user_content) + prompt = "{0}\n\nUser request:\n{1}".format(system_content, user_content) return { "model": self.config.get("model", "llama3.2"), "prompt": prompt, "format": self.config.get("format", "json"), - "stream": self.config.get("stream", False) + "stream": self.config.get("stream", False), } def process_response(self, response_data): @@ -39,13 +39,18 @@ def process_response(self, response_data): # Ollama /chat API adapter class + class OllamaChatAPIAdapter(BaseAPIAdapter): def prepare_request(self, system_content, user_content): total_input_size = len(system_content) + len(user_content) - max_tokens = self.config.get("max_input_size", 32000) # Default to 32k if not specified + max_tokens = self.config.get( + "max_input_size", 32000 + ) # Default to 32k if not specified if total_input_size > max_tokens: - raise ValueError("Input size ({total_input_size} chars) exceeds maximum allowed ({max_tokens})") + raise ValueError( + "Input size ({total_input_size} chars) exceeds maximum allowed ({max_tokens})" + ) model = self.config.get("model", "llama3.2") quantization = self.config.get("quantization") @@ -54,8 +59,12 @@ def prepare_request(self, system_content, user_content): model = "{0}:{1}".format(model, quantization) try: - system_content = system_content.encode('utf-8', errors='replace').decode('utf-8') - user_content = user_content.encode('utf-8', errors='replace').decode('utf-8') + system_content = system_content.encode("utf-8", errors="replace").decode( + "utf-8" + ) + user_content = user_content.encode("utf-8", errors="replace").decode( + "utf-8" + ) except Exception as e: raise ValueError("Error encoding content: {str(e)}") @@ -63,14 +72,15 @@ def prepare_request(self, system_content, user_content): "model": model, "messages": [ {"role": "system", "content": system_content}, - {"role": "user", "content": user_content} + {"role": "user", "content": user_content}, ], - "stream": self.config.get("stream", False) + "stream": self.config.get("stream", False), } def process_response(self, response_data): return json.loads(response_data) + # OpenAI /v1/chat/completions API adapter class @@ -80,15 +90,15 @@ def prepare_request(self, user_content, system_content=None): "model": self.config.get("model", "gpt-4o-mini"), "messages": [ {"role": "system", "content": system_content}, - {"role": "user", "content": user_content} - ] + {"role": "user", "content": user_content}, + ], } def process_response(self, response_data): response = json.loads(response_data) - if 'choices' in response and len(response['choices']) > 0: - if 'message' in response['choices'][0]: - return response['choices'][0]['message']['content'] + if "choices" in response and len(response["choices"]) > 0: + if "message" in response["choices"][0]: + return response["choices"][0]["message"]["content"] else: raise ValueError("Unexpected response format: {response}") else: @@ -97,56 +107,61 @@ def process_response(self, response_data): def send_request(self, request_payload): headers = { "Authorization": "Bearer {0}".format(self.config.get("api_key", "")), - "Content-Type": "application/json" + "Content-Type": "application/json", } - req = urllib2.Request(self.config.get( - "host"), json.dumps(request_payload), headers=headers) - req.get_method = lambda: 'POST' + req = urllib2.Request( + self.config.get("host"), json.dumps(request_payload), headers=headers + ) + req.get_method = lambda: "POST" response = urllib2.urlopen(req) return response.read() # Anthropic /v1/messages API adapter class + class AnthropicAPIAdapter(BaseAPIAdapter): def prepare_request(self, user_content, system_content=None): return { "model": self.config.get("model", "claude-3-5-sonnet-20241022"), "max_tokens": int(self.config.get("max_tokens", 1020)), "system": system_content, - "messages": [ - {"role": "user", "content": user_content} - ] + "messages": [{"role": "user", "content": user_content}], } def send_request(self, request_payload): headers = { "x-api-key": self.config.get("headers", {}).get("x-api-key", ""), "content-type": "application/json", - "anthropic-version": self.config.get("headers", {}).get("anthropic-version", "2023-06-01") + "anthropic-version": self.config.get("headers", {}).get( + "anthropic-version", "2023-06-01" + ), } - req = urllib2.Request(self.config.get("host"), - data=json.dumps(request_payload).encode('utf-8'), - headers=headers) - req.get_method = lambda: 'POST' + req = urllib2.Request( + self.config.get("host"), + data=json.dumps(request_payload).encode("utf-8"), + headers=headers, + ) + req.get_method = lambda: "POST" try: response = urllib2.urlopen(req) return response.read() except urllib2.HTTPError as e: - error_message = e.read().decode('utf-8') + error_message = e.read().decode("utf-8") raise ValueError("HTTP Error {e.code}: {error_message}") except Exception as e: raise ValueError("Error sending request: {str(e)}") def process_response(self, response_data): response = json.loads(response_data) - if 'message' in response: - return response['message']['content'] - elif 'content' in response: - return response['content'] + if "message" in response: + return response["message"]["content"] + elif "content" in response: + return response["content"] else: raise ValueError("Unexpected response format: {response}") + # Groq openai/v1/chat/completions @@ -157,22 +172,23 @@ def prepare_request(self, user_content, system_content=None): "max_tokens": int(self.config.get("max_tokens", 1020)), "messages": [ {"role": "system", "content": system_content}, - {"role": "user", "content": user_content} - ] + {"role": "user", "content": user_content}, + ], } def process_response(self, response_data): response = json.loads(response_data) - return response['choices'][0]['message']['content'] + return response["choices"][0]["message"]["content"] def send_request(self, request_payload): headers = { "x-api-key": "{0}".format(self.config.get("api_key", "")), - "Content-Type": "application/json" + "Content-Type": "application/json", } - req = urllib2.Request(self.config.get( - "host"), json.dumps(request_payload), headers=headers) - req.get_method = lambda: 'POST' + req = urllib2.Request( + self.config.get("host"), json.dumps(request_payload), headers=headers + ) + req.get_method = lambda: "POST" response = urllib2.urlopen(req) return response.read() @@ -184,25 +200,81 @@ def prepare_request(self, system_content, user_content): "max_tokens": int(self.config.get("max_tokens", 1020)), "messages": [ {"role": "system", "content": system_content}, - {"role": "user", "content": user_content} - ] + {"role": "user", "content": user_content}, + ], } def process_response(self, response_data): response = json.loads(response_data) - return response['choices'][0]['message']['content'] + return response["choices"][0]["message"]["content"] def send_request(self, request_payload): headers = { "x-api-key": "{0}".format(self.config.get("api_key", "")), - "Content-Type": "application/json" + "Content-Type": "application/json", } - req = urllib2.Request(self.config.get( - "host"), json.dumps(request_payload), headers=headers) - req.get_method = lambda: 'POST' + req = urllib2.Request( + self.config.get("host"), json.dumps(request_payload), headers=headers + ) + req.get_method = lambda: "POST" response = urllib2.urlopen(req) return response.read() + +# HuggingFace API adapter class /chat-completion + + +class HuggingFaceAPIAdapter(BaseAPIAdapter): + def prepare_request(self, user_content, system_content=None): + messages = [] + if system_content: + messages.append({"role": "system", "content": system_content}) + messages.append({"role": "user", "content": user_content}) + + return { + "inputs": {"messages": messages}, + "parameters": { + "max_length": self.config.get("parameters", {}).get("max_length", 512), + "temperature": self.config.get("parameters", {}).get( + "temperature", 0.7 + ), + "top_p": self.config.get("parameters", {}).get("top_p", 0.9), + "repetition_penalty": self.config.get("parameters", {}).get( + "repetition_penalty", 1.2 + ), + }, + } + + def send_request(self, request_payload): + headers = self.config.get("headers", {}) + + if "Authorization" not in headers: + headers["Authorization"] = "Bearer {}".format( + self.config.get("api_key", "") + ) + + req = urllib2.Request( + self.config.get("host"), + json.dumps(request_payload).encode("utf-8"), + headers=headers, + ) + + try: + response = urllib2.urlopen(req) + return response.read() + except urllib2.HTTPError as e: + error_message = e.read().decode("utf-8") + raise ValueError("HTTP Error {}: {}".format(e.code, error_message)) + + def process_response(self, response_data): + response = json.loads(response_data) + if isinstance(response, list) and len(response) > 0: + return response[0].get("generated_text", "") + elif isinstance(response, dict): + return response.get("generated_text", str(response)) + return str(response) + + # Generic other API base adapter @@ -218,6 +290,7 @@ def process_response(self, response_data): # Function to define and load the API adapter + def get_api_adapter(config): api_type = config.get("api_type", "").lower() endpoint = config.get("host", "").lower() @@ -237,6 +310,8 @@ def get_api_adapter(config): return GroqOpenAIChatAPIAdapter(config) elif api_type == "groq-openai-stream": return GroqOpenAIChatAPIStreamAdapter(config) + elif api_type == "huggingface": + return HuggingFaceAPIAdapter(config) elif api_type == "other": return OtherAPIAdapter(config) else: diff --git a/configs/README.md b/configs/README.md index 90ff1d1..012dcb6 100644 --- a/configs/README.md +++ b/configs/README.md @@ -15,6 +15,8 @@ If you intend to fork or contribute to burpference, ensure that you have exclude - [Example Anthropic `/messages` inference with `claude-3-5-sonnet-20241022`:](#example-anthropic-messages-inference-with-claude-3-5-sonnet-20241022) - [OpenAI Inference](#openai-inference) - [Example OpenAI `/completions` inference with `gpt-4o-mini`:](#example-openai-completions-inference-with-gpt-4o-mini) + - [HuggingFace Serveless Inference](#huggingface-serveless-inference) + - [Example HuggingFace `/text-generation` inference](#example-huggingface-text-generation-inference) - [Model System Prompts](#model-system-prompts) --- @@ -96,6 +98,31 @@ In order to serve inference as part of burpference, the model must be running on } ``` +### HuggingFace Serveless Inference + +#### Example HuggingFace `/text-generation` inference + +```json +{ + "api_type": "huggingface", + "name": "HuggingFace Code Review", + "model": "bigcode/starcoder", + "host": "https://api-inference.huggingface.co/models/bigcode/starcoder", + "api_key": "YOUR_HUGGINGFACE_API_KEY", + "headers": { + "Authorization": "YOUR_HUGGINGFACE_API_KEY", + "Content-Type": "application/json" + }, + "parameters": { + "max_tokens": 512, + "temperature": 0.7, + "top_p": 0.9, + "repetition_penalty": 1.2, + "do_sample": true + } +} +``` + ## Model System Prompts By default, the system prompt sent as pretext to the model is defined [here](../prompts/proxy_prompt.txt), feel free to edit, tune and tweak as you see fit. diff --git a/configs/huggingface_bigstar_coder.json b/configs/huggingface_bigstar_coder.json new file mode 100644 index 0000000..77dff12 --- /dev/null +++ b/configs/huggingface_bigstar_coder.json @@ -0,0 +1,16 @@ +{ + "api_type": "huggingface", + "name": "HuggingFace Code Review", + "model": "bigcode/starcoder", + "host": "https://api-inference.huggingface.co/models/bigcode/starcoder", + "headers": { + "Authorization": "Bearer YOUR_HUGGINGFACE_API_KEY", + "Content-Type": "application/json" + }, + "parameters": { + "max_length": 512, + "temperature": 0.7, + "top_p": 0.9, + "repetition_penalty": 1.2 + } +} From 7af0a99fede1af3bf304ff14fc24225c05435482 Mon Sep 17 00:00:00 2001 From: Ads Dawson <104169244+GangGreenTemperTatum@users.noreply.github.com> Date: Mon, 30 Dec 2024 14:08:31 -0500 Subject: [PATCH 2/6] fix: shellcode chars --- .github/workflows/rigging_pr_description.yml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/rigging_pr_description.yml b/.github/workflows/rigging_pr_description.yml index 4ef920d..b2d36ce 100644 --- a/.github/workflows/rigging_pr_description.yml +++ b/.github/workflows/rigging_pr_description.yml @@ -22,9 +22,9 @@ jobs: run: | git fetch origin "${{ github.base_ref }}" MERGE_BASE=$(git merge-base HEAD "origin/${{ github.base_ref }}") - # Encode the diff as base64 to preserve all characters - DIFF=$(git diff "$MERGE_BASE"..HEAD | base64 -w 0) - echo "diff=$DIFF" >> "$GITHUB_OUTPUT" + # Encode the diff as base64 with explicit newline handling + DIFF=$(git diff "$MERGE_BASE" HEAD | base64 --wrap=0) + echo "diff=${DIFF}" >> "$GITHUB_OUTPUT" - uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b #v5.0.3 with: python-version: "3.11" From 9c5a930282b3477ea7a955f97f13fbdd23433ab3 Mon Sep 17 00:00:00 2001 From: Ads Dawson <104169244+GangGreenTemperTatum@users.noreply.github.com> Date: Mon, 30 Dec 2024 14:12:37 -0500 Subject: [PATCH 3/6] fix: ruff lints --- burpference/api_adapters.py | 14 +- burpference/burpference.py | 330 ++++++++++++++++++++++-------------- 2 files changed, 206 insertions(+), 138 deletions(-) diff --git a/burpference/api_adapters.py b/burpference/api_adapters.py index 806fc11..f1733eb 100644 --- a/burpference/api_adapters.py +++ b/burpference/api_adapters.py @@ -65,8 +65,8 @@ def prepare_request(self, system_content, user_content): user_content = user_content.encode("utf-8", errors="replace").decode( "utf-8" ) - except Exception as e: - raise ValueError("Error encoding content: {str(e)}") + except Exception as err: + raise ValueError("Error encoding content: {str(err)}") return { "model": model, @@ -146,11 +146,11 @@ def send_request(self, request_payload): try: response = urllib2.urlopen(req) return response.read() - except urllib2.HTTPError as e: - error_message = e.read().decode("utf-8") - raise ValueError("HTTP Error {e.code}: {error_message}") - except Exception as e: - raise ValueError("Error sending request: {str(e)}") + except urllib2.HTTPError as err: + error_msg = err.read().decode("utf-8") + raise ValueError("HTTP Error {err.code}: {error_msg}") + except Exception as err: + raise ValueError("Error sending request: {str(err)}") def process_response(self, response_data): response = json.loads(response_data) diff --git a/burpference/burpference.py b/burpference/burpference.py index 94be978..c42ac0e 100644 --- a/burpference/burpference.py +++ b/burpference/burpference.py @@ -3,9 +3,20 @@ from burp import IBurpExtender, ITab, IHttpListener from java.awt import BorderLayout, GridBagLayout, GridBagConstraints, Font from javax.swing import ( - JPanel, JTextArea, JScrollPane, - BorderFactory, JSplitPane, JButton, JComboBox, - JTable, table, ListSelectionModel, JOptionPane, JTextField, JTabbedPane) + JPanel, + JTextArea, + JScrollPane, + BorderFactory, + JSplitPane, + JButton, + JComboBox, + JTable, + table, + ListSelectionModel, + JOptionPane, + JTextField, + JTabbedPane, +) from javax.swing.table import DefaultTableCellRenderer, TableRowSorter from javax.swing.border import TitledBorder from java.util import Comparator @@ -19,7 +30,7 @@ def load_ascii_art(file_path): try: - with open(file_path, 'r') as file: + with open(file_path, "r") as file: return file.read() except IOError: return "Failed to load ASCII art" @@ -29,19 +40,17 @@ def load_ascii_art(file_path): class BurpExtender(IBurpExtender, ITab, IHttpListener): - def __init__(self): self.popupShown = False timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") try: if not os.path.exists(LOG_DIR): - os.makedirs(LOG_DIR, 0755) + os.makedirs(LOG_DIR, 0o755) except OSError as e: print("Failed to create log directory: %s" % str(e)) self.log_file_path = os.path.join( - LOG_DIR, - "burpference_log_{}.txt".format(timestamp) + LOG_DIR, "burpference_log_{}.txt".format(timestamp) ) self.config = None self.api_adapter = None @@ -66,7 +75,7 @@ def registerExtenderCallbacks(self, callbacks): outerBorder = BorderFactory.createTitledBorder( BorderFactory.createLineBorder(DREADNODE_ORANGE), - "burpference, made with <3 by @dreadnode" + "burpference, made with <3 by @dreadnode", ) outerBorder.setTitleColor(DREADNODE_PURPLE) outerBorder.setTitleFont(Font(Font.SANS_SERIF, Font.BOLD, 12)) @@ -104,8 +113,7 @@ def registerExtenderCallbacks(self, callbacks): self.logArea.setCaretColor(DREADNODE_GREY) logScrollPane = JScrollPane(self.logArea) border = BorderFactory.createTitledBorder( - BorderFactory.createLineBorder(DREADNODE_ORANGE), - "Extension Log Output" + BorderFactory.createLineBorder(DREADNODE_ORANGE), "Extension Log Output" ) border.setTitleColor(DREADNODE_ORANGE) boldFont = border.getTitleFont().deriveFont(Font.BOLD, 14) @@ -118,8 +126,7 @@ def registerExtenderCallbacks(self, callbacks): self.temp_log_messages = [] # Create a split pane for input and log - splitPane = JSplitPane(JSplitPane.VERTICAL_SPLIT, - inputPanel, logScrollPane) + splitPane = JSplitPane(JSplitPane.VERTICAL_SPLIT, inputPanel, logScrollPane) splitPane.setBackground(DARK_BACKGROUND) splitPane.setDividerSize(0) splitPane.setEnabled(False) @@ -133,7 +140,8 @@ def registerExtenderCallbacks(self, callbacks): self.historyTable.setSelectionForeground(DREADNODE_GREY) self.historyTable.setGridColor(DREADNODE_GREY) self.historyTableModel = table.DefaultTableModel( - ["#", "Timestamp", "Host", "URL", "Request", "Response"], 0) + ["#", "Timestamp", "Host", "URL", "Request", "Response"], 0 + ) self.historyTable.setModel(self.historyTableModel) class NumericComparator(Comparator): @@ -156,11 +164,11 @@ def compare(self, s1, s2): # Set selection mode and listener self.historyTable.setSelectionMode(ListSelectionModel.SINGLE_SELECTION) self.historyTable.getSelectionModel().addListSelectionListener( - self.historyTableSelectionChanged) + self.historyTableSelectionChanged + ) historyScrollPane = JScrollPane(self.historyTable) border = BorderFactory.createTitledBorder( - BorderFactory.createLineBorder(DREADNODE_ORANGE), - "HTTP History" + BorderFactory.createLineBorder(DREADNODE_ORANGE), "HTTP History" ) border.setTitleColor(DREADNODE_ORANGE) boldFont = border.getTitleFont().deriveFont(Font.BOLD, 14) @@ -178,7 +186,7 @@ def compare(self, s1, s2): requestScrollPane = JScrollPane(self.requestArea) border = BorderFactory.createTitledBorder( BorderFactory.createLineBorder(DREADNODE_ORANGE), - "Inference Request - Live View" + "Inference Request - Live View", ) border.setTitleColor(DREADNODE_ORANGE) boldFont = border.getTitleFont().deriveFont(Font.BOLD, 14) @@ -195,7 +203,7 @@ def compare(self, s1, s2): responseScrollPane = JScrollPane(self.responseArea) border = BorderFactory.createTitledBorder( BorderFactory.createLineBorder(DREADNODE_ORANGE), - "Inference Response - Live View" + "Inference Response - Live View", ) border.setTitleColor(DREADNODE_ORANGE) boldFont = border.getTitleFont().deriveFont(Font.BOLD, 14) @@ -204,7 +212,8 @@ def compare(self, s1, s2): # Create split panes diffSplitPane = JSplitPane( - JSplitPane.HORIZONTAL_SPLIT, requestScrollPane, responseScrollPane) + JSplitPane.HORIZONTAL_SPLIT, requestScrollPane, responseScrollPane + ) diffSplitPane.setBackground(DARK_BACKGROUND) diffSplitPane.setDividerSize(2) diffSplitPane.setResizeWeight(0.5) @@ -220,7 +229,7 @@ def compare(self, s1, s2): selectedRequestScrollPane = JScrollPane(self.selectedRequestArea) border = BorderFactory.createTitledBorder( BorderFactory.createLineBorder(DREADNODE_ORANGE), - "Selected HTTP Request & Response" + "Selected HTTP Request & Response", ) border.setTitleColor(DREADNODE_ORANGE) boldFont = border.getTitleFont().deriveFont(Font.BOLD, 14) @@ -237,7 +246,7 @@ def compare(self, s1, s2): selectedResponseScrollPane = JScrollPane(self.selectedResponseArea) border = BorderFactory.createTitledBorder( BorderFactory.createLineBorder(DREADNODE_ORANGE), - "Selected Inference Response" + "Selected Inference Response", ) border.setTitleColor(DREADNODE_ORANGE) boldFont = border.getTitleFont().deriveFont(Font.BOLD, 14) @@ -245,14 +254,18 @@ def compare(self, s1, s2): selectedResponseScrollPane.setBorder(border) selectedDiffSplitPane = JSplitPane( - JSplitPane.HORIZONTAL_SPLIT, selectedRequestScrollPane, selectedResponseScrollPane) + JSplitPane.HORIZONTAL_SPLIT, + selectedRequestScrollPane, + selectedResponseScrollPane, + ) selectedDiffSplitPane.setBackground(DARK_BACKGROUND) selectedDiffSplitPane.setDividerSize(2) selectedDiffSplitPane.setResizeWeight(0.5) # Main split pane for history and selected entry mainSplitPane = JSplitPane( - JSplitPane.VERTICAL_SPLIT, historyScrollPane, selectedDiffSplitPane) + JSplitPane.VERTICAL_SPLIT, historyScrollPane, selectedDiffSplitPane + ) mainSplitPane.setBackground(DARK_BACKGROUND) mainSplitPane.setDividerSize(2) mainSplitPane.setResizeWeight(0.5) @@ -284,10 +297,8 @@ def compare(self, s1, s2): self.log_message("Extension initialized and running.") callbacks.printOutput(SQUID_ASCII + "\n\n") - callbacks.printOutput( - "Yer configs be stowed and ready from " + CONFIG_DIR) - callbacks.printOutput( - "\nNow ye be speakin' the pirate's tongue, savvy?") + callbacks.printOutput("Yer configs be stowed and ready from " + CONFIG_DIR) + callbacks.printOutput("\nNow ye be speakin' the pirate's tongue, savvy?") self.promptForConfiguration() self.applyDarkTheme(self.tabbedPane) @@ -301,10 +312,8 @@ def getUiComponent(self): def stopExtension(self, event): if self.is_running: self.is_running = False - self._callbacks.removeHttpListener( - self) - self.log_message( - "Extension stopped. No further traffic will be processed.") + self._callbacks.removeHttpListener(self) + self.log_message("Extension stopped. No further traffic will be processed.") self.updateUIState() def updateUIState(self): @@ -315,7 +324,8 @@ def updateUIState(self): self.stopButton.addActionListener(self.stopExtension) else: self.stopButton.setText( - "Extension Stopped - Unload and reload if required, doing so will remove displayed logs and state") + "Extension Stopped - Unload and reload if required, doing so will remove displayed logs and state" + ) for listener in self.stopButton.getActionListeners(): self.stopButton.removeActionListener(listener) @@ -323,17 +333,18 @@ def loadConfigFiles(self): if not os.path.exists(CONFIG_DIR): self.log_message("Config directory not found: {CONFIG_DIR}") return [] - return [f for f in os.listdir(CONFIG_DIR) if f.endswith('.json')] + return [f for f in os.listdir(CONFIG_DIR) if f.endswith(".json")] def loadConfiguration(self, event): selected_config = self.configSelector.getSelectedItem() config_path = os.path.join(CONFIG_DIR, selected_config) if os.path.exists(config_path): try: - with open(config_path, 'r') as config_file: + with open(config_path, "r") as config_file: self.config = json.load(config_file) - self.log_message("Loaded configuration: %s" % - json.dumps(self.config, indent=2)) + self.log_message( + "Loaded configuration: %s" % json.dumps(self.config, indent=2) + ) try: self.api_adapter = get_api_adapter(self.config) self.log_message("API adapter initialized successfully") @@ -342,21 +353,21 @@ def loadConfiguration(self, event): self.api_adapter = None except Exception as e: self.log_message( - "Unexpected error initializing API adapter: %s" % str(e)) + "Unexpected error initializing API adapter: %s" % str(e) + ) self.api_adapter = None except ValueError as e: self.log_message( - "Error parsing JSON in configuration file: %s" % str(e)) + "Error parsing JSON in configuration file: %s" % str(e) + ) self.config = None self.api_adapter = None except Exception as e: - self.log_message( - "Unexpected error loading configuration: %s" % str(e)) + self.log_message("Unexpected error loading configuration: %s" % str(e)) self.config = None self.api_adapter = None else: - self.log_message( - "Configuration file %s not found." % selected_config) + self.log_message("Configuration file %s not found." % selected_config) self.config = None self.api_adapter = None @@ -372,18 +383,19 @@ def create_inference_logger_tab(self): self.inferenceLogTable.setGridColor(DREADNODE_GREY) self.inferenceLogTableModel = table.DefaultTableModel( - ["Timestamp", "API Endpoint", "Proxy Request", "Model Response", "Status"], 0) + ["Timestamp", "API Endpoint", "Proxy Request", "Model Response", "Status"], + 0, + ) self.inferenceLogTable.setModel(self.inferenceLogTableModel) - self.inferenceLogTable.setSelectionMode( - ListSelectionModel.SINGLE_SELECTION) + self.inferenceLogTable.setSelectionMode(ListSelectionModel.SINGLE_SELECTION) self.inferenceLogTable.getSelectionModel().addListSelectionListener( - self.inferenceLogSelectionChanged) + self.inferenceLogSelectionChanged + ) inferenceLogScrollPane = JScrollPane(self.inferenceLogTable) border = BorderFactory.createTitledBorder( - BorderFactory.createLineBorder(DREADNODE_ORANGE), - "API Requests Log" + BorderFactory.createLineBorder(DREADNODE_ORANGE), "API Requests Log" ) border.setTitleColor(DREADNODE_ORANGE) border.setTitleFont(border.getTitleFont().deriveFont(Font.BOLD, 14)) @@ -396,10 +408,12 @@ def create_inference_logger_tab(self): self.inferenceRequestDetail.setBackground(LIGHTER_BACKGROUND) self.inferenceRequestDetail.setForeground(DREADNODE_ORANGE) requestDetailPane = JScrollPane(self.inferenceRequestDetail) - requestDetailPane.setBorder(BorderFactory.createTitledBorder( - BorderFactory.createLineBorder(DREADNODE_ORANGE), - "Inference Request Detail" - )) + requestDetailPane.setBorder( + BorderFactory.createTitledBorder( + BorderFactory.createLineBorder(DREADNODE_ORANGE), + "Inference Request Detail", + ) + ) self.inferenceResponseDetail = JTextArea(10, 30) self.inferenceResponseDetail.setEditable(False) @@ -408,24 +422,22 @@ def create_inference_logger_tab(self): self.inferenceResponseDetail.setBackground(LIGHTER_BACKGROUND) self.inferenceResponseDetail.setForeground(DREADNODE_ORANGE) responseDetailPane = JScrollPane(self.inferenceResponseDetail) - responseDetailPane.setBorder(BorderFactory.createTitledBorder( - BorderFactory.createLineBorder(DREADNODE_ORANGE), - "Inference Response Detail" - )) + responseDetailPane.setBorder( + BorderFactory.createTitledBorder( + BorderFactory.createLineBorder(DREADNODE_ORANGE), + "Inference Response Detail", + ) + ) # Create split pane for details detailsSplitPane = JSplitPane( - JSplitPane.HORIZONTAL_SPLIT, - requestDetailPane, - responseDetailPane + JSplitPane.HORIZONTAL_SPLIT, requestDetailPane, responseDetailPane ) detailsSplitPane.setResizeWeight(0.5) # Create main split pane mainSplitPane = JSplitPane( - JSplitPane.VERTICAL_SPLIT, - inferenceLogScrollPane, - detailsSplitPane + JSplitPane.VERTICAL_SPLIT, inferenceLogScrollPane, detailsSplitPane ) mainSplitPane.setResizeWeight(0.5) @@ -438,8 +450,12 @@ def inferenceLogSelectionChanged(self, event): if selectedRow != -1: try: # Get the request and response data from columns - request = self.inferenceLogTableModel.getValueAt(selectedRow, 2) # Proxy Request column - response = self.inferenceLogTableModel.getValueAt(selectedRow, 3) # Model Response column + request = self.inferenceLogTableModel.getValueAt( + selectedRow, 2 + ) # Proxy Request column + response = self.inferenceLogTableModel.getValueAt( + selectedRow, 3 + ) # Model Response column # Format the request JSON try: @@ -458,11 +474,15 @@ def inferenceLogSelectionChanged(self, event): else: response_obj = json.loads(response) - if isinstance(response_obj, dict) and 'message' in response_obj and 'content' in response_obj['message']: + if ( + isinstance(response_obj, dict) + and "message" in response_obj + and "content" in response_obj["message"] + ): # Get just the content string and handle unicode - content = response_obj['message']['content'] + content = response_obj["message"]["content"] # Handle unicode strings and newlines - formatted_response = content.replace('\\n', '\n') + formatted_response = content.replace("\\n", "\n") else: formatted_response = json.dumps(response_obj, indent=2) except (ValueError, AttributeError, TypeError): @@ -490,8 +510,8 @@ def historyTableSelectionChanged(self, event): http_pair = json.loads(http_pair_json) self.selectedRequestArea.setText(json.dumps(http_pair, indent=2)) - model_text = model_analysis.strip('"').decode('unicode_escape') - model_text = model_text.replace('\\n', '\n') + model_text = model_analysis.strip('"').decode("unicode_escape") + model_text = model_text.replace("\\n", "\n") self.selectedResponseArea.setText(model_text) self.selectedRequestArea.setCaretPosition(0) @@ -507,12 +527,19 @@ def historyTableSelectionChanged(self, event): def colorizeHistoryTable(self): renderer = self.SeverityCellRenderer() for column in range(self.historyTable.getColumnCount()): - self.historyTable.getColumnModel().getColumn(column).setCellRenderer(renderer) + self.historyTable.getColumnModel().getColumn(column).setCellRenderer( + renderer + ) class SeverityCellRenderer(DefaultTableCellRenderer): - def getTableCellRendererComponent(self, table, value, isSelected, hasFocus, row, column): - component = super(BurpExtender.SeverityCellRenderer, self).getTableCellRendererComponent( - table, value, isSelected, hasFocus, row, column) + def getTableCellRendererComponent( + self, table, value, isSelected, hasFocus, row, column + ): + component = super( + BurpExtender.SeverityCellRenderer, self + ).getTableCellRendererComponent( + table, value, isSelected, hasFocus, row, column + ) component.setForeground(DREADNODE_GREY) @@ -557,17 +584,16 @@ def log_message(self, message): self.temp_log_messages.append(log_entry) else: self.logArea.append(log_entry) - self.logArea.setCaretPosition( - self.logArea.getDocument().getLength()) + self.logArea.setCaretPosition(self.logArea.getDocument().getLength()) try: # Try to create/write to log file with explicit permissions log_dir = os.path.dirname(self.log_file_path) if not os.path.exists(log_dir): - os.makedirs(log_dir, 0755) # Python2 octal notation + os.makedirs(log_dir, 0o755) # Open with explicit write permissions - with open(self.log_file_path, 'a+') as log_file: + with open(self.log_file_path, "a+") as log_file: log_file.write(log_entry) except (IOError, OSError) as e: print("Warning: Could not write to log file: %s" % str(e)) @@ -594,7 +620,7 @@ def applyDarkTheme(self, component): component.setDividerSize(2) component.setDividerLocation(0.5) - if hasattr(component, 'getComponents'): + if hasattr(component, "getComponents"): for child in component.getComponents(): self.applyDarkTheme(child) @@ -602,8 +628,10 @@ def processHttpMessage(self, toolFlag, messageIsRequest, messageInfo): if not self.is_running: return if not self.api_adapter: - if not hasattr(self, '_no_adapter_logged'): - self.log_message("No API adapter configured. Please select a configuration file.") + if not hasattr(self, "_no_adapter_logged"): + self.log_message( + "No API adapter configured. Please select a configuration file." + ) self._no_adapter_logged = True return if messageIsRequest: @@ -614,10 +642,9 @@ def processHttpMessage(self, toolFlag, messageIsRequest, messageInfo): response = messageInfo # Filter MIME content types to reduce noise and exceeding tokens - responseInfo = self._helpers.analyzeResponse( - response.getResponse()) + responseInfo = self._helpers.analyzeResponse(response.getResponse()) contentType = responseInfo.getStatedMimeType().lower() - if contentType in ['css', 'image', 'script', 'video', 'audio', 'font']: + if contentType in ["css", "image", "script", "video", "audio", "font"]: return # Filter request size @@ -625,20 +652,24 @@ def processHttpMessage(self, toolFlag, messageIsRequest, messageInfo): return try: analyzed_request = self._helpers.analyzeRequest(request) - analyzed_response = self._helpers.analyzeResponse(response.getResponse()) + analyzed_response = self._helpers.analyzeResponse( + response.getResponse() + ) # Get the body bytes - request_body = request.getRequest()[analyzed_request.getBodyOffset():] - response_body = response.getResponse()[analyzed_response.getBodyOffset():] + request_body = request.getRequest()[analyzed_request.getBodyOffset() :] + response_body = response.getResponse()[ + analyzed_response.getBodyOffset() : + ] # Try to safely decode bodies, fallback to hex for binary data try: - request_body_str = request_body.tostring().decode('utf-8') + request_body_str = request_body.tostring().decode("utf-8") except UnicodeDecodeError: request_body_str = "" % len(request_body) try: - response_body_str = response_body.tostring().decode('utf-8') + response_body_str = response_body.tostring().decode("utf-8") except UnicodeDecodeError: response_body_str = "" % len(response_body) @@ -655,33 +686,41 @@ def processHttpMessage(self, toolFlag, messageIsRequest, messageInfo): request_data = { "method": analyzed_request.getMethod(), "url": str(request.getUrl()), - "headers": dict(header.split(': ', 1) for header in analyzed_request.getHeaders()[1:] if ': ' in header), - "body": request_body_str + "headers": dict( + header.split(": ", 1) + for header in analyzed_request.getHeaders()[1:] + if ": " in header + ), + "body": request_body_str, } # Package the original HTTP response response_data = { "status_code": analyzed_response.getStatusCode(), - "headers": dict(header.split(': ', 1) for header in analyzed_response.getHeaders()[1:] if ': ' in header), - "body": response_body_str + "headers": dict( + header.split(": ", 1) + for header in analyzed_response.getHeaders()[1:] + if ": " in header + ), + "body": response_body_str, } # Create the HTTP pair to send to the model - http_pair = { - "request": request_data, - "response": response_data - } + http_pair = {"request": request_data, "response": response_data} # Load prompt template for system role if os.path.exists(PROXY_PROMPT): - with open(PROXY_PROMPT, 'r') as prompt_file: + with open(PROXY_PROMPT, "r") as prompt_file: system_content = prompt_file.read().strip() # Only log if there's an issue or if it's different from last time - if not hasattr(self, '_last_system_content') or system_content != self._last_system_content: + if ( + not hasattr(self, "_last_system_content") + or system_content != self._last_system_content + ): self.log_message("Custom prompt loaded from " + PROXY_PROMPT) self._last_system_content = system_content else: - if not hasattr(self, '_prompt_missing_logged'): + if not hasattr(self, "_prompt_missing_logged"): self.log_message("No prompt file found. Using default prompt.") self._prompt_missing_logged = True system_content = "Examine this request and response pair for any security issues:" @@ -689,7 +728,7 @@ def processHttpMessage(self, toolFlag, messageIsRequest, messageInfo): # Prepare the request using the adapter remote_request = self.api_adapter.prepare_request( user_content=json.dumps(http_pair, indent=2), - system_content=system_content + system_content=system_content, ) # Send request to model @@ -700,18 +739,22 @@ def processHttpMessage(self, toolFlag, messageIsRequest, messageInfo): try: response = urllib2.urlopen(req, json.dumps(remote_request)) response_data = response.read() - processed_response = self.api_adapter.process_response(response_data) + processed_response = self.api_adapter.process_response( + response_data + ) status = "Success" except urllib2.HTTPError as e: - error_message = e.read().decode('utf-8') + error_message = e.read().decode("utf-8") error_details = { "status_code": e.code, "headers": dict(e.headers), - "body": error_message + "body": error_message, } processed_response = json.dumps(error_details, indent=2) - self.log_message("API Error Response:\nStatus Code: %d\nHeaders: %s\nBody: %s" % - (e.code, dict(e.headers), error_message)) + self.log_message( + "API Error Response:\nStatus Code: %d\nHeaders: %s\nBody: %s" + % (e.code, dict(e.headers), error_message) + ) status = "Failed (%d)" % e.code except Exception as e: processed_response = "Error: %s" % str(e) @@ -721,51 +764,75 @@ def processHttpMessage(self, toolFlag, messageIsRequest, messageInfo): # Always log to inference logger if self.is_running: timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S") - self.inferenceLogTableModel.addRow([ - timestamp, - self.config.get("host", ""), - json.dumps(remote_request), - processed_response, - status - ]) + self.inferenceLogTableModel.addRow( + [ + timestamp, + self.config.get("host", ""), + json.dumps(remote_request), + processed_response, + status, + ] + ) # Only update the main UI if we got a successful response if status == "Success" and self.is_running: # Add to history table with metadata - self.historyTableModel.addRow([ - table_metadata.get("id", ""), - timestamp, - table_metadata.get("host", ""), - table_metadata.get("url", ""), - json.dumps(http_pair, indent=2), # Store the HTTP request/response pair - json.dumps(processed_response, indent=2) # Store the model's analysis - ]) + self.historyTableModel.addRow( + [ + table_metadata.get("id", ""), + timestamp, + table_metadata.get("host", ""), + table_metadata.get("url", ""), + json.dumps( + http_pair, indent=2 + ), # Store the HTTP request/response pair + json.dumps( + processed_response, indent=2 + ), # Store the model's analysis + ] + ) self.colorizeHistoryTable() - self.requestArea.append("\n\n=== Request #" + str(self.request_counter) + " ===\n") + self.requestArea.append( + "\n\n=== Request #" + str(self.request_counter) + " ===\n" + ) try: # Format the request nicely formatted_request = json.dumps(http_pair, indent=2) - formatted_request = formatted_request.replace('\\n', '\n') + formatted_request = formatted_request.replace("\\n", "\n") formatted_request = formatted_request.replace('\\"', '"') self.requestArea.append(formatted_request) except Exception as e: self.requestArea.append(str(http_pair)) - self.requestArea.setCaretPosition(self.requestArea.getDocument().getLength()) + self.requestArea.setCaretPosition( + self.requestArea.getDocument().getLength() + ) - self.responseArea.append("\n\n=== Response #" + str(self.request_counter) + " ===\n") + self.responseArea.append( + "\n\n=== Response #" + str(self.request_counter) + " ===\n" + ) try: # Format the response nicely - if isinstance(processed_response, dict) and 'message' in processed_response and 'content' in processed_response['message']: - formatted_response = processed_response['message']['content'] + if ( + isinstance(processed_response, dict) + and "message" in processed_response + and "content" in processed_response["message"] + ): + formatted_response = processed_response["message"][ + "content" + ] else: - formatted_response = json.dumps(processed_response, indent=2) - formatted_response = formatted_response.replace('\\n', '\n') + formatted_response = json.dumps( + processed_response, indent=2 + ) + formatted_response = formatted_response.replace("\\n", "\n") formatted_response = formatted_response.replace('\\"', '"') self.responseArea.append(formatted_response) except Exception as e: self.responseArea.append(str(processed_response)) - self.responseArea.setCaretPosition(self.responseArea.getDocument().getLength()) + self.responseArea.setCaretPosition( + self.responseArea.getDocument().getLength() + ) except Exception as e: self.log_message("Error processing request: %s" % str(e)) @@ -776,4 +843,5 @@ def promptForConfiguration(self): "Select a configuration file to load in the burpference extension" " tab and go brrr", "burpference Configuration Required", - JOptionPane.INFORMATION_MESSAGE) \ No newline at end of file + JOptionPane.INFORMATION_MESSAGE, + ) From 26ba6160ff09212df5be9c2abe1fd84ecd81a767 Mon Sep 17 00:00:00 2001 From: Ads Dawson <104169244+GangGreenTemperTatum@users.noreply.github.com> Date: Mon, 30 Dec 2024 14:16:47 -0500 Subject: [PATCH 4/6] Revert "fix: ruff lints" This reverts commit 9c5a930282b3477ea7a955f97f13fbdd23433ab3. --- burpference/api_adapters.py | 14 +- burpference/burpference.py | 330 ++++++++++++++---------------------- 2 files changed, 138 insertions(+), 206 deletions(-) diff --git a/burpference/api_adapters.py b/burpference/api_adapters.py index f1733eb..806fc11 100644 --- a/burpference/api_adapters.py +++ b/burpference/api_adapters.py @@ -65,8 +65,8 @@ def prepare_request(self, system_content, user_content): user_content = user_content.encode("utf-8", errors="replace").decode( "utf-8" ) - except Exception as err: - raise ValueError("Error encoding content: {str(err)}") + except Exception as e: + raise ValueError("Error encoding content: {str(e)}") return { "model": model, @@ -146,11 +146,11 @@ def send_request(self, request_payload): try: response = urllib2.urlopen(req) return response.read() - except urllib2.HTTPError as err: - error_msg = err.read().decode("utf-8") - raise ValueError("HTTP Error {err.code}: {error_msg}") - except Exception as err: - raise ValueError("Error sending request: {str(err)}") + except urllib2.HTTPError as e: + error_message = e.read().decode("utf-8") + raise ValueError("HTTP Error {e.code}: {error_message}") + except Exception as e: + raise ValueError("Error sending request: {str(e)}") def process_response(self, response_data): response = json.loads(response_data) diff --git a/burpference/burpference.py b/burpference/burpference.py index c42ac0e..94be978 100644 --- a/burpference/burpference.py +++ b/burpference/burpference.py @@ -3,20 +3,9 @@ from burp import IBurpExtender, ITab, IHttpListener from java.awt import BorderLayout, GridBagLayout, GridBagConstraints, Font from javax.swing import ( - JPanel, - JTextArea, - JScrollPane, - BorderFactory, - JSplitPane, - JButton, - JComboBox, - JTable, - table, - ListSelectionModel, - JOptionPane, - JTextField, - JTabbedPane, -) + JPanel, JTextArea, JScrollPane, + BorderFactory, JSplitPane, JButton, JComboBox, + JTable, table, ListSelectionModel, JOptionPane, JTextField, JTabbedPane) from javax.swing.table import DefaultTableCellRenderer, TableRowSorter from javax.swing.border import TitledBorder from java.util import Comparator @@ -30,7 +19,7 @@ def load_ascii_art(file_path): try: - with open(file_path, "r") as file: + with open(file_path, 'r') as file: return file.read() except IOError: return "Failed to load ASCII art" @@ -40,17 +29,19 @@ def load_ascii_art(file_path): class BurpExtender(IBurpExtender, ITab, IHttpListener): + def __init__(self): self.popupShown = False timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") try: if not os.path.exists(LOG_DIR): - os.makedirs(LOG_DIR, 0o755) + os.makedirs(LOG_DIR, 0755) except OSError as e: print("Failed to create log directory: %s" % str(e)) self.log_file_path = os.path.join( - LOG_DIR, "burpference_log_{}.txt".format(timestamp) + LOG_DIR, + "burpference_log_{}.txt".format(timestamp) ) self.config = None self.api_adapter = None @@ -75,7 +66,7 @@ def registerExtenderCallbacks(self, callbacks): outerBorder = BorderFactory.createTitledBorder( BorderFactory.createLineBorder(DREADNODE_ORANGE), - "burpference, made with <3 by @dreadnode", + "burpference, made with <3 by @dreadnode" ) outerBorder.setTitleColor(DREADNODE_PURPLE) outerBorder.setTitleFont(Font(Font.SANS_SERIF, Font.BOLD, 12)) @@ -113,7 +104,8 @@ def registerExtenderCallbacks(self, callbacks): self.logArea.setCaretColor(DREADNODE_GREY) logScrollPane = JScrollPane(self.logArea) border = BorderFactory.createTitledBorder( - BorderFactory.createLineBorder(DREADNODE_ORANGE), "Extension Log Output" + BorderFactory.createLineBorder(DREADNODE_ORANGE), + "Extension Log Output" ) border.setTitleColor(DREADNODE_ORANGE) boldFont = border.getTitleFont().deriveFont(Font.BOLD, 14) @@ -126,7 +118,8 @@ def registerExtenderCallbacks(self, callbacks): self.temp_log_messages = [] # Create a split pane for input and log - splitPane = JSplitPane(JSplitPane.VERTICAL_SPLIT, inputPanel, logScrollPane) + splitPane = JSplitPane(JSplitPane.VERTICAL_SPLIT, + inputPanel, logScrollPane) splitPane.setBackground(DARK_BACKGROUND) splitPane.setDividerSize(0) splitPane.setEnabled(False) @@ -140,8 +133,7 @@ def registerExtenderCallbacks(self, callbacks): self.historyTable.setSelectionForeground(DREADNODE_GREY) self.historyTable.setGridColor(DREADNODE_GREY) self.historyTableModel = table.DefaultTableModel( - ["#", "Timestamp", "Host", "URL", "Request", "Response"], 0 - ) + ["#", "Timestamp", "Host", "URL", "Request", "Response"], 0) self.historyTable.setModel(self.historyTableModel) class NumericComparator(Comparator): @@ -164,11 +156,11 @@ def compare(self, s1, s2): # Set selection mode and listener self.historyTable.setSelectionMode(ListSelectionModel.SINGLE_SELECTION) self.historyTable.getSelectionModel().addListSelectionListener( - self.historyTableSelectionChanged - ) + self.historyTableSelectionChanged) historyScrollPane = JScrollPane(self.historyTable) border = BorderFactory.createTitledBorder( - BorderFactory.createLineBorder(DREADNODE_ORANGE), "HTTP History" + BorderFactory.createLineBorder(DREADNODE_ORANGE), + "HTTP History" ) border.setTitleColor(DREADNODE_ORANGE) boldFont = border.getTitleFont().deriveFont(Font.BOLD, 14) @@ -186,7 +178,7 @@ def compare(self, s1, s2): requestScrollPane = JScrollPane(self.requestArea) border = BorderFactory.createTitledBorder( BorderFactory.createLineBorder(DREADNODE_ORANGE), - "Inference Request - Live View", + "Inference Request - Live View" ) border.setTitleColor(DREADNODE_ORANGE) boldFont = border.getTitleFont().deriveFont(Font.BOLD, 14) @@ -203,7 +195,7 @@ def compare(self, s1, s2): responseScrollPane = JScrollPane(self.responseArea) border = BorderFactory.createTitledBorder( BorderFactory.createLineBorder(DREADNODE_ORANGE), - "Inference Response - Live View", + "Inference Response - Live View" ) border.setTitleColor(DREADNODE_ORANGE) boldFont = border.getTitleFont().deriveFont(Font.BOLD, 14) @@ -212,8 +204,7 @@ def compare(self, s1, s2): # Create split panes diffSplitPane = JSplitPane( - JSplitPane.HORIZONTAL_SPLIT, requestScrollPane, responseScrollPane - ) + JSplitPane.HORIZONTAL_SPLIT, requestScrollPane, responseScrollPane) diffSplitPane.setBackground(DARK_BACKGROUND) diffSplitPane.setDividerSize(2) diffSplitPane.setResizeWeight(0.5) @@ -229,7 +220,7 @@ def compare(self, s1, s2): selectedRequestScrollPane = JScrollPane(self.selectedRequestArea) border = BorderFactory.createTitledBorder( BorderFactory.createLineBorder(DREADNODE_ORANGE), - "Selected HTTP Request & Response", + "Selected HTTP Request & Response" ) border.setTitleColor(DREADNODE_ORANGE) boldFont = border.getTitleFont().deriveFont(Font.BOLD, 14) @@ -246,7 +237,7 @@ def compare(self, s1, s2): selectedResponseScrollPane = JScrollPane(self.selectedResponseArea) border = BorderFactory.createTitledBorder( BorderFactory.createLineBorder(DREADNODE_ORANGE), - "Selected Inference Response", + "Selected Inference Response" ) border.setTitleColor(DREADNODE_ORANGE) boldFont = border.getTitleFont().deriveFont(Font.BOLD, 14) @@ -254,18 +245,14 @@ def compare(self, s1, s2): selectedResponseScrollPane.setBorder(border) selectedDiffSplitPane = JSplitPane( - JSplitPane.HORIZONTAL_SPLIT, - selectedRequestScrollPane, - selectedResponseScrollPane, - ) + JSplitPane.HORIZONTAL_SPLIT, selectedRequestScrollPane, selectedResponseScrollPane) selectedDiffSplitPane.setBackground(DARK_BACKGROUND) selectedDiffSplitPane.setDividerSize(2) selectedDiffSplitPane.setResizeWeight(0.5) # Main split pane for history and selected entry mainSplitPane = JSplitPane( - JSplitPane.VERTICAL_SPLIT, historyScrollPane, selectedDiffSplitPane - ) + JSplitPane.VERTICAL_SPLIT, historyScrollPane, selectedDiffSplitPane) mainSplitPane.setBackground(DARK_BACKGROUND) mainSplitPane.setDividerSize(2) mainSplitPane.setResizeWeight(0.5) @@ -297,8 +284,10 @@ def compare(self, s1, s2): self.log_message("Extension initialized and running.") callbacks.printOutput(SQUID_ASCII + "\n\n") - callbacks.printOutput("Yer configs be stowed and ready from " + CONFIG_DIR) - callbacks.printOutput("\nNow ye be speakin' the pirate's tongue, savvy?") + callbacks.printOutput( + "Yer configs be stowed and ready from " + CONFIG_DIR) + callbacks.printOutput( + "\nNow ye be speakin' the pirate's tongue, savvy?") self.promptForConfiguration() self.applyDarkTheme(self.tabbedPane) @@ -312,8 +301,10 @@ def getUiComponent(self): def stopExtension(self, event): if self.is_running: self.is_running = False - self._callbacks.removeHttpListener(self) - self.log_message("Extension stopped. No further traffic will be processed.") + self._callbacks.removeHttpListener( + self) + self.log_message( + "Extension stopped. No further traffic will be processed.") self.updateUIState() def updateUIState(self): @@ -324,8 +315,7 @@ def updateUIState(self): self.stopButton.addActionListener(self.stopExtension) else: self.stopButton.setText( - "Extension Stopped - Unload and reload if required, doing so will remove displayed logs and state" - ) + "Extension Stopped - Unload and reload if required, doing so will remove displayed logs and state") for listener in self.stopButton.getActionListeners(): self.stopButton.removeActionListener(listener) @@ -333,18 +323,17 @@ def loadConfigFiles(self): if not os.path.exists(CONFIG_DIR): self.log_message("Config directory not found: {CONFIG_DIR}") return [] - return [f for f in os.listdir(CONFIG_DIR) if f.endswith(".json")] + return [f for f in os.listdir(CONFIG_DIR) if f.endswith('.json')] def loadConfiguration(self, event): selected_config = self.configSelector.getSelectedItem() config_path = os.path.join(CONFIG_DIR, selected_config) if os.path.exists(config_path): try: - with open(config_path, "r") as config_file: + with open(config_path, 'r') as config_file: self.config = json.load(config_file) - self.log_message( - "Loaded configuration: %s" % json.dumps(self.config, indent=2) - ) + self.log_message("Loaded configuration: %s" % + json.dumps(self.config, indent=2)) try: self.api_adapter = get_api_adapter(self.config) self.log_message("API adapter initialized successfully") @@ -353,21 +342,21 @@ def loadConfiguration(self, event): self.api_adapter = None except Exception as e: self.log_message( - "Unexpected error initializing API adapter: %s" % str(e) - ) + "Unexpected error initializing API adapter: %s" % str(e)) self.api_adapter = None except ValueError as e: self.log_message( - "Error parsing JSON in configuration file: %s" % str(e) - ) + "Error parsing JSON in configuration file: %s" % str(e)) self.config = None self.api_adapter = None except Exception as e: - self.log_message("Unexpected error loading configuration: %s" % str(e)) + self.log_message( + "Unexpected error loading configuration: %s" % str(e)) self.config = None self.api_adapter = None else: - self.log_message("Configuration file %s not found." % selected_config) + self.log_message( + "Configuration file %s not found." % selected_config) self.config = None self.api_adapter = None @@ -383,19 +372,18 @@ def create_inference_logger_tab(self): self.inferenceLogTable.setGridColor(DREADNODE_GREY) self.inferenceLogTableModel = table.DefaultTableModel( - ["Timestamp", "API Endpoint", "Proxy Request", "Model Response", "Status"], - 0, - ) + ["Timestamp", "API Endpoint", "Proxy Request", "Model Response", "Status"], 0) self.inferenceLogTable.setModel(self.inferenceLogTableModel) - self.inferenceLogTable.setSelectionMode(ListSelectionModel.SINGLE_SELECTION) + self.inferenceLogTable.setSelectionMode( + ListSelectionModel.SINGLE_SELECTION) self.inferenceLogTable.getSelectionModel().addListSelectionListener( - self.inferenceLogSelectionChanged - ) + self.inferenceLogSelectionChanged) inferenceLogScrollPane = JScrollPane(self.inferenceLogTable) border = BorderFactory.createTitledBorder( - BorderFactory.createLineBorder(DREADNODE_ORANGE), "API Requests Log" + BorderFactory.createLineBorder(DREADNODE_ORANGE), + "API Requests Log" ) border.setTitleColor(DREADNODE_ORANGE) border.setTitleFont(border.getTitleFont().deriveFont(Font.BOLD, 14)) @@ -408,12 +396,10 @@ def create_inference_logger_tab(self): self.inferenceRequestDetail.setBackground(LIGHTER_BACKGROUND) self.inferenceRequestDetail.setForeground(DREADNODE_ORANGE) requestDetailPane = JScrollPane(self.inferenceRequestDetail) - requestDetailPane.setBorder( - BorderFactory.createTitledBorder( - BorderFactory.createLineBorder(DREADNODE_ORANGE), - "Inference Request Detail", - ) - ) + requestDetailPane.setBorder(BorderFactory.createTitledBorder( + BorderFactory.createLineBorder(DREADNODE_ORANGE), + "Inference Request Detail" + )) self.inferenceResponseDetail = JTextArea(10, 30) self.inferenceResponseDetail.setEditable(False) @@ -422,22 +408,24 @@ def create_inference_logger_tab(self): self.inferenceResponseDetail.setBackground(LIGHTER_BACKGROUND) self.inferenceResponseDetail.setForeground(DREADNODE_ORANGE) responseDetailPane = JScrollPane(self.inferenceResponseDetail) - responseDetailPane.setBorder( - BorderFactory.createTitledBorder( - BorderFactory.createLineBorder(DREADNODE_ORANGE), - "Inference Response Detail", - ) - ) + responseDetailPane.setBorder(BorderFactory.createTitledBorder( + BorderFactory.createLineBorder(DREADNODE_ORANGE), + "Inference Response Detail" + )) # Create split pane for details detailsSplitPane = JSplitPane( - JSplitPane.HORIZONTAL_SPLIT, requestDetailPane, responseDetailPane + JSplitPane.HORIZONTAL_SPLIT, + requestDetailPane, + responseDetailPane ) detailsSplitPane.setResizeWeight(0.5) # Create main split pane mainSplitPane = JSplitPane( - JSplitPane.VERTICAL_SPLIT, inferenceLogScrollPane, detailsSplitPane + JSplitPane.VERTICAL_SPLIT, + inferenceLogScrollPane, + detailsSplitPane ) mainSplitPane.setResizeWeight(0.5) @@ -450,12 +438,8 @@ def inferenceLogSelectionChanged(self, event): if selectedRow != -1: try: # Get the request and response data from columns - request = self.inferenceLogTableModel.getValueAt( - selectedRow, 2 - ) # Proxy Request column - response = self.inferenceLogTableModel.getValueAt( - selectedRow, 3 - ) # Model Response column + request = self.inferenceLogTableModel.getValueAt(selectedRow, 2) # Proxy Request column + response = self.inferenceLogTableModel.getValueAt(selectedRow, 3) # Model Response column # Format the request JSON try: @@ -474,15 +458,11 @@ def inferenceLogSelectionChanged(self, event): else: response_obj = json.loads(response) - if ( - isinstance(response_obj, dict) - and "message" in response_obj - and "content" in response_obj["message"] - ): + if isinstance(response_obj, dict) and 'message' in response_obj and 'content' in response_obj['message']: # Get just the content string and handle unicode - content = response_obj["message"]["content"] + content = response_obj['message']['content'] # Handle unicode strings and newlines - formatted_response = content.replace("\\n", "\n") + formatted_response = content.replace('\\n', '\n') else: formatted_response = json.dumps(response_obj, indent=2) except (ValueError, AttributeError, TypeError): @@ -510,8 +490,8 @@ def historyTableSelectionChanged(self, event): http_pair = json.loads(http_pair_json) self.selectedRequestArea.setText(json.dumps(http_pair, indent=2)) - model_text = model_analysis.strip('"').decode("unicode_escape") - model_text = model_text.replace("\\n", "\n") + model_text = model_analysis.strip('"').decode('unicode_escape') + model_text = model_text.replace('\\n', '\n') self.selectedResponseArea.setText(model_text) self.selectedRequestArea.setCaretPosition(0) @@ -527,19 +507,12 @@ def historyTableSelectionChanged(self, event): def colorizeHistoryTable(self): renderer = self.SeverityCellRenderer() for column in range(self.historyTable.getColumnCount()): - self.historyTable.getColumnModel().getColumn(column).setCellRenderer( - renderer - ) + self.historyTable.getColumnModel().getColumn(column).setCellRenderer(renderer) class SeverityCellRenderer(DefaultTableCellRenderer): - def getTableCellRendererComponent( - self, table, value, isSelected, hasFocus, row, column - ): - component = super( - BurpExtender.SeverityCellRenderer, self - ).getTableCellRendererComponent( - table, value, isSelected, hasFocus, row, column - ) + def getTableCellRendererComponent(self, table, value, isSelected, hasFocus, row, column): + component = super(BurpExtender.SeverityCellRenderer, self).getTableCellRendererComponent( + table, value, isSelected, hasFocus, row, column) component.setForeground(DREADNODE_GREY) @@ -584,16 +557,17 @@ def log_message(self, message): self.temp_log_messages.append(log_entry) else: self.logArea.append(log_entry) - self.logArea.setCaretPosition(self.logArea.getDocument().getLength()) + self.logArea.setCaretPosition( + self.logArea.getDocument().getLength()) try: # Try to create/write to log file with explicit permissions log_dir = os.path.dirname(self.log_file_path) if not os.path.exists(log_dir): - os.makedirs(log_dir, 0o755) + os.makedirs(log_dir, 0755) # Python2 octal notation # Open with explicit write permissions - with open(self.log_file_path, "a+") as log_file: + with open(self.log_file_path, 'a+') as log_file: log_file.write(log_entry) except (IOError, OSError) as e: print("Warning: Could not write to log file: %s" % str(e)) @@ -620,7 +594,7 @@ def applyDarkTheme(self, component): component.setDividerSize(2) component.setDividerLocation(0.5) - if hasattr(component, "getComponents"): + if hasattr(component, 'getComponents'): for child in component.getComponents(): self.applyDarkTheme(child) @@ -628,10 +602,8 @@ def processHttpMessage(self, toolFlag, messageIsRequest, messageInfo): if not self.is_running: return if not self.api_adapter: - if not hasattr(self, "_no_adapter_logged"): - self.log_message( - "No API adapter configured. Please select a configuration file." - ) + if not hasattr(self, '_no_adapter_logged'): + self.log_message("No API adapter configured. Please select a configuration file.") self._no_adapter_logged = True return if messageIsRequest: @@ -642,9 +614,10 @@ def processHttpMessage(self, toolFlag, messageIsRequest, messageInfo): response = messageInfo # Filter MIME content types to reduce noise and exceeding tokens - responseInfo = self._helpers.analyzeResponse(response.getResponse()) + responseInfo = self._helpers.analyzeResponse( + response.getResponse()) contentType = responseInfo.getStatedMimeType().lower() - if contentType in ["css", "image", "script", "video", "audio", "font"]: + if contentType in ['css', 'image', 'script', 'video', 'audio', 'font']: return # Filter request size @@ -652,24 +625,20 @@ def processHttpMessage(self, toolFlag, messageIsRequest, messageInfo): return try: analyzed_request = self._helpers.analyzeRequest(request) - analyzed_response = self._helpers.analyzeResponse( - response.getResponse() - ) + analyzed_response = self._helpers.analyzeResponse(response.getResponse()) # Get the body bytes - request_body = request.getRequest()[analyzed_request.getBodyOffset() :] - response_body = response.getResponse()[ - analyzed_response.getBodyOffset() : - ] + request_body = request.getRequest()[analyzed_request.getBodyOffset():] + response_body = response.getResponse()[analyzed_response.getBodyOffset():] # Try to safely decode bodies, fallback to hex for binary data try: - request_body_str = request_body.tostring().decode("utf-8") + request_body_str = request_body.tostring().decode('utf-8') except UnicodeDecodeError: request_body_str = "" % len(request_body) try: - response_body_str = response_body.tostring().decode("utf-8") + response_body_str = response_body.tostring().decode('utf-8') except UnicodeDecodeError: response_body_str = "" % len(response_body) @@ -686,41 +655,33 @@ def processHttpMessage(self, toolFlag, messageIsRequest, messageInfo): request_data = { "method": analyzed_request.getMethod(), "url": str(request.getUrl()), - "headers": dict( - header.split(": ", 1) - for header in analyzed_request.getHeaders()[1:] - if ": " in header - ), - "body": request_body_str, + "headers": dict(header.split(': ', 1) for header in analyzed_request.getHeaders()[1:] if ': ' in header), + "body": request_body_str } # Package the original HTTP response response_data = { "status_code": analyzed_response.getStatusCode(), - "headers": dict( - header.split(": ", 1) - for header in analyzed_response.getHeaders()[1:] - if ": " in header - ), - "body": response_body_str, + "headers": dict(header.split(': ', 1) for header in analyzed_response.getHeaders()[1:] if ': ' in header), + "body": response_body_str } # Create the HTTP pair to send to the model - http_pair = {"request": request_data, "response": response_data} + http_pair = { + "request": request_data, + "response": response_data + } # Load prompt template for system role if os.path.exists(PROXY_PROMPT): - with open(PROXY_PROMPT, "r") as prompt_file: + with open(PROXY_PROMPT, 'r') as prompt_file: system_content = prompt_file.read().strip() # Only log if there's an issue or if it's different from last time - if ( - not hasattr(self, "_last_system_content") - or system_content != self._last_system_content - ): + if not hasattr(self, '_last_system_content') or system_content != self._last_system_content: self.log_message("Custom prompt loaded from " + PROXY_PROMPT) self._last_system_content = system_content else: - if not hasattr(self, "_prompt_missing_logged"): + if not hasattr(self, '_prompt_missing_logged'): self.log_message("No prompt file found. Using default prompt.") self._prompt_missing_logged = True system_content = "Examine this request and response pair for any security issues:" @@ -728,7 +689,7 @@ def processHttpMessage(self, toolFlag, messageIsRequest, messageInfo): # Prepare the request using the adapter remote_request = self.api_adapter.prepare_request( user_content=json.dumps(http_pair, indent=2), - system_content=system_content, + system_content=system_content ) # Send request to model @@ -739,22 +700,18 @@ def processHttpMessage(self, toolFlag, messageIsRequest, messageInfo): try: response = urllib2.urlopen(req, json.dumps(remote_request)) response_data = response.read() - processed_response = self.api_adapter.process_response( - response_data - ) + processed_response = self.api_adapter.process_response(response_data) status = "Success" except urllib2.HTTPError as e: - error_message = e.read().decode("utf-8") + error_message = e.read().decode('utf-8') error_details = { "status_code": e.code, "headers": dict(e.headers), - "body": error_message, + "body": error_message } processed_response = json.dumps(error_details, indent=2) - self.log_message( - "API Error Response:\nStatus Code: %d\nHeaders: %s\nBody: %s" - % (e.code, dict(e.headers), error_message) - ) + self.log_message("API Error Response:\nStatus Code: %d\nHeaders: %s\nBody: %s" % + (e.code, dict(e.headers), error_message)) status = "Failed (%d)" % e.code except Exception as e: processed_response = "Error: %s" % str(e) @@ -764,75 +721,51 @@ def processHttpMessage(self, toolFlag, messageIsRequest, messageInfo): # Always log to inference logger if self.is_running: timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S") - self.inferenceLogTableModel.addRow( - [ - timestamp, - self.config.get("host", ""), - json.dumps(remote_request), - processed_response, - status, - ] - ) + self.inferenceLogTableModel.addRow([ + timestamp, + self.config.get("host", ""), + json.dumps(remote_request), + processed_response, + status + ]) # Only update the main UI if we got a successful response if status == "Success" and self.is_running: # Add to history table with metadata - self.historyTableModel.addRow( - [ - table_metadata.get("id", ""), - timestamp, - table_metadata.get("host", ""), - table_metadata.get("url", ""), - json.dumps( - http_pair, indent=2 - ), # Store the HTTP request/response pair - json.dumps( - processed_response, indent=2 - ), # Store the model's analysis - ] - ) + self.historyTableModel.addRow([ + table_metadata.get("id", ""), + timestamp, + table_metadata.get("host", ""), + table_metadata.get("url", ""), + json.dumps(http_pair, indent=2), # Store the HTTP request/response pair + json.dumps(processed_response, indent=2) # Store the model's analysis + ]) self.colorizeHistoryTable() - self.requestArea.append( - "\n\n=== Request #" + str(self.request_counter) + " ===\n" - ) + self.requestArea.append("\n\n=== Request #" + str(self.request_counter) + " ===\n") try: # Format the request nicely formatted_request = json.dumps(http_pair, indent=2) - formatted_request = formatted_request.replace("\\n", "\n") + formatted_request = formatted_request.replace('\\n', '\n') formatted_request = formatted_request.replace('\\"', '"') self.requestArea.append(formatted_request) except Exception as e: self.requestArea.append(str(http_pair)) - self.requestArea.setCaretPosition( - self.requestArea.getDocument().getLength() - ) + self.requestArea.setCaretPosition(self.requestArea.getDocument().getLength()) - self.responseArea.append( - "\n\n=== Response #" + str(self.request_counter) + " ===\n" - ) + self.responseArea.append("\n\n=== Response #" + str(self.request_counter) + " ===\n") try: # Format the response nicely - if ( - isinstance(processed_response, dict) - and "message" in processed_response - and "content" in processed_response["message"] - ): - formatted_response = processed_response["message"][ - "content" - ] + if isinstance(processed_response, dict) and 'message' in processed_response and 'content' in processed_response['message']: + formatted_response = processed_response['message']['content'] else: - formatted_response = json.dumps( - processed_response, indent=2 - ) - formatted_response = formatted_response.replace("\\n", "\n") + formatted_response = json.dumps(processed_response, indent=2) + formatted_response = formatted_response.replace('\\n', '\n') formatted_response = formatted_response.replace('\\"', '"') self.responseArea.append(formatted_response) except Exception as e: self.responseArea.append(str(processed_response)) - self.responseArea.setCaretPosition( - self.responseArea.getDocument().getLength() - ) + self.responseArea.setCaretPosition(self.responseArea.getDocument().getLength()) except Exception as e: self.log_message("Error processing request: %s" % str(e)) @@ -843,5 +776,4 @@ def promptForConfiguration(self): "Select a configuration file to load in the burpference extension" " tab and go brrr", "burpference Configuration Required", - JOptionPane.INFORMATION_MESSAGE, - ) + JOptionPane.INFORMATION_MESSAGE) \ No newline at end of file From 187d9bc90b13e119df232d8633748761a96b2559 Mon Sep 17 00:00:00 2001 From: Ads Dawson <104169244+GangGreenTemperTatum@users.noreply.github.com> Date: Mon, 30 Dec 2024 14:17:34 -0500 Subject: [PATCH 5/6] chore: temp rm ruff for cleanup --- .pre-commit-config.yaml | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 18c7b6f..54313f5 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -18,16 +18,6 @@ repos: - id: actionlint name: Check Github Actions - # Python linting - - repo: https://github.com/astral-sh/ruff-pre-commit - # Ruff version. - rev: 8b76f04e7e5a9cd259e9d1db7799599355f97cdf # v0.8.2 - hooks: - # Run the linter. - - id: ruff - # Run the formatter. - - id: ruff-format - # Python code security - repo: https://github.com/PyCQA/bandit rev: 8fd258abbac759d62863779f946d6a88e8eabb0f #1.8.0 From 0600499e9dc0779cc28a79aee9fc6a868a45df6d Mon Sep 17 00:00:00 2001 From: Ads Dawson <104169244+GangGreenTemperTatum@users.noreply.github.com> Date: Mon, 30 Dec 2024 14:24:49 -0500 Subject: [PATCH 6/6] fix: whitelist fp shellcheck --- .github/workflows/rigging_pr_description.yml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.github/workflows/rigging_pr_description.yml b/.github/workflows/rigging_pr_description.yml index b2d36ce..5bb8b8f 100644 --- a/.github/workflows/rigging_pr_description.yml +++ b/.github/workflows/rigging_pr_description.yml @@ -19,10 +19,11 @@ jobs: # Get the diff first - name: Get Diff id: diff + # shellcheck disable=SC2102 run: | git fetch origin "${{ github.base_ref }}" MERGE_BASE=$(git merge-base HEAD "origin/${{ github.base_ref }}") - # Encode the diff as base64 with explicit newline handling + # Use separate diff arguments instead of range notation DIFF=$(git diff "$MERGE_BASE" HEAD | base64 --wrap=0) echo "diff=${DIFF}" >> "$GITHUB_OUTPUT" - uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b #v5.0.3