ㅎㅎ

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)

        # GUI 구성
        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

        request = messageInfo.getRequest()
        request_info = self.helpers.analyzeRequest(request)
        request_url = request_info.getUrl().toString()
        request_body = request[request_info.getBodyOffset():].tostring()
        request_content = (request_url + request_body).lower()

        response_content = ""
        if messageInfo.getResponse():
            response = messageInfo.getResponse()
            response_info = self.helpers.analyzeResponse(response)
            response_body = response[response_info.getBodyOffset():].tostring()
            response_content = response_body.lower()

        highlight_color = None

        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_content:
                matched = True
            elif scope == "Response" and keyword in response_content:
                matched = True
            elif scope == "Both" and (keyword in request_content or keyword in response_content):
                matched = True

            if matched:
                print("[+] Keyword matched: '{}' in {}".format(keyword, scope))
                print("    -> Applying highlight color: {}".format(color))
                highlight_color = color
                break

        if highlight_color:
            messageInfo.setHighlight(highlight_color)