Skip to content

Commit

Permalink
Version: 4.0 Update
Browse files Browse the repository at this point in the history
  • Loading branch information
@gh0stkey
gh0stkey committed Dec 21, 2024
1 parent 1f1ca99 commit daacb2e
Show file tree
Hide file tree
Showing 24 changed files with 469 additions and 1,427 deletions.
4 changes: 1 addition & 3 deletions build.gradle
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,8 +22,6 @@ dependencies {
implementation 'org.yaml:snakeyaml:2.0'
implementation 'dk.brics.automaton:automaton:1.11-8'
implementation 'com.github.ben-manes.caffeine:caffeine:3.1.8'
implementation 'com.google.code.gson:gson:2.11.0'
implementation 'com.squareup.okhttp3:okhttp:4.12.0'
}

test {
Expand All @@ -36,4 +34,4 @@ jar {
from {
configurations.runtimeClasspath.collect { it.isDirectory() ? it : zipTree(it) }
}
}
}
Binary file removed images/config-ai.png
View file
Binary file not shown.
Binary file removed images/config-setting.png
View file
Binary file not shown.
5 changes: 0 additions & 5 deletions src/main/java/hae/Config.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -61,11 +61,6 @@ public class Config {
"gray"
};

public static String prompt = "You are a data security expert in the field of cyber security. Your task is to optimize the information provided by the user and then output it in JSON format. The user-supplied information is data that has been extracted by regular expressions. The user-supplied information is divided into two parts, the first part is RuleName which represents the name of the regular expression and the second part is MarkInfo which represents the data extracted by the regular expression. You need to find the matching or similar data in MarkInfo according to the meaning of RuleName, and output the original rows of these data in JSON format.(garbled and meaningless data rows should be removed)\n" +
"You must ensure that the extracted data is accurately expressed and correctly formatted in the JSON structure. Your output data must comply with the original MarkInfo content rows without modification, and strictly adhere to the following JSON format for return, no other text, code and formatting (e.g., line breaks, carriage returns, indentation, spaces), once the return of other irrelevant content will cause irreparable damage to the user: {\"data\":[\"data1\", \"data2\"]}.";

public static String userTextFormat = "User Input: \r\nRuleName: %s\r\nMarkInfo: %s";

public static Map<String, Object[][]> globalRules = new HashMap<>();

public static ConcurrentHashMap<String, Map<String, List<String>>> globalDataMap = new ConcurrentHashMap<>();
Expand Down
10 changes: 8 additions & 2 deletions src/main/java/hae/HaE.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,19 +12,21 @@
import hae.instances.editor.WebSocketEditor;
import hae.instances.websocket.WebSocketMessageHandler;
import hae.utils.ConfigLoader;
import hae.utils.DataManager;

