Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion src/HaENet/sources/src/main/java/hae/component/Config.java
Original file line number Diff line number Diff line change
Expand Up @@ -468,7 +468,7 @@ private void reinitActionPerformed(ActionEvent e) {
JOptionPane.YES_NO_OPTION
);
if (retCode == JOptionPane.YES_OPTION) {
boolean ret = configLoader.initRules();
boolean ret = configLoader.initTemplates();
if (ret) {
rules.reloadRuleGroup();
}
Expand Down
171 changes: 107 additions & 64 deletions src/HaENet/sources/src/main/java/hae/utils/ConfigLoader.java
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ public class ConfigLoader {
private final MontoyaApi api;
private final Yaml yaml;
private final String configFilePath;
private final String rulesFilePath;
private final String templatesPath;
private volatile Map<String, Object> configCache;

public ConfigLoader(MontoyaApi api) {
Expand All @@ -30,7 +30,7 @@ public ConfigLoader(MontoyaApi api) {

String configPath = determineConfigPath();
this.configFilePath = String.format("%s/%s", configPath, "Config.yml");
this.rulesFilePath = String.format("%s/%s", configPath, "Rules.yml");
this.templatesPath = String.format("%s/%s", configPath, "templates");

// 构造函数,初始化配置
File configDir = new File(configPath);
Expand All @@ -43,9 +43,9 @@ public ConfigLoader(MontoyaApi api) {
initConfig();
}

File rulesFilePath = new File(this.rulesFilePath);
if (!(rulesFilePath.exists() && rulesFilePath.isFile())) {
initRules();
File templatesDir = new File(this.templatesPath);
if (!(templatesDir.exists() && templatesDir.isDirectory())) {
initTemplates();
}
}

Expand Down Expand Up @@ -118,49 +118,85 @@ public void initConfig() {
}

public String getRulesFilePath() {
return rulesFilePath;
return templatesPath;
}

public String getTemplatesPath() {
return templatesPath;
}

// 获取规则配置
public Map<String, List<RuleDefinition>> getRules() {
Map<String, List<RuleDefinition>> rules = new HashMap<>();
File templatesDir = new File(templatesPath);
if (!templatesDir.exists()) {
templatesDir.mkdirs();
}
return loadRulesFromTemplates(templatesDir);
}

try {
InputStream inputStream = Files.newInputStream(
Paths.get(getRulesFilePath())
);
Map<String, Object> rulesMap = yaml.load(inputStream);

Object rulesObj = rulesMap.get("rules");
if (rulesObj instanceof List) {
List<Map<String, Object>> groupData = (List<
Map<String, Object>
>) rulesObj;
for (Map<String, Object> groupFields : groupData) {
List<RuleDefinition> data = new ArrayList<>();

Object ruleObj = groupFields.get("rule");
if (ruleObj instanceof List) {
List<Map<String, Object>> ruleData = (List<
Map<String, Object>
>) ruleObj;
for (Map<String, Object> ruleFields : ruleData) {
data.add(RuleDefinition.fromYamlMap(ruleFields));
}
}
private Map<String, List<RuleDefinition>> loadRulesFromTemplates(File templatesDir) {
Map<String, List<RuleDefinition>> rules = new LinkedHashMap<>();

File[] groupDirs = templatesDir.listFiles(File::isDirectory);
if (groupDirs == null) return rules;

rules.put(groupFields.get("group").toString(), data);
for (File groupDir : groupDirs) {
String groupName = groupDirToName(groupDir.getName());
List<RuleDefinition> groupRules = new ArrayList<>();

File[] yamlFiles = groupDir.listFiles(f -> f.getName().endsWith(".yaml") || f.getName().endsWith(".yml"));
if (yamlFiles == null) continue;

for (File yamlFile : yamlFiles) {
try (InputStream in = Files.newInputStream(yamlFile.toPath())) {
Map<String, Object> tmpl = yaml.load(in);
if (tmpl != null && tmpl.containsKey("file")) {
groupRules.add(RuleDefinition.fromTemplateYaml(tmpl, groupName));
}
} catch (Exception e) {
api.logging().logToError("Failed to load template " + yamlFile.getName() + ": " + e.getMessage());
}
}

return rules;
} catch (Exception e) {
api.logging().logToError("Failed to load rules: " + e.getMessage());
if (!groupRules.isEmpty()) {
rules.put(groupName, groupRules);
}
}

return rules;
}

private static String groupDirToName(String dirName) {
switch (dirName) {
case "fingerprint": return "Fingerprint";
case "vulnerability": return "Maybe Vulnerability";
case "basic-info": return "Basic Information";
case "sensitive-info": return "Sensitive Information";
case "other": return "Other";
default:
StringBuilder sb = new StringBuilder();
for (String part : dirName.split("-")) {
if (!part.isEmpty()) {
sb.append(Character.toUpperCase(part.charAt(0)));
sb.append(part.substring(1));
sb.append(' ');
}
}
return sb.toString().trim();
}
}

public static String groupNameToDir(String groupName) {
switch (groupName) {
case "Fingerprint": return "fingerprint";
case "Maybe Vulnerability": return "vulnerability";
case "Basic Information": return "basic-info";
case "Sensitive Information": return "sensitive-info";
case "Other": return "other";
default:
return groupName.toLowerCase().replace(' ', '-');
}
}

public String getBlockHost() {
return getValueFromConfig("BlockHost", AppConstants.host);
}
Expand Down Expand Up @@ -277,40 +313,47 @@ private Map<String, Object> loadCurrentConfig() {
}
}

public boolean initRules() {
boolean ret = copyRulesToFile(this.rulesFilePath);
if (!ret) {
api.extension().unload();
}
return ret;
}

private boolean copyRulesToFile(String targetFilePath) {
InputStream inputStream = getClass()
.getClassLoader()
.getResourceAsStream("rules/Rules.yml");
File targetFile = new File(targetFilePath);
public boolean initTemplates() {
String[] groups = {"fingerprint", "vulnerability", "basic-info", "sensitive-info", "other"};
boolean success = true;

try (
inputStream;
OutputStream outputStream = new FileOutputStream(targetFile)
) {
if (inputStream != null) {
byte[] buffer = new byte[1024];
int length;
for (String group : groups) {
File groupDir = new File(templatesPath, group);
if (!groupDir.exists()) {
groupDir.mkdirs();
}

while ((length = inputStream.read(buffer)) > 0) {
outputStream.write(buffer, 0, length);
String resourceBase = "templates/" + group + "/";
try {
InputStream indexStream = getClass().getClassLoader().getResourceAsStream(resourceBase + "index.txt");
if (indexStream == null) continue;

String content = new String(indexStream.readAllBytes(), StandardCharsets.UTF_8);
for (String fileName : content.split("\n")) {
fileName = fileName.trim();
if (fileName.isEmpty()) continue;

InputStream fileStream = getClass().getClassLoader().getResourceAsStream(resourceBase + fileName);
if (fileStream == null) continue;

File target = new File(groupDir, fileName);
try (fileStream; OutputStream out = new FileOutputStream(target)) {
byte[] buffer = new byte[1024];
int len;
while ((len = fileStream.read(buffer)) > 0) {
out.write(buffer, 0, len);
}
}
}

return true;
} catch (Exception e) {
api.logging().logToError("Failed to init templates for " + group + ": " + e.getMessage());
success = false;
}
} catch (Exception e) {
api
.logging()
.logToError("Failed to copy rules file: " + e.getMessage());
}

return false;
if (!success) {
api.logging().logToError("Template init failed");
}
return success;
}
}
66 changes: 42 additions & 24 deletions src/HaENet/sources/src/main/java/hae/utils/rule/RuleProcessor.java
Original file line number Diff line number Diff line change
Expand Up @@ -5,22 +5,16 @@
import hae.cache.DataCache;
import hae.repository.RuleRepository;
import hae.utils.ConfigLoader;
import hae.utils.rule.model.Group;
import hae.utils.rule.model.RuleDefinition;
import org.yaml.snakeyaml.DumperOptions;
import org.yaml.snakeyaml.Yaml;
import org.yaml.snakeyaml.representer.Representer;

import java.io.File;
import java.io.OutputStreamWriter;
import java.io.Writer;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
import java.util.*;

public class RuleProcessor {
private final MontoyaApi api;
Expand All @@ -36,30 +30,54 @@ public RuleProcessor(MontoyaApi api, ConfigLoader configLoader, RuleRepository r
public void rulesFormatAndSave() {
DataCache.clear();

String templatesPath = configLoader.getTemplatesPath();
File templatesDir = new File(templatesPath);
if (!templatesDir.exists()) {
templatesDir.mkdirs();
}
saveAsTemplates(templatesDir);
}

private void saveAsTemplates(File templatesDir) {
DumperOptions dop = new DumperOptions();
dop.setDefaultFlowStyle(DumperOptions.FlowStyle.BLOCK);
Representer representer = new Representer(dop);
Yaml yaml = new Yaml(representer, dop);

List<Group> ruleGroupList = new ArrayList<>();

ruleRepository.getAll().forEach((k, v) ->
ruleGroupList.add(new Group(k, v))
);

List<Map<String, Object>> outputGroupsMap = ruleGroupList.stream()
.map(Group::getFields)
.collect(Collectors.toList());

Map<String, Object> outputMap = new LinkedHashMap<>();
outputMap.put("rules", outputGroupsMap);
Set<String> writtenFiles = new java.util.HashSet<>();

ruleRepository.getAll().forEach((groupName, rules) -> {
String dirName = ConfigLoader.groupNameToDir(groupName);
File groupDir = new File(templatesDir, dirName);
if (!groupDir.exists()) {
groupDir.mkdirs();
}

for (RuleDefinition rule : rules) {
Map<String, Object> tmpl = rule.toTemplateYaml(groupName);
String fileName = slugify(rule.getName()) + ".yaml";
File targetFile = new File(groupDir, fileName);
writtenFiles.add(targetFile.getAbsolutePath());

try (Writer writer = new java.io.OutputStreamWriter(
Files.newOutputStream(targetFile.toPath()), StandardCharsets.UTF_8)) {
yaml.dump(tmpl, writer);
} catch (Exception e) {
api.logging().logToError("Failed to save template " + fileName + ": " + e.getMessage());
}
}
});
}

File rulesFile = new File(configLoader.getRulesFilePath());
try (Writer writer = new OutputStreamWriter(Files.newOutputStream(rulesFile.toPath()), StandardCharsets.UTF_8)) {
yaml.dump(outputMap, writer);
} catch (Exception e) {
api.logging().logToError("Failed to save rules file: " + e.getMessage());
private static String slugify(String s) {
if (s == null) return "unknown";
StringBuilder sb = new StringBuilder(s.length());
for (char c : s.toCharArray()) {
if ((c >= 'a' && c <= 'z') || (c >= '0' && c <= '9') || c == '-') sb.append(c);
else if (c >= 'A' && c <= 'Z') sb.append((char) (c + 32));
else sb.append('-');
}
return sb.toString();
}

public void changeRule(RuleDefinition rule, int select, String type) {
Expand Down
Loading