diff --git a/src/main/java/hae/HaE.java b/src/main/java/hae/HaE.java index d8cae24..937c7b1 100644 --- a/src/main/java/hae/HaE.java +++ b/src/main/java/hae/HaE.java @@ -10,13 +10,13 @@ import hae.instances.editor.ResponseEditor; import hae.instances.editor.WebSocketEditor; import hae.instances.http.HttpMessageHandler; import hae.instances.websocket.WebSocketMessageHandler; -import hae.utils.config.ConfigLoader; +import hae.utils.ConfigLoader; public class HaE implements BurpExtension { @Override public void initialize(MontoyaApi api) { // 设置扩展名称 - String version = "3.1"; + String version = "3.2"; api.extension().setName(String.format("HaE (%s) - Highlighter and Extractor", version)); // 加载扩展后输出的项目信息 diff --git a/src/main/java/hae/component/Main.java b/src/main/java/hae/component/Main.java index 43e019a..60af1e6 100644 --- a/src/main/java/hae/component/Main.java +++ b/src/main/java/hae/component/Main.java @@ -5,7 +5,7 @@ import hae.component.board.Databoard; import hae.component.board.message.MessageTableModel; import hae.component.config.Config; import hae.component.rule.Rules; -import hae.utils.config.ConfigLoader; +import hae.utils.ConfigLoader; import javax.swing.*; import java.awt.*; diff --git a/src/main/java/hae/component/board/Databoard.java b/src/main/java/hae/component/board/Databoard.java index 004fd5a..2bdeedf 100644 --- a/src/main/java/hae/component/board/Databoard.java +++ b/src/main/java/hae/component/board/Databoard.java @@ -1,27 +1,39 @@ package hae.component.board; 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.responses.HttpResponse; import hae.Config; +import hae.component.board.message.MessageEntry; import hae.component.board.message.MessageTableModel; import hae.component.board.message.MessageTableModel.MessageTable; -import hae.utils.config.ConfigLoader; +import hae.instances.http.utils.RegularMatcher; +import hae.utils.ConfigLoader; +import hae.utils.project.ProjectProcessor; +import hae.utils.project.model.HaeFileContent; import hae.utils.string.StringProcessor; import javax.swing.*; import javax.swing.event.DocumentEvent; import javax.swing.event.DocumentListener; +import javax.swing.filechooser.FileNameExtensionFilter; import javax.swing.table.TableColumnModel; import javax.swing.table.TableModel; import javax.swing.table.TableRowSorter; import java.awt.*; import java.awt.event.*; +import java.io.File; import java.util.List; import java.util.*; import java.util.concurrent.ConcurrentHashMap; +import java.util.stream.Collectors; public class Databoard extends JPanel { private final MontoyaApi api; private final ConfigLoader configLoader; + private final ProjectProcessor projectProcessor; private final MessageTableModel messageTableModel; private JTextField hostTextField; private JTabbedPane dataTabbedPane; @@ -35,6 +47,7 @@ public class Databoard extends JPanel { public Databoard(MontoyaApi api, ConfigLoader configLoader, MessageTableModel messageTableModel) { this.api = api; this.configLoader = configLoader; + this.projectProcessor = new ProjectProcessor(api); this.messageTableModel = messageTableModel; initComponents(); @@ -50,11 +63,15 @@ public class Databoard extends JPanel { JLabel hostLabel = new JLabel("Host:"); JButton clearButton = new JButton("Clear"); + JButton exportButton = new JButton("Export"); + JButton importButton = new JButton("Import"); JButton actionButton = new JButton("Action"); - JPanel menuPanel = new JPanel(new GridLayout(1, 1)); + JPanel menuPanel = new JPanel(new GridLayout(3, 1, 0, 5)); menuPanel.setBorder(BorderFactory.createEmptyBorder(3, 3, 3, 3)); JPopupMenu menu = new JPopupMenu(); menuPanel.add(clearButton); + menuPanel.add(exportButton); + menuPanel.add(importButton); menu.add(menuPanel); hostTextField = new JTextField(); @@ -68,6 +85,8 @@ public class Databoard extends JPanel { }); clearButton.addActionListener(this::clearActionPerformed); + exportButton.addActionListener(this::exportActionPerformed); + importButton.addActionListener(this::importActionPerformed); splitPane.addComponentListener(new ComponentAdapter() { @Override @@ -204,9 +223,8 @@ public class Databoard extends JPanel { if (selectedHost.contains("*")) { // 通配符数据 selectedDataMap = new HashMap<>(); - String hostPattern = StringProcessor.replaceFirstOccurrence(selectedHost, "*.", ""); for (String key : dataMap.keySet()) { - if (key.contains(hostPattern) || selectedHost.equals("*")) { + if ((StringProcessor.matchesHostPattern(key, selectedHost) || selectedHost.equals("*")) && !key.contains("*")) { Map> ruleMap = dataMap.get(key); for (String ruleKey : ruleMap.keySet()) { List dataList = ruleMap.get(ruleKey); @@ -256,7 +274,8 @@ public class Databoard extends JPanel { return true; } else { String host = StringProcessor.getHostByUrl((String) entry.getValue(1)); - return StringProcessor.matchFromEnd(host, cleanedText); + + return StringProcessor.matchesHostPattern(host, filterText); } } }; @@ -267,7 +286,155 @@ public class Databoard extends JPanel { } private List getHostByList() { - return new ArrayList<>(Config.globalDataMap.keySet()); + if (!(Config.globalDataMap.keySet().size() == 1 && Config.globalDataMap.keySet().stream().anyMatch(key -> key.contains("*")))) { + return new ArrayList<>(Config.globalDataMap.keySet()); + } + return new ArrayList<>(); + } + + private void exportActionPerformed(ActionEvent e) { + String selectedHost = hostTextField.getText().trim(); + + if (selectedHost.isEmpty()) { + return; + } + + String exportDir = selectDirectory(true); + + if (exportDir.isEmpty()) { + return; + } + + ConcurrentHashMap>> dataMap = Config.globalDataMap; + List taskStatusList = exportData(selectedHost, exportDir, dataMap); + + if (!taskStatusList.isEmpty()) { + String exportStatusMessage = String.format("Exported File List Status:\n%s", String.join("\n", taskStatusList)); + JOptionPane.showConfirmDialog(null, exportStatusMessage, "Info", JOptionPane.YES_OPTION); + } + } + + 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("*")) + .map(entry -> exportEntry(entry, exportDir)) + .filter(Objects::nonNull) + .collect(Collectors.toList()); + } + + private String exportEntry(Map.Entry>> entry, String exportDir) { + String key = entry.getKey(); + Map> ruleMap = entry.getValue(); + + if (ruleMap == null || ruleMap.isEmpty()) { + return null; + } + + List messageEntryList = messageTableModel.getLogs(); + Map> httpMap = messageEntryList.stream() + .filter(messageEntry -> !StringProcessor.getHostByUrl(messageEntry.getUrl()).isEmpty()) + .filter(messageEntry -> StringProcessor.getHostByUrl(messageEntry.getUrl()).equals(key)) + .collect(Collectors.toMap( + MessageEntry::getUrl, + this::createHttpItemMap, + (existing, replacement) -> existing + )); + + String hostName = key.replace(":", "_"); + String filename = String.format("%s/%s.hae", exportDir, hostName); + boolean createdStatus = projectProcessor.createHaeFile(filename, key, ruleMap, httpMap); + + return String.format("Filename: %s, Status: %s", filename, createdStatus); + } + + private Map createHttpItemMap(MessageEntry entry) { + Map httpItemMap = new HashMap<>(); + httpItemMap.put("comment", entry.getComment()); + httpItemMap.put("color", entry.getColor()); + httpItemMap.put("request", entry.getRequestResponse().request().toString()); + httpItemMap.put("response", entry.getRequestResponse().response().toString()); + return httpItemMap; + } + + private void importActionPerformed(ActionEvent e) { + String exportDir = selectDirectory(false); + if (exportDir.isEmpty()) { + return; + } + + List filesWithExtension = findFilesWithExtension(new File(exportDir), ".hae"); + List taskStatusList = filesWithExtension.stream() + .map(this::importData) + .collect(Collectors.toList()); + + if (!taskStatusList.isEmpty()) { + String importStatusMessage = "Imported File List Status:\n" + String.join("\n", taskStatusList); + JOptionPane.showConfirmDialog(null, importStatusMessage, "Info", JOptionPane.YES_OPTION); + } + } + + private String importData(String filename) { + HaeFileContent haeFileContent = projectProcessor.readHaeFile(filename); + boolean readStatus = haeFileContent != null; + + if (readStatus) { + String host = haeFileContent.getHost(); + haeFileContent.getDataMap().forEach((key, value) -> RegularMatcher.putDataToGlobalMap(host, key, value)); + + haeFileContent.getHttpMap().forEach((key, httpItemMap) -> { + String comment = httpItemMap.get("comment"); + String color = httpItemMap.get("color"); + HttpRequestResponse httpRequestResponse = createHttpRequestResponse(key, httpItemMap); + messageTableModel.add(httpRequestResponse, comment, color); + }); + } + + return String.format("Filename: %s, Status: %s", filename, readStatus); + } + + private HttpRequestResponse createHttpRequestResponse(String key, Map httpItemMap) { + HttpService httpService = HttpService.httpService(key); + HttpRequest httpRequest = HttpRequest.httpRequest(httpService, httpItemMap.get("request")); + HttpResponse httpResponse = HttpResponse.httpResponse(httpItemMap.get("response")); + return HttpRequestResponse.httpRequestResponse(httpRequest, httpResponse); + } + + private List findFilesWithExtension(File directory, String extension) { + List filePaths = new ArrayList<>(); + if (directory.isDirectory()) { + File[] files = directory.listFiles(); + if (files != null) { + for (File file : files) { + if (file.isDirectory()) { + filePaths.addAll(findFilesWithExtension(file, extension)); + } else if (file.isFile() && file.getName().toLowerCase().endsWith(extension)) { + filePaths.add(file.getAbsolutePath()); + } + } + } + } + filePaths.add(directory.getAbsolutePath()); + return filePaths; + } + + private String selectDirectory(boolean forDirectories) { + JFileChooser chooser = new JFileChooser(); + chooser.setCurrentDirectory(new java.io.File(configLoader.getRulesFilePath())); + chooser.setDialogTitle(String.format("Select a Directory%s", forDirectories ? "" : " or File")); + FileNameExtensionFilter filter = new FileNameExtensionFilter(".hae Files", "hae"); + chooser.addChoosableFileFilter(filter); + chooser.setFileFilter(filter); + + chooser.setFileSelectionMode(forDirectories ? JFileChooser.DIRECTORIES_ONLY : JFileChooser.FILES_AND_DIRECTORIES); + chooser.setAcceptAllFileFilterUsed(!forDirectories); + + if (chooser.showOpenDialog(this) == JFileChooser.APPROVE_OPTION) { + File selectedDirectory = chooser.getSelectedFile(); + return selectedDirectory.getAbsolutePath(); + } + + return ""; } private void clearActionPerformed(ActionEvent e) { @@ -286,7 +453,7 @@ public class Databoard extends JPanel { Config.globalDataMap.remove(host); } - messageTableModel.deleteByHost(cleanedHost); + messageTableModel.deleteByHost(host); } } } diff --git a/src/main/java/hae/component/board/Datatable.java b/src/main/java/hae/component/board/Datatable.java index 621eb35..d0fccee 100644 --- a/src/main/java/hae/component/board/Datatable.java +++ b/src/main/java/hae/component/board/Datatable.java @@ -2,7 +2,7 @@ package hae.component.board; import burp.api.montoya.MontoyaApi; import hae.component.board.message.MessageTableModel; -import hae.utils.ui.UIEnhancer; +import hae.utils.UIEnhancer; import javax.swing.*; import javax.swing.event.DocumentEvent; diff --git a/src/main/java/hae/component/board/message/MessageTableModel.java b/src/main/java/hae/component/board/message/MessageTableModel.java index 03adb0e..48c6bc4 100644 --- a/src/main/java/hae/component/board/message/MessageTableModel.java +++ b/src/main/java/hae/component/board/message/MessageTableModel.java @@ -142,7 +142,7 @@ public class MessageTableModel extends AbstractTableModel { MessageEntry entry = log.get(i); String host = StringProcessor.getHostByUrl(entry.getUrl()); if (!host.isEmpty()) { - if (StringProcessor.matchFromEnd(host, filterText) || filterText.contains("*")) { + if (StringProcessor.matchesHostPattern(host, filterText) || filterText.contains("*")) { rowsToRemove.add(i); } } @@ -167,7 +167,7 @@ public class MessageTableModel extends AbstractTableModel { for (MessageEntry entry : log) { String host = StringProcessor.getHostByUrl(entry.getUrl()); if (!host.isEmpty()) { - if (filterText.contains("*.") && StringProcessor.matchFromEnd(host, cleanedText)) { + if (filterText.contains("*.") && StringProcessor.matchFromEnd(StringProcessor.extractHostname(host), cleanedText)) { filteredLog.add(entry); } else if (host.equals(filterText) || filterText.contains("*")) { filteredLog.add(entry); @@ -351,7 +351,6 @@ public class MessageTableModel extends AbstractTableModel { return log; } - @Override public int getRowCount() { return filteredLog.size(); diff --git a/src/main/java/hae/component/config/Config.java b/src/main/java/hae/component/config/Config.java index f4bae28..7375c83 100644 --- a/src/main/java/hae/component/config/Config.java +++ b/src/main/java/hae/component/config/Config.java @@ -2,8 +2,8 @@ package hae.component.config; import burp.api.montoya.MontoyaApi; import hae.component.rule.Rules; -import hae.utils.config.ConfigLoader; -import hae.utils.ui.UIEnhancer; +import hae.utils.ConfigLoader; +import hae.utils.UIEnhancer; import javax.swing.*; import javax.swing.border.EmptyBorder; @@ -14,6 +14,8 @@ import java.awt.*; import java.awt.datatransfer.Clipboard; import java.awt.datatransfer.DataFlavor; import java.awt.event.ActionEvent; +import java.awt.event.KeyAdapter; +import java.awt.event.KeyEvent; import java.util.ArrayList; import java.util.LinkedHashSet; import java.util.List; @@ -23,6 +25,8 @@ 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"; public Config(MontoyaApi api, ConfigLoader configLoader, Rules rules) { this.api = api; @@ -40,7 +44,7 @@ public class Config extends JPanel { constraints.fill = GridBagConstraints.HORIZONTAL; JPanel ruleInfoPanel = new JPanel(new GridBagLayout()); - ruleInfoPanel.setBorder(new EmptyBorder(5, 5, 5, 5)); + ruleInfoPanel.setBorder(new EmptyBorder(10, 15, 5, 15)); JLabel ruleLabel = new JLabel("Path:"); JTextField pathTextField = new JTextField(); @@ -112,7 +116,7 @@ public class Config extends JPanel { } if (selected.equals("Block host")) { - if (!values.equals(configLoader.getExcludeSuffix()) && !values.isEmpty()) { + if (!values.equals(configLoader.getBlockHost()) && !values.isEmpty()) { configLoader.setBlockHost(values); } } @@ -131,8 +135,7 @@ public class Config extends JPanel { constraints.gridy = 4; buttonPanel.add(clearButton, constraints); - JTextField addTextField = new JTextField(); - String defaultText = "Enter a new item"; + addTextField = new JTextField(); UIEnhancer.setTextFieldPlaceholder(addTextField, defaultText); inputPanelB.add(addTextField, BorderLayout.CENTER); @@ -142,13 +145,16 @@ public class Config extends JPanel { settingPanel.add(buttonPanel, BorderLayout.EAST); settingPanel.add(inputPanel, BorderLayout.CENTER); - addButton.addActionListener(e -> { - String addTextFieldText = addTextField.getText(); - if (!addTextFieldText.equals(defaultText)) { - addDataToTable(addTextFieldText, model); + + addButton.addActionListener(e -> addActionPerformedAction(e, model)); + + addTextField.addKeyListener(new KeyAdapter() { + @Override + public void keyPressed(KeyEvent e) { + if (e.getKeyCode() == KeyEvent.VK_ENTER) { + addActionPerformedAction(null, model); + } } - addTextField.setText(""); - addTextField.requestFocusInWindow(); }); pasteButton.addActionListener(e -> { @@ -176,7 +182,7 @@ public class Config extends JPanel { JLabel settingLabel = new JLabel("Setting:"); JPanel settingLabelPanel = new JPanel(new BorderLayout()); settingLabelPanel.add(settingLabel, BorderLayout.WEST); - settingMainPanel.setBorder(new EmptyBorder(0, 5, 10, 5)); + settingMainPanel.setBorder(new EmptyBorder(0, 15, 10, 15)); settingMainPanel.add(settingLabelPanel, BorderLayout.NORTH); settingMainPanel.add(settingPanel, BorderLayout.CENTER); @@ -184,6 +190,7 @@ public class Config extends JPanel { add(settingMainPanel, BorderLayout.CENTER); } + private String getFirstColumnDataAsString(DefaultTableModel model) { StringBuilder firstColumnData = new StringBuilder(); int numRows = model.getRowCount(); @@ -232,6 +239,15 @@ public class Config extends JPanel { } } + private void addActionPerformedAction(ActionEvent e, DefaultTableModel model) { + String addTextFieldText = addTextField.getText(); + if (!addTextFieldText.equals(defaultText)) { + addDataToTable(addTextFieldText, model); + } + addTextField.setText(""); + addTextField.requestFocusInWindow(); + } + private void onlineUpdateActionPerformed(ActionEvent e) { // 添加提示框防止用户误触导致配置更新 int retCode = JOptionPane.showConfirmDialog(null, "Do you want to update rules?", "Info", JOptionPane.YES_NO_OPTION); diff --git a/src/main/java/hae/component/rule/Rule.java b/src/main/java/hae/component/rule/Rule.java index 20b6338..72143b3 100644 --- a/src/main/java/hae/component/rule/Rule.java +++ b/src/main/java/hae/component/rule/Rule.java @@ -2,7 +2,7 @@ package hae.component.rule; import burp.api.montoya.MontoyaApi; import hae.Config; -import hae.utils.config.ConfigLoader; +import hae.utils.ConfigLoader; import hae.utils.rule.RuleProcessor; import javax.swing.*; @@ -151,7 +151,7 @@ public class Rule extends JPanel { private void ruleRemoveActionPerformed(ActionEvent e, JTable ruleTable, JTabbedPane tabbedPane) { if (ruleTable.getSelectedRowCount() >= 1) { - if (JOptionPane.showConfirmDialog(null, "Are you sure you want to delete this rule?", "Info", JOptionPane.OK_OPTION) == 0) { + if (JOptionPane.showConfirmDialog(null, "Are you sure you want to remove this rule?", "Info", JOptionPane.YES_NO_OPTION) == 0) { DefaultTableModel model = (DefaultTableModel) ruleTable.getModel(); int select = ruleTable.convertRowIndexToModel(ruleTable.getSelectedRow()); diff --git a/src/main/java/hae/component/rule/Rules.java b/src/main/java/hae/component/rule/Rules.java index e800358..1bc87f7 100644 --- a/src/main/java/hae/component/rule/Rules.java +++ b/src/main/java/hae/component/rule/Rules.java @@ -2,7 +2,7 @@ package hae.component.rule; import burp.api.montoya.MontoyaApi; import hae.Config; -import hae.utils.config.ConfigLoader; +import hae.utils.ConfigLoader; import hae.utils.rule.RuleProcessor; import javax.swing.*; diff --git a/src/main/java/hae/instances/editor/RequestEditor.java b/src/main/java/hae/instances/editor/RequestEditor.java index c9c01f3..a1d8325 100644 --- a/src/main/java/hae/instances/editor/RequestEditor.java +++ b/src/main/java/hae/instances/editor/RequestEditor.java @@ -11,7 +11,7 @@ import burp.api.montoya.ui.editor.extension.ExtensionProvidedHttpRequestEditor; import burp.api.montoya.ui.editor.extension.HttpRequestEditorProvider; import hae.component.board.Datatable; import hae.instances.http.utils.MessageProcessor; -import hae.utils.config.ConfigLoader; +import hae.utils.ConfigLoader; import hae.utils.string.StringProcessor; import javax.swing.*; diff --git a/src/main/java/hae/instances/editor/ResponseEditor.java b/src/main/java/hae/instances/editor/ResponseEditor.java index 033c368..2e5c1e5 100644 --- a/src/main/java/hae/instances/editor/ResponseEditor.java +++ b/src/main/java/hae/instances/editor/ResponseEditor.java @@ -12,7 +12,7 @@ import burp.api.montoya.ui.editor.extension.ExtensionProvidedHttpResponseEditor; import burp.api.montoya.ui.editor.extension.HttpResponseEditorProvider; import hae.component.board.Datatable; import hae.instances.http.utils.MessageProcessor; -import hae.utils.config.ConfigLoader; +import hae.utils.ConfigLoader; import hae.utils.string.StringProcessor; import javax.swing.*; diff --git a/src/main/java/hae/instances/http/HttpMessageHandler.java b/src/main/java/hae/instances/http/HttpMessageHandler.java index b5fbfcb..8610e4f 100644 --- a/src/main/java/hae/instances/http/HttpMessageHandler.java +++ b/src/main/java/hae/instances/http/HttpMessageHandler.java @@ -9,7 +9,7 @@ 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.config.ConfigLoader; +import hae.utils.ConfigLoader; import hae.utils.string.StringProcessor; import java.util.ArrayList; diff --git a/src/main/java/hae/instances/http/utils/RegularMatcher.java b/src/main/java/hae/instances/http/utils/RegularMatcher.java index dfbc21c..b92b8b6 100644 --- a/src/main/java/hae/instances/http/utils/RegularMatcher.java +++ b/src/main/java/hae/instances/http/utils/RegularMatcher.java @@ -90,43 +90,15 @@ public class RegularMatcher { result.clear(); result.addAll(tmpList); - String nameAndSize = String.format("%s (%s)", name, result.size()); if (!result.isEmpty()) { tmpMap.put("color", color); String dataStr = String.join("\n", result); tmpMap.put("data", dataStr); + + String nameAndSize = String.format("%s (%s)", name, result.size()); finalMap.put(nameAndSize, tmpMap); - // 添加到全局变量中，便于Databoard检索 - if (!Objects.equals(host, "") && host != null) { - List dataList = Arrays.asList(dataStr.split("\n")); - Config.globalDataMap.compute(host, (existingHost, existingMap) -> { - Map> gRuleMap = Optional.ofNullable(existingMap).orElse(new ConcurrentHashMap<>()); - - gRuleMap.merge(name, new ArrayList<>(dataList), (existingList, newList) -> { - Set combinedSet = new LinkedHashSet<>(existingList); - combinedSet.addAll(newList); - return new ArrayList<>(combinedSet); - }); - - return gRuleMap; - }); - - String[] splitHost = host.split("\\."); - String onlyHost = host.split(":")[0]; - - String anyHost = (splitHost.length > 2 && !onlyHost.matches("\\b(?:\\d{1,3}\\.){3}\\d{1,3}\\b")) ? StringProcessor.replaceFirstOccurrence(onlyHost, splitHost[0], "*") : ""; - - if (!Config.globalDataMap.containsKey(anyHost) && anyHost.length() > 0) { - // 添加通配符Host，实际数据从查询哪里将所有数据提取 - Config.globalDataMap.put(anyHost, new HashMap<>()); - } - - if (!Config.globalDataMap.containsKey("*")) { - // 添加通配符全匹配，同上 - Config.globalDataMap.put("*", new HashMap<>()); - } - } + putDataToGlobalMap(host, name, result); } } } @@ -136,6 +108,38 @@ public class RegularMatcher { } } + public static void putDataToGlobalMap(String host, String name, List dataList) { + // 添加到全局变量中，便于Databoard检索 + if (!Objects.equals(host, "") && host != null) { + Config.globalDataMap.compute(host, (existingHost, existingMap) -> { + Map> gRuleMap = Optional.ofNullable(existingMap).orElse(new ConcurrentHashMap<>()); + + gRuleMap.merge(name, new ArrayList<>(dataList), (existingList, newList) -> { + Set combinedSet = new LinkedHashSet<>(existingList); + combinedSet.addAll(newList); + return new ArrayList<>(combinedSet); + }); + + return gRuleMap; + }); + + String[] splitHost = host.split("\\."); + String onlyHost = host.split(":")[0]; + + String anyHost = (splitHost.length > 2 && !onlyHost.matches("\\b(?:\\d{1,3}\\.){3}\\d{1,3}\\b")) ? StringProcessor.replaceFirstOccurrence(onlyHost, splitHost[0], "*") : ""; + + if (!Config.globalDataMap.containsKey(anyHost) && anyHost.length() > 0) { + // 添加通配符Host，实际数据从查询哪里将所有数据提取 + Config.globalDataMap.put(anyHost, new HashMap<>()); + } + + if (!Config.globalDataMap.containsKey("*")) { + // 添加通配符全匹配，同上 + Config.globalDataMap.put("*", new HashMap<>()); + } + } + } + private List matchByRegex(String f_regex, String s_regex, String content, String format, String engine, boolean sensitive) { List retList = new ArrayList<>(); if ("nfa".equals(engine)) { diff --git a/src/main/java/hae/utils/config/ConfigLoader.java b/src/main/java/hae/utils/ConfigLoader.java similarity index 98% rename from src/main/java/hae/utils/config/ConfigLoader.java rename to src/main/java/hae/utils/ConfigLoader.java index 49fda87..d1ad57c 100644 --- a/src/main/java/hae/utils/config/ConfigLoader.java +++ b/src/main/java/hae/utils/ConfigLoader.java @@ -1,4 +1,4 @@ -package hae.utils.config; +package hae.utils; import burp.api.montoya.MontoyaApi; import burp.api.montoya.http.RequestOptions; @@ -178,7 +178,7 @@ public class ConfigLoader { try (InputStream in = Files.newInputStream(path)) { return yaml.load(in); - } catch (IOException e) { + } catch (Exception e) { return new LinkedHashMap<>(); // 读取失败时也返回空的Map } } @@ -189,7 +189,7 @@ public class ConfigLoader { try (Writer ws = new OutputStreamWriter(Files.newOutputStream(Paths.get(configFilePath)), StandardCharsets.UTF_8)) { yaml.dump(currentConfig, ws); - } catch (IOException ignored) { + } catch (Exception ignored) { } } @@ -199,7 +199,7 @@ public class ConfigLoader { try (Writer ws = new OutputStreamWriter(Files.newOutputStream(Paths.get(configFilePath)), StandardCharsets.UTF_8)) { yaml.dump(currentConfig, ws); - } catch (IOException ignored) { + } catch (Exception ignored) { } } @@ -225,7 +225,7 @@ public class ConfigLoader { return true; } - } catch (IOException ignored) { + } catch (Exception ignored) { } return false; diff --git a/src/main/java/hae/utils/ui/UIEnhancer.java b/src/main/java/hae/utils/UIEnhancer.java similarity index 97% rename from src/main/java/hae/utils/ui/UIEnhancer.java rename to src/main/java/hae/utils/UIEnhancer.java index ee7da70..332a73e 100644 --- a/src/main/java/hae/utils/ui/UIEnhancer.java +++ b/src/main/java/hae/utils/UIEnhancer.java @@ -1,4 +1,4 @@ -package hae.utils.ui; +package hae.utils; import javax.swing.*; import java.awt.*; diff --git a/src/main/java/hae/utils/project/ProjectProcessor.java b/src/main/java/hae/utils/project/ProjectProcessor.java new file mode 100644 index 0000000..d47cc14 --- /dev/null +++ b/src/main/java/hae/utils/project/ProjectProcessor.java @@ -0,0 +1,77 @@ +package hae.utils.project; + +import burp.api.montoya.MontoyaApi; +import hae.utils.project.model.HaeFileContent; +import org.yaml.snakeyaml.Yaml; + +import java.io.*; +import java.nio.charset.StandardCharsets; +import java.util.List; +import java.util.Map; +import java.util.zip.ZipEntry; +import java.util.zip.ZipInputStream; +import java.util.zip.ZipOutputStream; + +public class ProjectProcessor { + private final MontoyaApi api; + + public ProjectProcessor(MontoyaApi api) { + this.api = api; + } + + public boolean createHaeFile(String haeFilePath, String host, Map> dataMap, Map> httpMap) { + ByteArrayOutputStream dataYamlStream = new ByteArrayOutputStream(); + ByteArrayOutputStream httpYamlStream = new ByteArrayOutputStream(); + Yaml yaml = new Yaml(); + + yaml.dump(dataMap, new OutputStreamWriter(dataYamlStream, StandardCharsets.UTF_8)); + yaml.dump(httpMap, new OutputStreamWriter(httpYamlStream, StandardCharsets.UTF_8)); + + try (ZipOutputStream zipOut = new ZipOutputStream(new FileOutputStream(haeFilePath))) { + zipOut.putNextEntry(new ZipEntry("info.txt")); + zipOut.write(host.getBytes(StandardCharsets.UTF_8)); + zipOut.closeEntry(); + + zipOut.putNextEntry(new ZipEntry("data.yml")); + zipOut.write(dataYamlStream.toByteArray()); + zipOut.closeEntry(); + + zipOut.putNextEntry(new ZipEntry("http.yml")); + zipOut.write(httpYamlStream.toByteArray()); + zipOut.closeEntry(); + } catch (Exception e) { + api.logging().logToOutput(e.getMessage()); + return false; + } + + return true; + } + + public HaeFileContent readHaeFile(String haeFilePath) { + HaeFileContent haeFileContent = new HaeFileContent(api); + Yaml yaml = new Yaml(); + + try (ZipInputStream zipIn = new ZipInputStream(new FileInputStream(haeFilePath))) { + ZipEntry entry; + while ((entry = zipIn.getNextEntry()) != null) { + switch (entry.getName()) { + case "info.txt": + haeFileContent.setHost(new String(zipIn.readAllBytes(), StandardCharsets.UTF_8)); + break; + case "data.yml": + haeFileContent.setDataMap(yaml.load(new InputStreamReader(zipIn, StandardCharsets.UTF_8))); + break; + case "http.yml": + haeFileContent.setHttpMap(yaml.load(new InputStreamReader(zipIn, StandardCharsets.UTF_8))); + break; + } + zipIn.closeEntry(); + } + } catch (Exception e) { + api.logging().logToOutput(e.getMessage()); + return null; + } + return haeFileContent; + } +} + diff --git a/src/main/java/hae/utils/project/model/HaeFileContent.java b/src/main/java/hae/utils/project/model/HaeFileContent.java new file mode 100644 index 0000000..60e8eaa --- /dev/null +++ b/src/main/java/hae/utils/project/model/HaeFileContent.java @@ -0,0 +1,67 @@ +package hae.utils.project.model; + +import burp.api.montoya.MontoyaApi; + +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +public class HaeFileContent { + private final MontoyaApi api; + private String host; + private final Map> dataMap; + private final Map> httpMap; + + public HaeFileContent(MontoyaApi api) { + this.api = api; + this.dataMap = new HashMap<>(); + this.httpMap = new HashMap<>(); + } + + public String getHost() { + return host; + } + + public Map> getDataMap() { + return dataMap; + } + + public Map> getHttpMap() { + return httpMap; + } + + public void setHost(String host) { + this.host = host; + } + + public void setDataMap(Map> dataMap) { + for (Map.Entry> entry : dataMap.entrySet()) { + List values = new ArrayList<>(); + for (Object value : entry.getValue()) { + try { + values.add(new String((byte[]) value, StandardCharsets.UTF_8)); + } catch (Exception e) { + values.add(value.toString()); + } + } + this.dataMap.put(entry.getKey(), values); + } + } + + public void setHttpMap(Map> httpMap) { + for (Map.Entry> entry : httpMap.entrySet()) { + Map newValues = new HashMap<>(); + Map values = entry.getValue(); + for (String key : values.keySet()) { + try { + newValues.put(key, new String((byte[]) values.get(key), StandardCharsets.UTF_8)); + } catch (Exception e) { + newValues.put(key, values.get(key).toString()); + } + } + this.httpMap.put(entry.getKey(), newValues); + } + } +} \ No newline at end of file diff --git a/src/main/java/hae/utils/rule/RuleProcessor.java b/src/main/java/hae/utils/rule/RuleProcessor.java index f8769b9..ed03608 100644 --- a/src/main/java/hae/utils/rule/RuleProcessor.java +++ b/src/main/java/hae/utils/rule/RuleProcessor.java @@ -2,7 +2,7 @@ package hae.utils.rule; import burp.api.montoya.MontoyaApi; import hae.Config; -import hae.utils.config.ConfigLoader; +import hae.utils.ConfigLoader; import hae.utils.rule.model.Group; import hae.utils.rule.model.Info; import org.yaml.snakeyaml.DumperOptions; diff --git a/src/main/java/hae/utils/string/StringProcessor.java b/src/main/java/hae/utils/string/StringProcessor.java index 254fea7..7e9644b 100644 --- a/src/main/java/hae/utils/string/StringProcessor.java +++ b/src/main/java/hae/utils/string/StringProcessor.java @@ -32,6 +32,29 @@ public class StringProcessor { return patternIndex == -1; } + public static String extractHostname(String hostWithPort) { + if (hostWithPort == null || hostWithPort.isEmpty()) { + return ""; + } + int colonIndex = hostWithPort.indexOf(":"); + if (colonIndex != -1) { + return hostWithPort.substring(0, colonIndex); + } else { + return hostWithPort; + } + } + + public static boolean matchesHostPattern(String host, String selectedHost) { + String hostname = StringProcessor.extractHostname(host); + String hostPattern = selectedHost.replace("*.", ""); + boolean matchesDirectly = selectedHost.equals("*") || host.equals(selectedHost); + boolean matchesPattern = !host.contains("*") && + (hostPattern.equals(selectedHost) ? + StringProcessor.matchFromEnd(host, hostPattern) : + StringProcessor.matchFromEnd(hostname, hostPattern)); + return matchesDirectly || matchesPattern; + } + public static String mergeComment(String comment) { if (!comment.contains(",")) { return comment;