liuyuqi-cnb 1 day ago
parent
commit
eabe9ae8bd
2 changed files with 77 additions and 18 deletions
  1. 33 15
      gui.py
  2. 44 3
      searchdomain/searchdomain.py

+ 33 - 15
gui.py

@@ -147,6 +147,9 @@ class SearchDomainGUI:
         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):
@@ -315,11 +318,19 @@ class SearchDomainGUI:
         # 在后台线程中运行
         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):
         """搜索域名的后台线程"""
@@ -331,32 +342,39 @@ class SearchDomainGUI:
             self.log(f"导出所有: {export_all}")
             
             # 传入日志回调函数,将搜索进度显示在运行日志中
-            searchdomain = SearchDomain(params=params, debug=True, export_all=export_all, log_callback=self.log)
-            searchdomain.run()
-            
-            self.log("域名搜索完成!")
-            output_path = os.path.join(params["app_path"], params["output"])
-            self.log(f"结果已保存到: {output_path}")
+            self.search_domain_instance = SearchDomain(params=params, debug=True, export_all=export_all, log_callback=self.log)
+            self.search_domain_instance.run()
             
-            # 验证文件是否真的存在
-            if os.path.exists(output_path):
-                file_size = os.path.getsize(output_path)
-                self.log(f"文件大小: {file_size} 字节")
-                messagebox.showinfo("成功", f"域名搜索完成!\n结果已保存到: {output_path}\n文件大小: {file_size} 字节")
+            # 检查是否被取消
+            if self.search_domain_instance._cancelled:
+                self.log("搜索任务已取消")
+                self.root.after(0, lambda: messagebox.showinfo("提示", "搜索任务已取消"))
             else:
-                error_msg = f"警告:输出文件不存在: {output_path}"
-                self.log(error_msg)
-                messagebox.showwarning("警告", error_msg)
+                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())
-            messagebox.showerror("错误", error_msg)
+            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("就绪"))
 
 

+ 44 - 3
searchdomain/searchdomain.py

@@ -22,6 +22,7 @@ class SearchDomain(object):
         self.output=params["output"]
         self.log_callback = log_callback
         self._output_file_checked = False  # 标记输出文件是否已检查
+        self._cancelled = False  # 取消标志
         
         # 获取当前模块的日志记录器
         self.logger = logging.getLogger(__name__)
@@ -55,11 +56,20 @@ class SearchDomain(object):
         else:
             logging.basicConfig(level=logging.INFO, format='%(message)s', force=True)
 
+    def cancel(self):
+        """取消搜索任务"""
+        self._cancelled = True
+        self.logger.info("搜索任务已取消")
+
     def crawl(self, domain: str, index:int) -> None:
         '''
         检测域名是否可用
         :params domain 域名:
         :return true or false'''
+        # 检查是否已取消
+        if self._cancelled:
+            return
+        
         res = False
         try:
             whi = whois.whois(domain)
@@ -74,6 +84,11 @@ class SearchDomain(object):
             else:
                 res = False
                 self.logger.error(f"Error checking {domain}: {error_str}")
+        
+        # 再次检查是否已取消
+        if self._cancelled:
+            return
+            
         if self.export_all:
             self.saveRes(domain, res)
         else:
@@ -100,6 +115,18 @@ class SearchDomain(object):
 
     def run(self):
         '''begin search domain'''
+        # 先创建输出文件(如果不存在),在保存结果前创建
+        output_path = os.path.join(self.params["app_path"], self.output)
+        if not os.path.exists(output_path):
+            # 确保目录存在
+            dir_path = os.path.dirname(output_path)
+            if dir_path and not os.path.exists(dir_path):
+                os.makedirs(dir_path, exist_ok=True)
+            # 创建文件
+            with open(output_path, 'w', encoding='utf-8') as f:
+                pass
+            self.logger.warning(f"警告:输出文件不存在,已创建: {output_path}")
+        
         # 支持 input 为完整路径或相对路径
         input_path = self.input if os.path.isabs(self.input) else os.path.join(self.params["app_path"], self.input)
         with open(input_path, "r", encoding="utf8", errors="ignore") as file:
@@ -107,14 +134,28 @@ class SearchDomain(object):
                 index = 0
                 futures = []
                 for line in file.readlines():
+                    # 检查是否已取消
+                    if self._cancelled:
+                        break
                     domain = line.strip()
                     if domain:  # 跳过空行
                         index = index + 1
                         future = pool.submit(self.crawl, domain, index)
                         futures.append(future)
-                # 等待所有任务完成
-                for future in futures:
-                    future.result()
+                # 等待所有任务完成,但如果已取消则取消未完成的任务
+                if self._cancelled:
+                    # 取消所有未完成的任务
+                    for future in futures:
+                        future.cancel()
+                    self.logger.info("已取消所有未完成的任务")
+                else:
+                    # 等待所有任务完成
+                    for future in futures:
+                        try:
+                            future.result()
+                        except Exception as e:
+                            if not self._cancelled:
+                                self.logger.error(f"任务执行出错: {str(e)}")
 
 if __name__ == '__main__':
     sd = SearchDomain()