ㄱㄷㄷ

from burp import IBurpExtender, IHttpListener, ITab
from javax.swing import JPanel, JButton, JTextField, JLabel, JComboBox, JTable, JScrollPane
from javax.swing.table import DefaultTableModel
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"])  # 추가된 scope 선택

        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)
        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:
            request = messageInfo.getRequest()
            response = messageInfo.getResponse()

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

            if response:
                analyzedResponse = self.helpers.analyzeResponse(response)
                response_body = response[analyzedResponse.getBodyOffset():].tostring()
                response_content = response_body.lower()
            else:
                response_content = ""

            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)

                if scope == "Request" and keyword in request_content:
                    messageInfo.setHighlight(color)
                    return
                elif scope == "Response" and keyword in response_content:
                    messageInfo.setHighlight(color)
                    return
                elif scope == "Both" and (keyword in request_content or keyword in response_content):
                    messageInfo.setHighlight(color)
                    return