107 lines
3.8 KiB
Python
107 lines
3.8 KiB
Python
import json
|
|
from pathlib import Path
|
|
import re
|
|
|
|
class KeywordManager:
|
|
def __init__(self):
|
|
self.exact_keywords = {} # 精准匹配关键词库
|
|
self.fuzzy_keywords = {} # 模糊匹配关键词库
|
|
self.load_keywords()
|
|
|
|
def load_keywords(self):
|
|
"""从文件加载关键词库"""
|
|
try:
|
|
if Path('exact_keywords.json').exists():
|
|
with open('exact_keywords.json', 'r', encoding='utf-8') as f:
|
|
self.exact_keywords = json.load(f)
|
|
if Path('fuzzy_keywords.json').exists():
|
|
with open('fuzzy_keywords.json', 'r', encoding='utf-8') as f:
|
|
self.fuzzy_keywords = json.load(f)
|
|
except Exception as e:
|
|
print(f"加载关键词库失败: {e}")
|
|
|
|
def save_keywords(self):
|
|
"""保存关键词库到文件"""
|
|
try:
|
|
with open('exact_keywords.json', 'w', encoding='utf-8') as f:
|
|
json.dump(self.exact_keywords, f, ensure_ascii=False, indent=4)
|
|
with open('fuzzy_keywords.json', 'w', encoding='utf-8') as f:
|
|
json.dump(self.fuzzy_keywords, f, ensure_ascii=False, indent=4)
|
|
except Exception as e:
|
|
print(f"保存关键词库失败: {e}")
|
|
|
|
def add_keyword(self, keyword, type_name, is_fuzzy=False):
|
|
"""添加关键词"""
|
|
if is_fuzzy:
|
|
self.fuzzy_keywords[keyword] = type_name
|
|
else:
|
|
self.exact_keywords[keyword] = type_name
|
|
self.save_keywords()
|
|
|
|
def remove_keyword(self, keyword, is_fuzzy=False):
|
|
"""删除关键词"""
|
|
if is_fuzzy:
|
|
self.fuzzy_keywords.pop(keyword, None)
|
|
else:
|
|
self.exact_keywords.pop(keyword, None)
|
|
self.save_keywords()
|
|
|
|
def get_type(self, vuln_name, match_mode='both'):
|
|
"""
|
|
根据漏洞名称获取类型
|
|
match_mode: 'exact'(仅精准匹配), 'fuzzy'(仅模糊匹配), 'both'(先精准后模糊)
|
|
"""
|
|
if match_mode in ['exact', 'both']:
|
|
# 精准匹配
|
|
if vuln_name in self.exact_keywords:
|
|
return self.exact_keywords[vuln_name]
|
|
|
|
if match_mode in ['fuzzy', 'both']:
|
|
# 模糊匹配
|
|
for keyword, type_name in self.fuzzy_keywords.items():
|
|
if keyword in vuln_name:
|
|
return type_name
|
|
|
|
return "未知" # 修改默认返回值
|
|
|
|
def batch_import(self, keywords_data, is_fuzzy=False, overwrite=True):
|
|
"""
|
|
批量导入关键词
|
|
|
|
Args:
|
|
keywords_data: list of tuples [(keyword, type_name), ...]
|
|
is_fuzzy: 是否为模糊匹配关键词
|
|
overwrite: 是否覆盖已存在的关键词
|
|
|
|
Returns:
|
|
tuple: (成功数量, 跳过数量, 错误信息列表)
|
|
"""
|
|
success_count = 0
|
|
skip_count = 0
|
|
errors = []
|
|
|
|
target_dict = self.fuzzy_keywords if is_fuzzy else self.exact_keywords
|
|
|
|
for keyword, type_name in keywords_data:
|
|
try:
|
|
keyword = str(keyword).strip()
|
|
type_name = str(type_name).strip()
|
|
|
|
if not keyword or not type_name:
|
|
errors.append(f"无效的数据: {keyword} -> {type_name}")
|
|
continue
|
|
|
|
if keyword in target_dict and not overwrite:
|
|
skip_count += 1
|
|
continue
|
|
|
|
target_dict[keyword] = type_name
|
|
success_count += 1
|
|
|
|
except Exception as e:
|
|
errors.append(f"处理 {keyword} 时出错: {str(e)}")
|
|
|
|
if success_count > 0:
|
|
self.save_keywords()
|
|
|
|
return success_count, skip_count, errors |