SearchDomain-Python/gui.py

#!/usr/bin/env python
# -*- encoding: utf-8 -*-
"""
@Contact :   liuyuqi.gov@msn.cn
@Time    :   2024/12/19
@License :   Copyright © 2017-2022 liuyuqi. All Rights Reserved.
@Desc    :   GUI界面入口
"""

import tkinter as tk
from tkinter import ttk, scrolledtext, filedialog, messagebox
import threading
import sys
import os
from collections import OrderedDict
from searchdomain import SearchDomain, GenerateDomain, GenerateEnDomain
from searchdomain.utils.frozen_dir import get_app_path


class SearchDomainGUI:
    def __init__(self, root):
        self.root = root
        self.root.title("域名批量检索工具")
        self.root.geometry("800x700")
        self.root.resizable(True, True)
        
        # 设置应用路径
        self.app_path = get_app_path()
        
        # 创建主界面
        self.create_widgets()
        
        # 运行状态
        self.is_running = False
        
    def create_widgets(self):
        # 创建Notebook（标签页）
        notebook = ttk.Notebook(self.root)
        notebook.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)
        
        # 生成域名标签页
        self.generate_frame = ttk.Frame(notebook)
        notebook.add(self.generate_frame, text="生成域名")
        self.create_generate_tab()
        
        # 搜索域名标签页
        self.search_frame = ttk.Frame(notebook)
        notebook.add(self.search_frame, text="搜索域名")
        self.create_search_tab()
        
        # 日志输出区域（共用）
        log_frame = ttk.LabelFrame(self.root, text="运行日志")
        log_frame.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)
        
        self.log_text = scrolledtext.ScrolledText(log_frame, height=10, wrap=tk.WORD)
        self.log_text.pack(fill=tk.BOTH, expand=True, padx=5, pady=5)
        
        # 状态栏
        self.status_var = tk.StringVar(value="就绪")
        status_bar = ttk.Label(self.root, textvariable=self.status_var, relief=tk.SUNKEN)
        status_bar.pack(side=tk.BOTTOM, fill=tk.X)
        
    def create_generate_tab(self):
        # 语言选择
        lang_frame = ttk.LabelFrame(self.generate_frame, text="语言设置")
        lang_frame.pack(fill=tk.X, padx=10, pady=5)
        
        self.lang_var = tk.StringVar(value="en")
        ttk.Radiobutton(lang_frame, text="中文", variable=self.lang_var, value="zh").pack(side=tk.LEFT, padx=10, pady=5)
        ttk.Radiobutton(lang_frame, text="英文", variable=self.lang_var, value="en").pack(side=tk.LEFT, padx=10, pady=5)
        
        # 关键词输入
        keyword_frame = ttk.LabelFrame(self.generate_frame, text="关键词（用逗号分隔）")
        keyword_frame.pack(fill=tk.X, padx=10, pady=5)
        
        self.keyword_entry = ttk.Entry(keyword_frame, width=50)
        self.keyword_entry.pack(side=tk.LEFT, padx=5, pady=5, fill=tk.X, expand=True)
        self.keyword_entry.insert(0, "gpt,go,ai")
        
        # 域名后缀
        domain_frame = ttk.LabelFrame(self.generate_frame, text="域名后缀（用逗号分隔）")
        domain_frame.pack(fill=tk.X, padx=10, pady=5)
        
        self.domain_entry = ttk.Entry(domain_frame, width=50)
        self.domain_entry.pack(side=tk.LEFT, padx=5, pady=5, fill=tk.X, expand=True)
        self.domain_entry.insert(0, "com,cn,io")
        
        # 位置选择
        position_frame = ttk.LabelFrame(self.generate_frame, text="关键词位置")
        position_frame.pack(fill=tk.X, padx=10, pady=5)
        
        self.position_var = tk.StringVar(value="prefix")
        ttk.Radiobutton(position_frame, text="前缀", variable=self.position_var, value="prefix").pack(side=tk.LEFT, padx=10, pady=5)
        ttk.Radiobutton(position_frame, text="后缀", variable=self.position_var, value="suffix").pack(side=tk.LEFT, padx=10, pady=5)
        
        # 输出文件
        output_frame = ttk.LabelFrame(self.generate_frame, text="输出文件")
        output_frame.pack(fill=tk.X, padx=10, pady=5)
        
        self.generate_output_entry = ttk.Entry(output_frame, width=40)
        self.generate_output_entry.pack(side=tk.LEFT, padx=5, pady=5, fill=tk.X, expand=True)
        self.generate_output_entry.insert(0, "domain.txt")
        
        ttk.Button(output_frame, text="浏览", command=self.browse_generate_output).pack(side=tk.LEFT, padx=5, pady=5)
        
        # 生成按钮
        button_frame = ttk.Frame(self.generate_frame)
        button_frame.pack(fill=tk.X, padx=10, pady=10)
        
        self.generate_button = ttk.Button(button_frame, text="开始生成", command=self.start_generate)
        self.generate_button.pack(side=tk.LEFT, padx=5)
        
        ttk.Button(button_frame, text="清空日志", command=self.clear_log).pack(side=tk.LEFT, padx=5)
        
    def create_search_tab(self):
        # 输入文件
        input_frame = ttk.LabelFrame(self.search_frame, text="输入文件（域名列表）")
        input_frame.pack(fill=tk.X, padx=10, pady=5)
        
        self.input_entry = ttk.Entry(input_frame, width=40)
        self.input_entry.pack(side=tk.LEFT, padx=5, pady=5, fill=tk.X, expand=True)
        self.input_entry.insert(0, "domain.txt")
        
        ttk.Button(input_frame, text="浏览", command=self.browse_input_file).pack(side=tk.LEFT, padx=5, pady=5)
        
        # 输出文件
        output_frame = ttk.LabelFrame(self.search_frame, text="输出文件（结果）")
        output_frame.pack(fill=tk.X, padx=10, pady=5)
        
        self.search_output_entry = ttk.Entry(output_frame, width=40)
        self.search_output_entry.pack(side=tk.LEFT, padx=5, pady=5, fill=tk.X, expand=True)
        self.search_output_entry.insert(0, "result.txt")
        
        ttk.Button(output_frame, text="浏览", command=self.browse_search_output).pack(side=tk.LEFT, padx=5, pady=5)
        
        # 导出选项
        export_frame = ttk.LabelFrame(self.search_frame, text="导出选项")
        export_frame.pack(fill=tk.X, padx=10, pady=5)
        
        self.export_all_var = tk.BooleanVar(value=False)
        ttk.Checkbutton(export_frame, text="导出所有域名（包括不可用的）", variable=self.export_all_var).pack(side=tk.LEFT, padx=10, pady=5)
        
        # 搜索按钮
        button_frame = ttk.Frame(self.search_frame)
        button_frame.pack(fill=tk.X, padx=10, pady=10)
        
        self.search_button = ttk.Button(button_frame, text="开始搜索", command=self.start_search)
        self.search_button.pack(side=tk.LEFT, padx=5)
        
        self.cancel_search_button = ttk.Button(button_frame, text="终止", command=self.cancel_search, state=tk.DISABLED)
        self.cancel_search_button.pack(side=tk.LEFT, padx=5)
        
        ttk.Button(button_frame, text="清空日志", command=self.clear_log).pack(side=tk.LEFT, padx=5)
        
    def browse_input_file(self):
        filename = filedialog.askopenfilename(
            title="选择输入文件",
            filetypes=[("文本文件", "*.txt"), ("所有文件", "*.*")]
        )
        if filename:
            self.input_entry.delete(0, tk.END)
            self.input_entry.insert(0, filename)
            
    def browse_search_output(self):
        filename = filedialog.asksaveasfilename(
            title="选择输出文件",
            defaultextension=".txt",
            filetypes=[("文本文件", "*.txt"), ("所有文件", "*.*")]
        )
        if filename:
            self.search_output_entry.delete(0, tk.END)
            self.search_output_entry.insert(0, filename)
            
    def browse_generate_output(self):
        filename = filedialog.asksaveasfilename(
            title="选择输出文件",
            defaultextension=".txt",
            filetypes=[("文本文件", "*.txt"), ("所有文件", "*.*")]
        )
        if filename:
            self.generate_output_entry.delete(0, tk.END)
            self.generate_output_entry.insert(0, filename)
            
    def log(self, message):
        """添加日志"""
        self.log_text.insert(tk.END, message + "\n")
        self.log_text.see(tk.END)
        self.root.update_idletasks()
        
    def clear_log(self):
        """清空日志"""
        self.log_text.delete(1.0, tk.END)
        
    def start_generate(self):
        """开始生成域名"""
        if self.is_running:
            messagebox.showwarning("警告", "任务正在运行中，请等待完成")
            return
            
        # 获取参数
        lang = self.lang_var.get()
        keyword = self.keyword_entry.get().strip()
        domain = self.domain_entry.get().strip()
        position = self.position_var.get()
        output = self.generate_output_entry.get().strip()
        
        if not keyword:
            messagebox.showerror("错误", "请输入关键词")
            return
            
        if not domain:
            messagebox.showerror("错误", "请输入域名后缀")
            return
            
        if not output:
            messagebox.showerror("错误", "请输入输出文件名")
            return
            
        # 准备参数
        params = OrderedDict({
            "lang": lang,
            "keyword": keyword,
            "domain": domain,
            "position": position,
            "app_path": self.app_path
        })
        
        # 更新输出文件路径为绝对路径
        if not os.path.isabs(output):
            output = os.path.join(self.app_path, output)
        else:
            params["app_path"] = os.path.dirname(output)
            
        # 在后台线程中运行
        self.is_running = True
        self.generate_button.config(state=tk.DISABLED)
        self.status_var.set("正在生成域名...")
        
        thread = threading.Thread(target=self._generate_thread, args=(params, lang))
        thread.daemon = True
        thread.start()
        
    def _generate_thread(self, params, lang):
        """生成域名的后台线程"""
        try:
            self.log(f"开始生成域名...")
            self.log(f"语言: {lang}, 关键词: {params['keyword']}, 域名后缀: {params['domain']}, 位置: {params['position']}")
            
            if lang == "en":
                generateDomain = GenerateEnDomain(params=params)
            else:
                generateDomain = GenerateDomain(params=params)
                
            generateDomain.run()
            
            output_file = params.get("output", "domain.txt")
            self.log("域名生成完成！")
            self.log(f"结果已保存到: {output_file}")
            messagebox.showinfo("成功", f"域名生成完成！\n结果已保存到: {output_file}")
            
        except Exception as e:
            error_msg = f"生成域名时出错: {str(e)}"
            self.log(error_msg)
            messagebox.showerror("错误", error_msg)
        finally:
            self.is_running = False
            self.root.after(0, lambda: self.generate_button.config(state=tk.NORMAL))
            self.root.after(0, lambda: self.status_var.set("就绪"))
            
    def start_search(self):
        """开始搜索域名"""
        if self.is_running:
            messagebox.showwarning("警告", "任务正在运行中，请等待完成")
            return
            
        # 获取参数
        input_file = self.input_entry.get().strip()
        output_file = self.search_output_entry.get().strip()
        export_all = self.export_all_var.get()
        
        if not input_file:
            messagebox.showerror("错误", "请输入输入文件名")
            return
            
        if not output_file:
            messagebox.showerror("错误", "请输入输出文件名")
            return
            
        # 检查输入文件是否存在
        if not os.path.isabs(input_file):
            input_file = os.path.join(self.app_path, input_file)
            
        if not os.path.exists(input_file):
            messagebox.showerror("错误", f"输入文件不存在: {input_file}")
            return
            
        # 准备输出文件路径
        if not os.path.isabs(output_file):
            output_file = os.path.join(self.app_path, output_file)
        
        # 确保输出目录存在
        output_dir = os.path.dirname(output_file)
        if output_dir and not os.path.exists(output_dir):
            try:
                os.makedirs(output_dir, exist_ok=True)
            except Exception as e:
                messagebox.showerror("错误", f"无法创建输出目录: {output_dir}\n{str(e)}")
                return
            
        # 准备参数
        # app_path 应该设置为输出文件的目录，因为 saveRes 使用 app_path + output
        params = OrderedDict({
            "input": input_file,  # 使用完整路径
            "output": os.path.basename(output_file),  # 只使用文件名
            "app_path": os.path.dirname(output_file)  # 输出文件的目录
        })
        
        # 在后台线程中运行
        self.is_running = True
        self.search_button.config(state=tk.DISABLED)
        self.cancel_search_button.config(state=tk.NORMAL)
        self.status_var.set("正在搜索域名...")
        
        thread = threading.Thread(target=self._search_thread, args=(params, export_all))
        thread.daemon = True
        thread.start()
    
    def cancel_search(self):
        """取消搜索任务"""
        if self.search_domain_instance and self.is_running:
            self.search_domain_instance.cancel()
            self.log("正在终止搜索任务...")
            self.status_var.set("正在终止...")
        
    def _search_thread(self, params, export_all):
        """搜索域名的后台线程"""
        try:
            self.log(f"开始搜索域名...")
            self.log(f"输入文件: {params['input']}")
            self.log(f"输出目录: {params['app_path']}")
            self.log(f"输出文件: {params['output']}")
            self.log(f"导出所有: {export_all}")
            
            # 传入日志回调函数，将搜索进度显示在运行日志中
            self.search_domain_instance = SearchDomain(params=params, debug=True, export_all=export_all, log_callback=self.log)
            self.search_domain_instance.run()
            
            # 检查是否被取消
            if self.search_domain_instance._cancelled:
                self.log("搜索任务已取消")
                self.root.after(0, lambda: messagebox.showinfo("提示", "搜索任务已取消"))
            else:
                self.log("域名搜索完成！")
                output_path = os.path.join(params["app_path"], params["output"])
                self.log(f"结果已保存到: {output_path}")
                
                # 验证文件是否真的存在
                if os.path.exists(output_path):
                    file_size = os.path.getsize(output_path)
                    self.log(f"文件大小: {file_size} 字节")
                    self.root.after(0, lambda: messagebox.showinfo("成功", f"域名搜索完成！\n结果已保存到: {output_path}\n文件大小: {file_size} 字节"))
                else:
                    error_msg = f"警告：输出文件不存在: {output_path}"
                    self.log(error_msg)
                    self.root.after(0, lambda: messagebox.showwarning("警告", error_msg))
            
        except Exception as e:
            error_msg = f"搜索域名时出错: {str(e)}"
            self.log(error_msg)
            import traceback
            self.log(traceback.format_exc())
            self.root.after(0, lambda: messagebox.showerror("错误", error_msg))
        finally:
            self.is_running = False
            self.search_domain_instance = None
            self.root.after(0, lambda: self.search_button.config(state=tk.NORMAL))
            self.root.after(0, lambda: self.cancel_search_button.config(state=tk.DISABLED))
            self.root.after(0, lambda: self.status_var.set("就绪"))


def main():
    root = tk.Tk()
    app = SearchDomainGUI(root)
    root.mainloop()


if __name__ == "__main__":
    main()