public class HaE implements BurpExtension {
@Override
public void initialize(MontoyaApi api) {
// 设置扩展名称
String version = "3.4";
api.extension().setName(String.format("HaE (%s) - Highlighter and Extractor", version));
String version = "4.0";
api.extension().setName("HaE - Highlighter and Extractor");

// 加载扩展后输出的项目信息
Logging logging = api.logging();
logging.logToOutput("[ HACK THE WORLD - TO DO IT ]");
logging.logToOutput("[#] Author: EvilChen && 0chencc && vaycore");
logging.logToOutput("[#] Github: https://github.com/gh0stkey/HaE");
logging.logToOutput("[#] Version: " + version);

// 配置文件加载
ConfigLoader configLoader = new ConfigLoader(api);
Expand All @@ -42,6 +44,10 @@ public void initialize(MontoyaApi api) {
api.userInterface().registerHttpResponseEditorProvider(new ResponseEditor(api, configLoader));
api.userInterface().registerWebSocketMessageEditorProvider(new WebSocketEditor(api, configLoader));

// 从BurpSuite里加载数据
DataManager dataManager = new DataManager(api);
dataManager.loadData(messageTableModel);

api.extension().registerUnloadingHandler(new ExtensionUnloadingHandler() {
@Override
public void extensionUnloaded() {
Expand Down
110 changes: 20 additions & 90 deletions src/main/java/hae/component/Config.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,6 @@ public class Config extends JPanel {
private final ConfigLoader configLoader;
private final MessageTableModel messageTableModel;
private final Rules rules;
private final String defaultText = "Enter a new item";

private Registration activeHandler;
private Registration passiveHandler;
Expand Down Expand Up @@ -61,22 +60,22 @@ private void initComponents() {
pathTextField.setEditable(false);
pathTextField.setText(configLoader.getRulesFilePath());
JButton reloadButton = new JButton("Reload");
JButton updateButton = new JButton("Update");
JButton reinitButton = new JButton("Reinit");
ruleInfoPanel.add(ruleLabel);
ruleInfoPanel.add(pathTextField, constraints);
ruleInfoPanel.add(Box.createHorizontalStrut(5));
ruleInfoPanel.add(reloadButton);
ruleInfoPanel.add(reinitButton);
ruleInfoPanel.add(Box.createHorizontalStrut(5));
ruleInfoPanel.add(updateButton);
ruleInfoPanel.add(reloadButton);

reloadButton.addActionListener(this::reloadActionPerformed);
updateButton.addActionListener(this::onlineUpdateActionPerformed);
reinitButton.addActionListener(this::reinitActionPerformed);

constraints.gridx = 1;
JTabbedPane configTabbedPanel = new JTabbedPane();

String[] settingMode = new String[]{"Exclude suffix", "Block host", "Exclude status"};
JPanel settingPanel = createConfigTablePanel(settingMode, "Setting");
JPanel settingPanel = createConfigTablePanel(settingMode);

JPanel northPanel = new JPanel(new BorderLayout());

Expand Down Expand Up @@ -105,38 +104,6 @@ public void componentResized(ComponentEvent e) {
settingPanel.add(northPanel, BorderLayout.NORTH);

configTabbedPanel.add("Setting", settingPanel);

String[] aiMode = new String[]{"Alibaba", "Moonshot"};
JPanel aiPanel = createConfigTablePanel(aiMode, "AI+");
JTextArea promptTextArea = new JTextArea();
promptTextArea.setLineWrap(true);
promptTextArea.getDocument().addDocumentListener(new DocumentListener() {
@Override
public void insertUpdate(DocumentEvent e) {
onTextChange();
}

@Override
public void removeUpdate(DocumentEvent e) {
onTextChange();
}

@Override
public void changedUpdate(DocumentEvent e) {
onTextChange();
}

private void onTextChange() {
String promptText = promptTextArea.getText();
configLoader.setAIPrompt(promptText);
}
});
promptTextArea.setText(configLoader.getAIPrompt());
JScrollPane promptScrollPane = new JScrollPane(promptTextArea);
promptScrollPane.setBorder(new TitledBorder("Prompt"));
promptScrollPane.setPreferredSize(new Dimension(0, 100));
aiPanel.add(promptScrollPane, BorderLayout.NORTH);
configTabbedPanel.add("AI+", aiPanel);
add(ruleInfoPanel, BorderLayout.NORTH);
add(configTabbedPanel, BorderLayout.CENTER);
}
Expand Down Expand Up @@ -256,47 +223,8 @@ public void actionPerformed(ActionEvent e) {
};
}

private TableModelListener craeteAITableModelListener(JComboBox<String> setTypeComboBox, DefaultTableModel model) {
return new TableModelListener() {
@Override
public void tableChanged(TableModelEvent e) {
String selected = (String) setTypeComboBox.getSelectedItem();
String values = getFirstColumnDataAsString(model);

if (selected.equals("Alibaba")) {
if (!values.equals(configLoader.getAlibabaAIAPIKey()) && !values.isEmpty()) {
configLoader.setAlibabaAIAPIKey(values);
}
}

if (selected.equals("Moonshot")) {
if (!values.equals(configLoader.getMoonshotAIAPIKey()) && !values.isEmpty()) {
configLoader.setMoonshotAIAPIKey(values);
}
}
}
};
}

private ActionListener createAIActionListener(JComboBox<String> setTypeComboBox, DefaultTableModel model) {
return new ActionListener() {
@Override
public void actionPerformed(ActionEvent e) {
String selected = (String) setTypeComboBox.getSelectedItem();
model.setRowCount(0);

if (selected.equals("Alibaba")) {
addDataToTable(configLoader.getAlibabaAIAPIKey().replaceAll("\\|", "\r\n"), model);
}

if (selected.equals("Moonshot")) {
addDataToTable(configLoader.getMoonshotAIAPIKey().replaceAll("\\|", "\r\n"), model);
}
}
};
}

private JPanel createConfigTablePanel(String[] mode, String type) {
private JPanel createConfigTablePanel(String[] mode) {
GridBagConstraints constraints = new GridBagConstraints();
constraints.weightx = 1.0;
constraints.fill = GridBagConstraints.HORIZONTAL;
Expand Down Expand Up @@ -327,9 +255,9 @@ private JPanel createConfigTablePanel(String[] mode, String type) {
JComboBox<String> setTypeComboBox = new JComboBox<>();
setTypeComboBox.setModel(new DefaultComboBoxModel<>(mode));

model.addTableModelListener(type.equals("AI+") ? craeteAITableModelListener(setTypeComboBox, model) : craeteSettingTableModelListener(setTypeComboBox, model));
model.addTableModelListener(craeteSettingTableModelListener(setTypeComboBox, model));

setTypeComboBox.addActionListener(type.equals("AI+") ? createAIActionListener(setTypeComboBox, model) : createSettingActionListener(setTypeComboBox, model));
setTypeComboBox.addActionListener(createSettingActionListener(setTypeComboBox, model));

setTypeComboBox.setSelectedItem(mode[0]);

Expand All @@ -346,6 +274,7 @@ private JPanel createConfigTablePanel(String[] mode, String type) {
buttonPanel.add(clearButton, constraints);

JTextField addTextField = new JTextField();
String defaultText = "Enter a new item";
UIEnhancer.setTextFieldPlaceholder(addTextField, defaultText);

inputPanelB.add(addTextField, BorderLayout.CENTER);
Expand Down Expand Up @@ -390,7 +319,7 @@ public void keyPressed(KeyEvent e) {
JPanel settingMainPanel = new JPanel(new BorderLayout());
settingMainPanel.setBorder(new EmptyBorder(5, 15, 10, 15));
JScrollPane settingScroller = new JScrollPane(settingPanel);
settingScroller.setBorder(new TitledBorder(type.equals("AI+") ? "API Key" : "Setting"));
settingScroller.setBorder(new TitledBorder("Setting"));
settingMainPanel.add(settingScroller, BorderLayout.CENTER);

return settingMainPanel;
Expand Down Expand Up @@ -492,16 +421,17 @@ private void addActionPerformed(ActionEvent e, DefaultTableModel model, JTextFie
}
}

private void onlineUpdateActionPerformed(ActionEvent e) {
// 添加提示框防止用户误触导致配置更新
int retCode = JOptionPane.showConfirmDialog(this, "Do you want to update rules?", "Info", JOptionPane.YES_NO_OPTION);
if (retCode == JOptionPane.YES_OPTION) {
configLoader.initRulesByNet();
reloadActionPerformed(null);
}
}

private void reloadActionPerformed(ActionEvent e) {
rules.reloadRuleGroup();
}

private void reinitActionPerformed(ActionEvent e) {
int retCode = JOptionPane.showConfirmDialog(this, "Do you want to reinitialize rules? This action will overwrite your existing rules.", "Info", JOptionPane.YES_NO_OPTION);
if (retCode == JOptionPane.YES_OPTION) {
boolean ret = configLoader.initRules();
if (ret) {
rules.reloadRuleGroup();
}
}
}
}
3 changes: 1 addition & 2 deletions src/main/java/hae/component/Main.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -81,7 +81,6 @@ private ImageIcon getImageIcon(boolean isDark) {
ImageIcon originalIcon = new ImageIcon(imageURL);
Image originalImage = originalIcon.getImage();
Image scaledImage = originalImage.getScaledInstance(30, 20, Image.SCALE_FAST);
ImageIcon scaledIcon = new ImageIcon(scaledImage);
return scaledIcon;
return new ImageIcon(scaledImage);
}
}
Loading

0 comments on commit daacb2e

Please sign in to comment.