Felix-sec/VulnDataExtr
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Optimized UI layout

Browse Source
This commit is contained in:
Felix
2025-01-22 17:46:01 +08:00
parent 4e93437599
commit 2b57df40d8
6 changed files with 1949 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
.gitignore vendored
View File

@@ -3,6 +3,8 @@
/dist/
/venv/
/.idea/
/__pycache__/
/.idea/
# 忽略日志文件
*.log
@@ -11,7 +13,9 @@
*说明*.txt
# 忽略源文件
*.html
index.html
# 忽略测试输出的文件
*.json
data.json
index.json
*.xlsx

1582
exact_keywords.json Normal file
View File

File diff suppressed because it is too large Load Diff

9
fuzzy_keywords.json Normal file
View File

@@ -0,0 +1,9 @@
{
"Java": "应用",
"MySQL": "数据库",
"Samba": "应用",
"SMB": "应用",
"Python": "应用",
"Redis": "应用",
"Oracle": "数据库"
}

54
gui.py
View File

@@ -39,29 +39,22 @@ class JsonExtractorGUI:
self.save_btn = tk.Button(self.output_frame, text="浏览", command=self.browse_output)
self.save_btn.pack(side="left", padx=5)
# 将提取JSON按钮移到输出文件框中
self.extract_btn = tk.Button(self.output_frame, text="提取JSON",
command=self.extract_json)
self.extract_btn.pack(side="left", padx=5)
# 添加关键词管理器
self.keyword_manager = KeywordManager()
# 添加按钮框架
self.button_frame = tk.Frame(root)
self.button_frame.pack(pady=5)
# 提取按钮移到按钮框架
self.extract_btn = tk.Button(self.button_frame, text="提取JSON", command=self.extract_json)
self.extract_btn.pack(side="left", padx=5)
# 添加漏洞分类按钮
self.vuln_btn = tk.Button(self.button_frame, text="漏洞类型分类", command=self.export_vuln_types)
self.vuln_btn.pack(side="left", padx=5)
# 添加关键词管理按钮
self.keyword_btn = tk.Button(self.button_frame, text="关键词管理", command=self.show_keyword_dialog)
self.keyword_btn.pack(side="left", padx=5)
# 创建漏洞分类框架
vuln_frame = tk.LabelFrame(root, text="漏洞类型分类", padx=5, pady=5)
vuln_frame.pack(fill="x", padx=10, pady=5)
# 添加匹配模式选择
self.match_mode = tk.StringVar(value="both")
match_frame = tk.LabelFrame(root, text="匹配模式", padx=5, pady=5)
match_frame.pack(fill="x", padx=10, pady=5)
match_label = tk.Label(vuln_frame, text="匹配模式:")
match_label.pack(side="left", padx=5)
modes = [
("精准匹配", "exact"),
@@ -70,9 +63,19 @@ class JsonExtractorGUI:
]
for text, mode in modes:
tk.Radiobutton(match_frame, text=text, variable=self.match_mode,
tk.Radiobutton(vuln_frame, text=text, variable=self.match_mode,
value=mode).pack(side="left", padx=5)
# 添加关键词管理按钮
self.keyword_btn = tk.Button(vuln_frame, text="关键词管理",
command=self.show_keyword_dialog)
self.keyword_btn.pack(side="left", padx=15)
# 添加导出按钮
self.vuln_btn = tk.Button(vuln_frame, text="导出分类",
command=self.export_vuln_types)
self.vuln_btn.pack(side="left", padx=5)
# 日志框
self.log_frame = tk.LabelFrame(root, text="运行日志", padx=10, pady=5)
self.log_frame.pack(fill="both", expand=True, padx=10, pady=5)
@@ -189,6 +192,20 @@ class JsonExtractorGUI:
def export_vuln_types(self):
try:
# 先让用户选择保存位置
output_excel = filedialog.asksaveasfilename(
title="保存漏洞分类Excel",
defaultextension=".xlsx",
initialfile="漏洞类型分类.xlsx",
filetypes=[
("Excel文件", "*.xlsx"),
("所有文件", "*.*")
]
)
if not output_excel: # 用户取消选择
return
# 读取JSON文件
with open(self.output_path.get(), 'r', encoding='utf-8') as f:
data = json.load(f)
@@ -226,7 +243,6 @@ class JsonExtractorGUI:
# 创建DataFrame并导出到Excel
df = pd.DataFrame(excel_data)
output_excel = Path(self.output_path.get()).with_suffix('.xlsx')
df.to_excel(output_excel, index=False)
self.log_message(f"漏洞类型分类已导出到: {output_excel}", "SUCCESS")

210
keyword_dialog.py Normal file
View File

@@ -0,0 +1,210 @@
import tkinter as tk
from tkinter import ttk, messagebox, filedialog
import pandas as pd
class KeywordDialog:
def __init__(self, parent, keyword_manager):
self.dialog = tk.Toplevel(parent)
self.dialog.title("关键词管理")
self.dialog.geometry("800x600")
self.keyword_manager = keyword_manager
# 创建选项卡
self.notebook = ttk.Notebook(self.dialog)
self.exact_frame = ttk.Frame(self.notebook)
self.fuzzy_frame = ttk.Frame(self.notebook)
self.notebook.add(self.exact_frame, text="精准匹配")
self.notebook.add(self.fuzzy_frame, text="模糊匹配")
self.notebook.pack(expand=True, fill="both")
# 创建精准匹配界面
self.create_keyword_frame(self.exact_frame, False)
# 创建模糊匹配界面
self.create_keyword_frame(self.fuzzy_frame, True)
# 在每个标签页添加批量导入按钮
self.create_import_buttons(self.exact_frame, False)
self.create_import_buttons(self.fuzzy_frame, True)
def create_keyword_frame(self, frame, is_fuzzy):
# 输入区域
input_frame = ttk.LabelFrame(frame, text="添加关键词", padding=5)
input_frame.pack(fill="x", padx=5, pady=5)
ttk.Label(input_frame, text="关键词:").grid(row=0, column=0, padx=5)
keyword_entry = ttk.Entry(input_frame, width=30)
keyword_entry.grid(row=0, column=1, padx=5)
ttk.Label(input_frame, text="类型:").grid(row=0, column=2, padx=5)
type_entry = ttk.Entry(input_frame, width=20)
type_entry.grid(row=0, column=3, padx=5)
def add_keyword():
keyword = keyword_entry.get().strip()
type_name = type_entry.get().strip()
if keyword and type_name:
self.keyword_manager.add_keyword(keyword, type_name, is_fuzzy)
keyword_entry.delete(0, tk.END)
type_entry.delete(0, tk.END)
refresh_list()
else:
messagebox.showwarning("警告", "关键词和类型不能为空！")
ttk.Button(input_frame, text="添加", command=add_keyword).grid(row=0, column=4, padx=5)
# 列表区域
list_frame = ttk.Frame(frame)
list_frame.pack(fill="both", expand=True, padx=5, pady=5)
columns = ("关键词", "类型")
tree = ttk.Treeview(list_frame, columns=columns, show="headings")
for col in columns:
tree.heading(col, text=col)
tree.column(col, width=100)
scrollbar = ttk.Scrollbar(list_frame, orient="vertical", command=tree.yview)
tree.configure(yscrollcommand=scrollbar.set)
tree.pack(side="left", fill="both", expand=True)
scrollbar.pack(side="right", fill="y")
def refresh_list():
tree.delete(*tree.get_children())
keywords = self.keyword_manager.fuzzy_keywords if is_fuzzy else self.keyword_manager.exact_keywords
for keyword, type_name in keywords.items():
tree.insert("", "end", values=(keyword, type_name))
def remove_selected():
selected = tree.selection()
if selected:
item = tree.item(selected[0])
keyword = item['values'][0]
self.keyword_manager.remove_keyword(keyword, is_fuzzy)
refresh_list()
ttk.Button(frame, text="删除选中", command=remove_selected).pack(pady=5)
refresh_list()
def create_import_buttons(self, frame, is_fuzzy):
# 创建按钮框架并居中
import_frame = ttk.Frame(frame)
import_frame.pack(fill="x", padx=5, pady=10)
# 创建一个子框架来容纳按钮，并使其居中
button_frame = ttk.Frame(import_frame)
button_frame.pack(anchor="center")
# 导入按钮 - 设置宽度和高度
ttk.Button(button_frame, text="导入Excel", width=15, padding=(5, 8),
command=lambda: self.import_from_excel(is_fuzzy)).pack(side="left", padx=10)
# 导出按钮 - 设置宽度和高度
ttk.Button(button_frame, text="导出Excel", width=15, padding=(5, 8),
command=lambda: self.export_to_excel(is_fuzzy)).pack(side="left", padx=10)
def import_from_excel(self, is_fuzzy):
filename = filedialog.askopenfilename(
title="选择Excel文件",
filetypes=[
("Excel文件", "*.xlsx *.xls"),
("所有文件", "*.*")
]
)
if not filename:
return
try:
df = pd.read_excel(filename)
# 检查必要的列
required_columns = ['关键词', '类型']
if not all(col in df.columns for col in required_columns):
messagebox.showerror("错误", "Excel文件必须包含'关键词'和'类型'列！")
return
# 准备导入数据
keywords_data = list(zip(df['关键词'], df['类型']))
# 创建导入选项对话框
option_dialog = tk.Toplevel(self.dialog)
option_dialog.title("导入选项")
option_dialog.geometry("300x150")
overwrite_var = tk.BooleanVar(value=True)
ttk.Checkbutton(option_dialog, text="覆盖已存在的关键词",
variable=overwrite_var).pack(pady=10)
def do_import():
success_count, skip_count, errors = self.keyword_manager.batch_import(
keywords_data, is_fuzzy, overwrite_var.get()
)
result_msg = f"成功导入: {success_count}\n跳过: {skip_count}"
if errors:
result_msg += f"\n\n错误信息:\n" + "\n".join(errors)
messagebox.showinfo("导入结果", result_msg)
option_dialog.destroy()
# 刷新显示
for frame in [self.exact_frame, self.fuzzy_frame]:
for child in frame.winfo_children():
if isinstance(child, ttk.Frame):
for widget in child.winfo_children():
if isinstance(widget, ttk.Treeview):
self.refresh_list(widget, frame == self.fuzzy_frame)
ttk.Button(option_dialog, text="开始导入",
command=do_import).pack(pady=10)
ttk.Button(option_dialog, text="取消",
command=option_dialog.destroy).pack(pady=5)
# 使对话框模态
option_dialog.transient(self.dialog)
option_dialog.grab_set()
self.dialog.wait_window(option_dialog)
except Exception as e:
messagebox.showerror("错误", f"导入过程中出错：\n{str(e)}")
def export_to_excel(self, is_fuzzy):
try:
# 获取要导出的关键词库
keywords = self.keyword_manager.fuzzy_keywords if is_fuzzy else self.keyword_manager.exact_keywords
# 检查关键词库是否为空
if not keywords:
messagebox.showwarning("警告", "关键词库为空，无法导出！")
return
# 创建DataFrame
df = pd.DataFrame([
{'关键词': keyword, '类型': type_name}
for keyword, type_name in keywords.items()
])
# 让用户选择保存位置
filename = filedialog.asksaveasfilename(
title="保存Excel文件",
defaultextension=".xlsx",
filetypes=[
("Excel文件", "*.xlsx"),
("所有文件", "*.*")
]
)
if filename:
# 导出到Excel
df.to_excel(filename, index=False)
messagebox.showinfo("成功", f"关键词库已导出到：\n{filename}")
except Exception as e:
messagebox.showerror("错误", f"导出过程中出错：\n{str(e)}")
def refresh_list(self, tree, is_fuzzy):
"""刷新指定树形视图的显示"""
tree.delete(*tree.get_children())
keywords = self.keyword_manager.fuzzy_keywords if is_fuzzy else self.keyword_manager.exact_keywords
for keyword, type_name in keywords.items():
tree.insert("", "end", values=(keyword, type_name))

107
keyword_manager.py Normal file
View File

@@ -0,0 +1,107 @@
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