ㅎㅎㅎ

# -*- coding: utf-8 -*-
from burp import IBurpExtender, IHttpListener, ITab
from javax.swing import JPanel, JButton, JTextField, JLabel, JComboBox, JTable, JScrollPane
from javax.swing.table import DefaultTableModel
from javax.swing import DefaultCellEditor
from java.awt import BorderLayout, FlowLayout

class BurpExtender(IBurpExtender, IHttpListener, ITab):

    def registerExtenderCallbacks(self, callbacks):
        self.callbacks = callbacks
        self.helpers = callbacks.getHelpers()
        callbacks.setExtensionName("Keyword Highlighter Advanced")
        callbacks.registerHttpListener(self)

        self.panel = JPanel(BorderLayout())
        self.top_panel = JPanel(FlowLayout())

        self.keyword_field = JTextField(15)
        self.color_selector = JComboBox(["red", "orange", "yellow", "green", "cyan", "blue", "pink", "magenta", "gray"])
        self.scope_selector = JComboBox(["Request", "Response", "Both"])

        self.add_button = JButton("Add", actionPerformed=self.add_keyword)
        self.delete_button = JButton("Delete Selected", actionPerformed=self.delete_keyword)

        self.top_panel.add(JLabel("Keyword:"))
        self.top_panel.add(self.keyword_field)
        self.top_panel.add(JLabel("Color:"))
        self.top_panel.add(self.color_selector)
        self.top_panel.add(JLabel("Scope:"))
        self.top_panel.add(self.scope_selector)
        self.top_panel.add(self.add_button)
        self.top_panel.add(self.delete_button)

        self.table_model = DefaultTableModel(["Keyword", "Color", "Scope"], 0)
        self.table = JTable(self.table_model)

        color_editor = DefaultCellEditor(JComboBox(["red", "orange", "yellow", "green", "cyan", "blue", "pink", "magenta", "gray"]))
        scope_editor = DefaultCellEditor(JComboBox(["Request", "Response", "Both"]))
        self.table.getColumnModel().getColumn(1).setCellEditor(color_editor)
        self.table.getColumnModel().getColumn(2).setCellEditor(scope_editor)

        scroll_pane = JScrollPane(self.table)

        self.panel.add(self.top_panel, BorderLayout.NORTH)
        self.panel.add(scroll_pane, BorderLayout.CENTER)

        callbacks.addSuiteTab(self)
        print("[+] Keyword Highlighter Advanced loaded.")

    def getTabCaption(self):
        return "Keyword Highlighter"

    def getUiComponent(self):
        return self.panel

    def add_keyword(self, event):
        keyword = self.keyword_field.getText().strip()
        color = self.color_selector.getSelectedItem()
        scope = self.scope_selector.getSelectedItem()
        if keyword:
            self.table_model.addRow([keyword, color, scope])

    def delete_keyword(self, event):
        selected_row = self.table.getSelectedRow()
        if selected_row != -1:
            self.table_model.removeRow(selected_row)

    def processHttpMessage(self, toolFlag, messageIsRequest, messageInfo):
        if toolFlag != self.callbacks.TOOL_PROXY:
            return

        try:
            request_bytes = messageInfo.getRequest()
            request_info = self.helpers.analyzeRequest(messageInfo)
            request_url = str(request_info.getUrl())
            request_body = request_bytes[request_info.getBodyOffset():].tostring()
            request_str = request_url + " " + request_body.decode("utf-8", "ignore")

            response_str = ""
            if messageInfo.getResponse() is not None:
                response_bytes = messageInfo.getResponse()
                response_info = self.helpers.analyzeResponse(response_bytes)
                response_body = response_bytes[response_info.getBodyOffset():].tostring()
                response_str = response_body.decode("utf-8", "ignore")

            # 전부 소문자로 변환
            request_str = request_str.lower()
            response_str = response_str.lower()

            for row in range(self.table_model.getRowCount()):
                keyword = self.table_model.getValueAt(row, 0).lower()
                color = self.table_model.getValueAt(row, 1).lower()
                scope = self.table_model.getValueAt(row, 2)

                matched = False

                if scope == "Request" and keyword in request_str:
                    matched = True
                elif scope == "Response" and keyword in response_str:
                    matched = True
                elif scope == "Both" and (keyword in request_str or keyword in response_str):
                    matched = True

                if matched:
                    messageInfo.setHighlight(color)
                    print("[+] Highlight applied: '{}' matched in {} - Color: {}".format(keyword, scope, color))
                    break

        except Exception as e:
            print("[!] Exception in processHttpMessage: {}".format(str(e)))