diff --git a/src/main/java/hae/Config.java b/src/main/java/hae/Config.java index 67e1059..1b8a9a8 100644 --- a/src/main/java/hae/Config.java +++ b/src/main/java/hae/Config.java @@ -10,6 +10,8 @@ public class Config { public static String host = "gh0st.cn"; + public static String status = "404"; + public static String[] scope = new String[]{ "any", "any header", diff --git a/src/main/java/hae/HaE.java b/src/main/java/hae/HaE.java index 2ab3ff7..0432ace 100644 --- a/src/main/java/hae/HaE.java +++ b/src/main/java/hae/HaE.java @@ -18,7 +18,7 @@ public class HaE implements BurpExtension { @Override public void initialize(MontoyaApi api) { // 设置扩展名称 - String version = "3.3"; + String version = "3.3.1"; api.extension().setName(String.format("HaE (%s) - Highlighter and Extractor", version)); // 加载扩展后输出的项目信息 diff --git a/src/main/java/hae/component/board/Databoard.java b/src/main/java/hae/component/board/Databoard.java index 34e8d59..2270474 100644 --- a/src/main/java/hae/component/board/Databoard.java +++ b/src/main/java/hae/component/board/Databoard.java @@ -13,9 +13,11 @@ import hae.utils.string.StringProcessor; import javax.swing.*; +import javax.swing.border.TitledBorder; import javax.swing.event.DocumentEvent; import javax.swing.event.DocumentListener; import javax.swing.filechooser.FileNameExtensionFilter; +import javax.swing.table.DefaultTableModel; import javax.swing.table.TableColumnModel; import javax.swing.table.TableModel; import javax.swing.table.TableRowSorter; @@ -49,8 +51,8 @@ public class Databoard extends JPanel { private SwingWorker>, Void> handleComboBoxWorker; private SwingWorker applyHostFilterWorker; - private SwingWorker, Void> exportActionWorker; - private SwingWorker, Void> importActionWorker; + private SwingWorker, Void> exportActionWorker; + private SwingWorker, Void> importActionWorker; public Databoard(MontoyaApi api, ConfigLoader configLoader, MessageTableModel messageTableModel) { this.api = api; @@ -193,11 +195,11 @@ public void changedUpdate(DocumentEvent e) { private void handleComboBoxAction(ActionEvent e) { if (!isMatchHost && hostComboBox.getSelectedItem() != null) { - progressBar.setVisible(true); - setProgressBar(true); String selectedHost = hostComboBox.getSelectedItem().toString(); if (getHostByList().contains(selectedHost)) { + progressBar.setVisible(true); + setProgressBar(true); hostTextField.setText(selectedHost); if (handleComboBoxWorker != null && !handleComboBoxWorker.isDone()) { @@ -385,9 +387,9 @@ private void exportActionPerformed(ActionEvent e) { exportActionWorker.cancel(true); } - exportActionWorker = new SwingWorker, Void>() { + exportActionWorker = new SwingWorker, Void>() { @Override - protected List doInBackground() { + protected List doInBackground() { ConcurrentHashMap>> dataMap = Config.globalDataMap; return exportData(selectedHost, exportDir, dataMap); } @@ -395,11 +397,9 @@ protected List doInBackground() { @Override protected void done() { try { - List taskStatusList = get(); + List taskStatusList = get(); if (!taskStatusList.isEmpty()) { - String exportStatusMessage = String.format("Exported File List Status:\n%s", String.join("\n", taskStatusList)); - - JOptionPane.showMessageDialog(Databoard.this, generateTaskStatusPane(exportStatusMessage), "Info", JOptionPane.INFORMATION_MESSAGE); + JOptionPane.showMessageDialog(Databoard.this, generateTaskStatusPane(taskStatusList), "Info", JOptionPane.INFORMATION_MESSAGE); } } catch (Exception ignored) { } @@ -409,17 +409,36 @@ protected void done() { exportActionWorker.execute(); } - private JScrollPane generateTaskStatusPane(String message) { - JTextArea textArea = new JTextArea(message); - textArea.setEditable(false); - textArea.setLineWrap(true); - JScrollPane scrollPane = new JScrollPane(textArea); - scrollPane.setPreferredSize(new Dimension(400, 200)); + private JScrollPane generateTaskStatusPane(List dataList) { + String[] columnNames = {"#", "Filename", "Status"}; + DefaultTableModel taskStatusTableModel = new DefaultTableModel(columnNames, 0); + JTable taskStatusTable = new JTable(taskStatusTableModel); + + for (Object[] data : dataList) { + int rowCount = taskStatusTable.getRowCount(); + int id = rowCount > 0 ? (Integer) taskStatusTable.getValueAt(rowCount - 1, 0) + 1 : 1; + Object[] rowData = new Object[data.length + 1]; + rowData[0] = id; + System.arraycopy(data, 0, rowData, 1, data.length); + taskStatusTableModel.addRow(rowData); + } + + TableRowSorter sorter = new TableRowSorter<>(taskStatusTableModel); + taskStatusTable.setRowSorter(sorter); + + JScrollPane scrollPane = new JScrollPane(taskStatusTable); + scrollPane.setBorder(new TitledBorder("Task status")); + scrollPane.setPreferredSize(new Dimension(500, 300)); + + int paneWidth = scrollPane.getPreferredSize().width; + taskStatusTable.getColumnModel().getColumn(0).setPreferredWidth((int) (paneWidth * 0.1)); + taskStatusTable.getColumnModel().getColumn(1).setPreferredWidth((int) (paneWidth * 0.7)); + taskStatusTable.getColumnModel().getColumn(2).setPreferredWidth((int) (paneWidth * 0.2)); return scrollPane; } - private List exportData(String selectedHost, String exportDir, Map>> dataMap) { + private List exportData(String selectedHost, String exportDir, Map>> dataMap) { return dataMap.entrySet().stream() .filter(entry -> selectedHost.equals("*") || StringProcessor.matchesHostPattern(entry.getKey(), selectedHost)) .filter(entry -> !entry.getKey().contains("*")) @@ -428,7 +447,7 @@ private List exportData(String selectedHost, String exportDir, Map>> entry, String exportDir) { + private Object[] exportEntry(Map.Entry>> entry, String exportDir) { String key = entry.getKey(); Map> ruleMap = entry.getValue(); @@ -442,7 +461,7 @@ private String exportEntry(Map.Entry>> entry, S .collect(Collectors.toMap( messageEntry -> messageEntry, messageEntry -> StringProcessor.getRandomUUID(), - (existing, replacement) -> existing // 在冲突时保留现有的映射 + (existing, replacement) -> existing )); Map> httpMap = processEntries( @@ -463,7 +482,7 @@ private String exportEntry(Map.Entry>> entry, S String filename = String.format("%s/%s-%s.hae", exportDir, StringProcessor.getCurrentTime(), hostName); boolean createdStatus = projectProcessor.createHaeFile(filename, key, ruleMap, urlMap, httpMap); - return String.format("Filename: %s, Status: %s", filename, createdStatus); + return new Object[]{filename, createdStatus}; } @@ -507,9 +526,9 @@ private void importActionPerformed(ActionEvent e) { importActionWorker.cancel(true); } - importActionWorker = new SwingWorker, Void>() { + importActionWorker = new SwingWorker, Void>() { @Override - protected List doInBackground() { + protected List doInBackground() { List filesWithExtension = findFilesWithExtension(new File(exportDir), ".hae"); return filesWithExtension.stream() .map(Databoard.this::importData) @@ -519,10 +538,9 @@ protected List doInBackground() { @Override protected void done() { try { - List taskStatusList = get(); + List taskStatusList = get(); if (!taskStatusList.isEmpty()) { - String importStatusMessage = "Imported File List Status:\n" + String.join("\n", taskStatusList); - JOptionPane.showMessageDialog(Databoard.this, generateTaskStatusPane(importStatusMessage), "Info", JOptionPane.INFORMATION_MESSAGE); + JOptionPane.showMessageDialog(Databoard.this, generateTaskStatusPane(taskStatusList), "Info", JOptionPane.INFORMATION_MESSAGE); } } catch (Exception ignored) { } @@ -532,7 +550,7 @@ protected void done() { importActionWorker.execute(); } - private String importData(String filename) { + private Object[] importData(String filename) { ExecutorService executor = Executors.newFixedThreadPool(Runtime.getRuntime().availableProcessors() * 2); HaeFileContent haeFileContent = projectProcessor.readHaeFile(filename); @@ -568,7 +586,7 @@ private String importData(String filename) { } } - return String.format("Filename: %s, Status: %s", filename, readStatus); + return new Object[]{filename, readStatus}; } private List findFilesWithExtension(File directory, String extension) { @@ -648,6 +666,8 @@ private void clearActionPerformed(ActionEvent e) { } messageTableModel.deleteByHost(host); + + hostTextField.setText(""); } } } diff --git a/src/main/java/hae/component/board/message/MessageTableModel.java b/src/main/java/hae/component/board/message/MessageTableModel.java index e0b9bb2..3943a8f 100644 --- a/src/main/java/hae/component/board/message/MessageTableModel.java +++ b/src/main/java/hae/component/board/message/MessageTableModel.java @@ -435,7 +435,7 @@ public String getColumnName(int columnIndex) { public class MessageTable extends JTable { private MessageEntry messageEntry; - private SwingWorker currentWorker; + private SwingWorker currentWorker; private int lastSelectedIndex = -1; private final HttpRequestEditor requestEditor; private final HttpResponseEditor responseEditor; @@ -450,16 +450,13 @@ public MessageTable(TableModel messageTableModel, HttpRequestEditor requestEdito public void changeSelection(int row, int col, boolean toggle, boolean extend) { super.changeSelection(row, col, toggle, extend); - requestEditor.setRequest(HttpRequest.httpRequest("Loading...")); - responseEditor.setResponse(HttpResponse.httpResponse("Loading...")); - if (currentWorker != null && !currentWorker.isDone()) { currentWorker.cancel(true); } currentWorker = new SwingWorker<>() { @Override - protected Void doInBackground() { + protected ByteArray[] doInBackground() { int selectedIndex = convertRowIndexToModel(row); if (lastSelectedIndex != selectedIndex) { lastSelectedIndex = selectedIndex; @@ -470,14 +467,28 @@ protected Void doInBackground() { ByteArray requestByte = httpRequestResponse.request().toByteArray(); ByteArray responseByte = httpRequestResponse.response().toByteArray(); - requestEditor.setRequest(HttpRequest.httpRequest(messageEntry.getRequestResponse().httpService(), requestByte)); - responseEditor.setResponse(HttpResponse.httpResponse(responseByte)); - + ByteArray[] httpByteArray = new ByteArray[2]; + httpByteArray[0] = requestByte; + httpByteArray[1] = responseByte; + return httpByteArray; } return null; } + + @Override + protected void done() { + try { + ByteArray[] retByteArray = get(); + if (retByteArray != null) { + requestEditor.setRequest(HttpRequest.httpRequest(messageEntry.getRequestResponse().httpService(), retByteArray[0])); + responseEditor.setResponse(HttpResponse.httpResponse(retByteArray[1])); + } + } catch (Exception ignored) { + } + } }; + currentWorker.execute(); } } diff --git a/src/main/java/hae/component/board/table/AIPower.java b/src/main/java/hae/component/board/table/AIPower.java index c05969a..24af1b9 100644 --- a/src/main/java/hae/component/board/table/AIPower.java +++ b/src/main/java/hae/component/board/table/AIPower.java @@ -34,7 +34,7 @@ public class AIPower { public AIPower(MontoyaApi api, ConfigLoader configLoader, String aiModel, String aiBaseUrl, String[] apiKey) { this.api = api; this.configLoader = configLoader; - this.httpUtils = new HttpUtils(api); + this.httpUtils = new HttpUtils(api, configLoader); this.aiModel = aiModel; this.aiBaseUrl = aiBaseUrl; diff --git a/src/main/java/hae/component/config/Config.java b/src/main/java/hae/component/config/Config.java index 2b3d61c..b4dc82c 100644 --- a/src/main/java/hae/component/config/Config.java +++ b/src/main/java/hae/component/config/Config.java @@ -27,24 +27,23 @@ public class Config extends JPanel { private final MontoyaApi api; private final ConfigLoader configLoader; private final Rules rules; - private JTextField addTextField; private final String defaultText = "Enter a new item"; - private final GridBagConstraints constraints = new GridBagConstraints(); public Config(MontoyaApi api, ConfigLoader configLoader, Rules rules) { this.api = api; this.configLoader = configLoader; this.rules = rules; - constraints.weightx = 1.0; - constraints.fill = GridBagConstraints.HORIZONTAL; - initComponents(); } private void initComponents() { setLayout(new BorderLayout()); + GridBagConstraints constraints = new GridBagConstraints(); + constraints.weightx = 1.0; + constraints.fill = GridBagConstraints.HORIZONTAL; + JPanel ruleInfoPanel = new JPanel(new GridBagLayout()); ruleInfoPanel.setBorder(new EmptyBorder(10, 15, 5, 15)); @@ -67,7 +66,7 @@ private void initComponents() { constraints.gridx = 1; JTabbedPane configTabbedPanel = new JTabbedPane(); - String[] settingMode = new String[]{"Exclude suffix", "Block host"}; + String[] settingMode = new String[]{"Exclude suffix", "Block host", "Exclude status"}; JPanel settingPanel = createConfigTablePanel(settingMode, "Setting"); JPanel scopePanel = getScopePanel(); JScrollPane scopeScrollPane = new JScrollPane(scopePanel); @@ -148,6 +147,12 @@ public void tableChanged(TableModelEvent e) { configLoader.setBlockHost(values); } } + + if (selected.equals("Exclude status")) { + if (!values.equals(configLoader.getExcludeStatus()) && !values.isEmpty()) { + configLoader.setExcludeStatus(values); + } + } } }; } @@ -166,6 +171,10 @@ public void actionPerformed(ActionEvent e) { if (selected.equals("Block host")) { addDataToTable(configLoader.getBlockHost().replaceAll("\\|", "\r\n"), model); } + + if (selected.equals("Exclude status")) { + addDataToTable(configLoader.getExcludeStatus().replaceAll("\\|", "\r\n"), model); + } } }; } @@ -211,6 +220,10 @@ public void actionPerformed(ActionEvent e) { } private JPanel createConfigTablePanel(String[] mode, String type) { + GridBagConstraints constraints = new GridBagConstraints(); + constraints.weightx = 1.0; + constraints.fill = GridBagConstraints.HORIZONTAL; + JPanel settingPanel = new JPanel(new BorderLayout()); DefaultTableModel model = new DefaultTableModel(); @@ -255,7 +268,7 @@ private JPanel createConfigTablePanel(String[] mode, String type) { constraints.gridy = 4; buttonPanel.add(clearButton, constraints); - addTextField = new JTextField(); + JTextField addTextField = new JTextField(); UIEnhancer.setTextFieldPlaceholder(addTextField, defaultText); inputPanelB.add(addTextField, BorderLayout.CENTER); @@ -266,13 +279,13 @@ private JPanel createConfigTablePanel(String[] mode, String type) { settingPanel.add(inputPanel, BorderLayout.CENTER); - addButton.addActionListener(e -> addActionPerformed(e, model)); + addButton.addActionListener(e -> addActionPerformed(e, model, addTextField)); addTextField.addKeyListener(new KeyAdapter() { @Override public void keyPressed(KeyEvent e) { if (e.getKeyCode() == KeyEvent.VK_ENTER) { - addActionPerformed(null, model); + addActionPerformed(null, model, addTextField); } } }); @@ -372,8 +385,9 @@ public void updateScope(JCheckBox checkBox) { configLoader.setScope(String.join("|", HaEScope)); } - private void addActionPerformed(ActionEvent e, DefaultTableModel model) { + private void addActionPerformed(ActionEvent e, DefaultTableModel model, JTextField addTextField) { String addTextFieldText = addTextField.getText(); + api.logging().logToOutput(addTextFieldText); if (!addTextFieldText.equals(defaultText)) { addDataToTable(addTextFieldText, model); } diff --git a/src/main/java/hae/instances/editor/RequestEditor.java b/src/main/java/hae/instances/editor/RequestEditor.java index 75cb9a7..0853e8e 100644 --- a/src/main/java/hae/instances/editor/RequestEditor.java +++ b/src/main/java/hae/instances/editor/RequestEditor.java @@ -12,6 +12,7 @@ import hae.component.board.table.Datatable; import hae.instances.http.utils.MessageProcessor; import hae.utils.ConfigLoader; +import hae.utils.http.HttpUtils; import hae.utils.string.StringProcessor; import javax.swing.*; @@ -37,6 +38,7 @@ public ExtensionProvidedHttpRequestEditor provideHttpRequestEditor(EditorCreatio private static class Editor implements ExtensionProvidedHttpRequestEditor { private final MontoyaApi api; private final ConfigLoader configLoader; + private final HttpUtils httpUtils; private final EditorCreationContext creationContext; private final MessageProcessor messageProcessor; private HttpRequestResponse requestResponse; @@ -47,6 +49,7 @@ private static class Editor implements ExtensionProvidedHttpRequestEditor { public Editor(MontoyaApi api, ConfigLoader configLoader, EditorCreationContext creationContext) { this.api = api; this.configLoader = configLoader; + this.httpUtils = new HttpUtils(api, configLoader); this.creationContext = creationContext; this.messageProcessor = new MessageProcessor(api); } @@ -69,16 +72,10 @@ public synchronized boolean isEnabledFor(HttpRequestResponse requestResponse) { try { String host = StringProcessor.getHostByUrl(request.url()); if (!host.isEmpty()) { - String[] hostList = configLoader.getBlockHost().split("\\|"); - boolean isBlockHost = isBlockHost(hostList, host); - - List suffixList = Arrays.asList(configLoader.getExcludeSuffix().split("\\|")); String toolType = creationContext.toolSource().toolType().toolName(); - boolean isToolScope = configLoader.getScope().contains(toolType); - - boolean matches = suffixList.contains(request.fileExtension().toLowerCase()) || isBlockHost || !isToolScope; + boolean matches = httpUtils.verifyHttpRequestResponse(requestResponse, toolType); - if (!matches && !request.bodyToString().equals("Loading...")) { + if (!matches) { this.dataList = messageProcessor.processRequest("", request, false); return isListHasData(this.dataList); } @@ -121,19 +118,6 @@ public boolean isModified() { } } - public static boolean isBlockHost(String[] hostList, String host) { - boolean isBlockHost = false; - for (String hostName : hostList) { - String cleanedHost = StringProcessor.replaceFirstOccurrence(hostName, "*.", ""); - if (hostName.contains("*.") && StringProcessor.matchFromEnd(host, cleanedHost)) { - isBlockHost = true; - } else if (host.equals(hostName) || hostName.equals("*")) { - isBlockHost = true; - } - } - return isBlockHost; - } - public static boolean isListHasData(List> dataList) { if (dataList != null && !dataList.isEmpty()) { Map dataMap = dataList.get(0); diff --git a/src/main/java/hae/instances/editor/ResponseEditor.java b/src/main/java/hae/instances/editor/ResponseEditor.java index e7eccd8..d8d4af2 100644 --- a/src/main/java/hae/instances/editor/ResponseEditor.java +++ b/src/main/java/hae/instances/editor/ResponseEditor.java @@ -13,11 +13,11 @@ import hae.component.board.table.Datatable; import hae.instances.http.utils.MessageProcessor; import hae.utils.ConfigLoader; +import hae.utils.http.HttpUtils; import hae.utils.string.StringProcessor; import javax.swing.*; import java.awt.*; -import java.util.Arrays; import java.util.List; import java.util.Map; @@ -38,6 +38,7 @@ public ExtensionProvidedHttpResponseEditor provideHttpResponseEditor(EditorCreat private static class Editor implements ExtensionProvidedHttpResponseEditor { private final MontoyaApi api; private final ConfigLoader configLoader; + private final HttpUtils httpUtils; private final EditorCreationContext creationContext; private final MessageProcessor messageProcessor; private HttpRequestResponse requestResponse; @@ -48,6 +49,7 @@ private static class Editor implements ExtensionProvidedHttpResponseEditor { public Editor(MontoyaApi api, ConfigLoader configLoader, EditorCreationContext creationContext) { this.api = api; this.configLoader = configLoader; + this.httpUtils = new HttpUtils(api, configLoader); this.creationContext = creationContext; this.messageProcessor = new MessageProcessor(api); } @@ -75,20 +77,14 @@ public synchronized boolean isEnabledFor(HttpRequestResponse requestResponse) { try { String host = StringProcessor.getHostByUrl(request.url()); if (!host.isEmpty()) { - String[] hostList = configLoader.getBlockHost().split("\\|"); - boolean isBlockHost = RequestEditor.isBlockHost(hostList, host); - - List suffixList = Arrays.asList(configLoader.getExcludeSuffix().split("\\|")); String toolType = creationContext.toolSource().toolType().toolName(); - boolean isToolScope = configLoader.getScope().contains(toolType); - - matches = suffixList.contains(request.fileExtension().toLowerCase()) || isBlockHost || !isToolScope; + matches = httpUtils.verifyHttpRequestResponse(requestResponse, toolType); } } catch (Exception ignored) { } } - if (!matches && !response.bodyToString().equals("Loading...")) { + if (!matches) { this.dataList = messageProcessor.processResponse("", response, false); return RequestEditor.isListHasData(this.dataList); } diff --git a/src/main/java/hae/instances/http/HttpMessageHandler.java b/src/main/java/hae/instances/http/HttpMessageHandler.java index d382a42..bb3c10d 100644 --- a/src/main/java/hae/instances/http/HttpMessageHandler.java +++ b/src/main/java/hae/instances/http/HttpMessageHandler.java @@ -7,20 +7,20 @@ import burp.api.montoya.http.message.HttpRequestResponse; import burp.api.montoya.http.message.requests.HttpRequest; import hae.component.board.message.MessageTableModel; -import hae.instances.editor.RequestEditor; import hae.instances.http.utils.MessageProcessor; import hae.utils.ConfigLoader; +import hae.utils.http.HttpUtils; import hae.utils.string.StringProcessor; import javax.swing.*; import java.util.ArrayList; -import java.util.Arrays; import java.util.List; import java.util.Map; public class HttpMessageHandler implements HttpHandler { private final MontoyaApi api; private final ConfigLoader configLoader; + private final HttpUtils httpUtils; private final MessageTableModel messageTableModel; private final MessageProcessor messageProcessor; @@ -29,12 +29,12 @@ public class HttpMessageHandler implements HttpHandler { private final ThreadLocal host = ThreadLocal.withInitial(() -> ""); private final ThreadLocal> colorList = ThreadLocal.withInitial(ArrayList::new); private final ThreadLocal> commentList = ThreadLocal.withInitial(ArrayList::new); - private final ThreadLocal matches = ThreadLocal.withInitial(() -> false); private final ThreadLocal httpRequest = new ThreadLocal<>(); public HttpMessageHandler(MontoyaApi api, ConfigLoader configLoader, MessageTableModel messageTableModel) { this.api = api; this.configLoader = configLoader; + this.httpUtils = new HttpUtils(api, configLoader); this.messageTableModel = messageTableModel; this.messageProcessor = new MessageProcessor(api); } @@ -49,20 +49,6 @@ public RequestToBeSentAction handleHttpRequestToBeSent(HttpRequestToBeSent httpR try { httpRequest.set(httpRequestToBeSent); host.set(StringProcessor.getHostByUrl(httpRequestToBeSent.url())); - - String[] hostList = configLoader.getBlockHost().split("\\|"); - boolean isBlockHost = RequestEditor.isBlockHost(hostList, host.get()); - - String toolType = httpRequestToBeSent.toolSource().toolType().toolName(); - boolean isToolScope = configLoader.getScope().contains(toolType); - - List suffixList = Arrays.asList(configLoader.getExcludeSuffix().split("\\|")); - matches.set(suffixList.contains(httpRequestToBeSent.fileExtension().toLowerCase()) || isBlockHost || !isToolScope); - - if (!matches.get()) { - List> result = messageProcessor.processRequest(host.get(), httpRequestToBeSent, true); - setColorAndCommentList(result); - } } catch (Exception e) { api.logging().logToError("handleHttpRequestToBeSent: " + e.getMessage()); } @@ -73,33 +59,43 @@ public RequestToBeSentAction handleHttpRequestToBeSent(HttpRequestToBeSent httpR @Override public ResponseReceivedAction handleHttpResponseReceived(HttpResponseReceived httpResponseReceived) { Annotations annotations = httpResponseReceived.annotations(); - - if (!matches.get()) { - List> result = messageProcessor.processResponse(host.get(), httpResponseReceived, true); - setColorAndCommentList(result); - // 设置高亮颜色和注释 - if (!colorList.get().isEmpty() && !commentList.get().isEmpty()) { - String color = messageProcessor.retrieveFinalColor(messageProcessor.retrieveColorIndices(colorList.get())); - annotations.setHighlightColor(HighlightColor.highlightColor(color)); - String comment = StringProcessor.mergeComment(String.join(", ", commentList.get())); - annotations.setNotes(comment); - - HttpRequestResponse httpRequestResponse = HttpRequestResponse.httpRequestResponse(httpRequest.get(), httpResponseReceived); - - // 添加到Databoard - String method = httpRequest.get().method(); - String url = httpRequest.get().url(); - String status = String.valueOf(httpResponseReceived.statusCode()); - String length = String.valueOf(httpResponseReceived.toByteArray().length()); - - // 后台提交,防止线程阻塞 - new SwingWorker() { - @Override - protected Void doInBackground() { - messageTableModel.add(httpRequestResponse, url, method, status, length, comment, color, "", ""); - return null; - } - }.run(); + HttpRequest request = httpResponseReceived.initiatingRequest(); + HttpRequestResponse requestResponse = HttpRequestResponse.httpRequestResponse(request, httpResponseReceived); + String toolType = httpResponseReceived.toolSource().toolType().toolName(); + + boolean matches = httpUtils.verifyHttpRequestResponse(requestResponse, toolType); + + if (!matches) { + try { + setColorAndCommentList(messageProcessor.processRequest(host.get(), request, true)); + setColorAndCommentList(messageProcessor.processResponse(host.get(), httpResponseReceived, true)); + + // 设置高亮颜色和注释 + if (!colorList.get().isEmpty() && !commentList.get().isEmpty()) { + String color = messageProcessor.retrieveFinalColor(messageProcessor.retrieveColorIndices(colorList.get())); + annotations.setHighlightColor(HighlightColor.highlightColor(color)); + String comment = StringProcessor.mergeComment(String.join(", ", commentList.get())); + annotations.setNotes(comment); + + HttpRequestResponse httpRequestResponse = HttpRequestResponse.httpRequestResponse(httpRequest.get(), httpResponseReceived); + + // 添加到Databoard + String method = httpRequest.get().method(); + String url = httpRequest.get().url(); + String status = String.valueOf(httpResponseReceived.statusCode()); + String length = String.valueOf(httpResponseReceived.toByteArray().length()); + + // 后台提交,防止线程阻塞 + new SwingWorker() { + @Override + protected Void doInBackground() { + messageTableModel.add(httpRequestResponse, url, method, status, length, comment, color, "", ""); + return null; + } + }.run(); + } + } catch (Exception e) { + api.logging().logToError("handleHttpResponseReceived: " + e.getMessage()); } } diff --git a/src/main/java/hae/utils/ConfigLoader.java b/src/main/java/hae/utils/ConfigLoader.java index fbdd01e..ca48565 100644 --- a/src/main/java/hae/utils/ConfigLoader.java +++ b/src/main/java/hae/utils/ConfigLoader.java @@ -147,11 +147,15 @@ public String getAIPrompt() { } public String getBlockHost() { - return getValueFromConfig("blockHost", Config.host); + return getValueFromConfig("BlockHost", Config.host); } public String getExcludeSuffix() { - return getValueFromConfig("excludeSuffix", Config.suffix); + return getValueFromConfig("ExcludeSuffix", Config.suffix); + } + + public String getExcludeStatus() { + return getValueFromConfig("ExcludeStatus", Config.status); } public String getScope() { @@ -189,11 +193,15 @@ public void setAIPrompt(String prompt) { } public void setExcludeSuffix(String excludeSuffix) { - setValueToConfig("excludeSuffix", excludeSuffix); + setValueToConfig("ExcludeSuffix", excludeSuffix); } public void setBlockHost(String blockHost) { - setValueToConfig("blockHost", blockHost); + setValueToConfig("BlockHost", blockHost); + } + + public void setExcludeStatus(String status) { + setValueToConfig("ExcludeStatus", status); } public void setScope(String scope) { diff --git a/src/main/java/hae/utils/http/HttpUtils.java b/src/main/java/hae/utils/http/HttpUtils.java index efd1116..79e0651 100644 --- a/src/main/java/hae/utils/http/HttpUtils.java +++ b/src/main/java/hae/utils/http/HttpUtils.java @@ -1,16 +1,24 @@ package hae.utils.http; import burp.api.montoya.MontoyaApi; -import burp.api.montoya.http.HttpService; +import burp.api.montoya.http.message.HttpRequestResponse; import burp.api.montoya.http.message.requests.HttpRequest; import burp.api.montoya.http.message.requests.HttpTransformation; +import burp.api.montoya.http.message.responses.HttpResponse; import burp.api.montoya.utilities.RandomUtils; +import hae.utils.ConfigLoader; +import hae.utils.string.StringProcessor; + +import java.util.Arrays; +import java.util.List; public class HttpUtils { private final MontoyaApi api; + private final ConfigLoader configLoader; - public HttpUtils(MontoyaApi api) { + public HttpUtils(MontoyaApi api, ConfigLoader configLoader) { this.api = api; + this.configLoader = configLoader; } public HttpRequest generateRequestByMultipartUploadMethod(String url, String name, String filename, String content) { @@ -28,15 +36,40 @@ public HttpRequest generateRequestByMultipartUploadMethod(String url, String nam return baseRequest; } - public HttpRequest generateRequestByJsonMethod(String url, String data) { - HttpRequest baseRequest = HttpRequest.httpRequestFromUrl(url).withTransformationApplied(HttpTransformation.TOGGLE_METHOD); - HttpService baseService = baseRequest.httpService(); - String requestString = baseRequest.toString().replace("application/x-www-form-urlencoded", "application/json"); - baseRequest = HttpRequest.httpRequest(baseService, requestString).withBody(data); - return baseRequest; - } public HttpRequest generateRequestByDeleteMethod(String url) { return HttpRequest.httpRequestFromUrl(url).withMethod("DELETE"); } + + public boolean verifyHttpRequestResponse(HttpRequestResponse requestResponse, String toolType) { + HttpRequest request = requestResponse.request(); + HttpResponse response = requestResponse.response(); + + String host = StringProcessor.getHostByUrl(request.url()); + String[] hostList = configLoader.getBlockHost().split("\\|"); + boolean isBlockHost = isBlockHost(hostList, host); + + List suffixList = Arrays.asList(configLoader.getExcludeSuffix().split("\\|")); + boolean isExcludeSuffix = suffixList.contains(request.fileExtension().toLowerCase()); + + boolean isToolScope = !configLoader.getScope().contains(toolType); + + List statusList = Arrays.asList(configLoader.getExcludeStatus().split("\\|")); + boolean isExcludeStatus = statusList.contains(String.valueOf(response.statusCode())); + + return isExcludeSuffix || isBlockHost || isToolScope || isExcludeStatus; + } + + private boolean isBlockHost(String[] hostList, String host) { + boolean isBlockHost = false; + for (String hostName : hostList) { + String cleanedHost = StringProcessor.replaceFirstOccurrence(hostName, "*.", ""); + if (hostName.contains("*.") && StringProcessor.matchFromEnd(host, cleanedHost)) { + isBlockHost = true; + } else if (host.equals(hostName) || hostName.equals("*")) { + isBlockHost = true; + } + } + return isBlockHost; + } }