Optimized UI layout

.gitignore

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

fuzzy_keywords.json

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

gui.py

@@ -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)
+        vuln_frame = tk.LabelFrame(root, text="漏洞类型分类", padx=5, pady=5)
-        self.button_frame.pack(pady=5)
+        vuln_frame.pack(fill="x", padx=10, 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)
         
         # 添加匹配模式选择
         self.match_mode = tk.StringVar(value="both")
-        match_frame = tk.LabelFrame(root, text="匹配模式", padx=5, pady=5)
+        match_label = tk.Label(vuln_frame, text="匹配模式:")
-        match_frame.pack(fill="x", padx=10, pady=5)
+        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")

keyword_dialog.py

@@ -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)) 

keyword_manager.py

@@ -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