Merge pull request #9 from s1kSec/1.1

增加了进度展示以及预测剩余时间

ollamaMain.py

@@ -1,11 +1,13 @@
 import os
 import re
 import json
+import time
 import requests
 from PyQt5 import QtCore, QtGui, QtWidgets
 from PyQt5.QtCore import QThread, pyqtSignal
 from config import OLLAMA_API_URL, OLLAMA_MODEL  # 用户自定义配置
 
+
 class CyberTextEdit(QtWidgets.QTextEdit):
     def __init__(self, parent=None):
         super().__init__(parent)
@@ -21,9 +23,30 @@ class CyberTextEdit(QtWidgets.QTextEdit):
                 }
             """)
 
+
+def seconds_utils(seconds):
+    # 定义各个时间单位的秒数
+    units = {
+        '年': 365 * 86400,
+        '月': 30 * 86400,
+        '天': 86400,
+        '小时': 3600,
+        '分钟': 60,
+        '秒': 1
+    }
+    time_str = []
+    for unit, unit_seconds in units.items():
+        count = seconds // unit_seconds
+        if count > 0:
+            time_str.append(f"{count}{unit}")
+        seconds %= unit_seconds
+    return "".join(time_str)
+
+
 class HackerWorker(QThread):
     analysis_complete = pyqtSignal(str)
     progress_update = pyqtSignal(str)
+    button_text_update = pyqtSignal(str)
 
     def __init__(self, files_content):
         super().__init__()
@@ -31,24 +54,26 @@ class HackerWorker(QThread):
 
     def run(self):
         full_report = []
-        for filepath, content in self.files_content.items():
-            self.progress_update.emit(f"🔍 Analyzing {os.path.basename(filepath)}...")
+        start_time = time.time()
+        for index, (filepath, content) in enumerate(self.files_content.items(), start=1):
+            self.progress_update.emit(
+                f"🔍 Analyzing {os.path.basename(filepath)} ({index}/{len(self.files_content)})...")
 
             prompt = f"""【强制指令】你是一个专业的安全审计AI，请按以下要求分析代码：
 
-1. 漏洞分析流程：
+    1. 漏洞分析流程：
        1.1 识别潜在风险点（SQL操作、文件操作、用户输入点、文件上传漏洞、CSRF、SSRF、XSS、RCE、OWASP top10等漏洞）
        1.2 验证漏洞可利用性
        1.3 按CVSS评分标准评估风险等级
 
-2. 输出规则：
+    2. 输出规则：
        - 仅输出确认存在的高危/中危漏洞
        - 使用严格格式：[风险等级] 类型 - 位置:行号 - 50字内描述
        - 禁止解释漏洞原理
        - 禁止给出修复建议
        - 如果有可能，给出POC（HTTP请求数据包）
 
-3. 输出示例（除此外不要有任何输出）：
+    3. 输出示例（除此外不要有任何输出）：
        [高危] SQL注入 - user_login.php:32 - 未过滤的$_GET参数直接拼接SQL查询
        [POC]POST /login.php HTTP/1.1
        Host: example.com
@@ -58,8 +83,8 @@ class HackerWorker(QThread):
        Host: example.com
        Content-Type: application/x-www-form-urlencoded
 
-4. 当前代码（仅限分析）：
-{content[:3000]}"""
+    4. 当前代码（仅限分析）：
+    {content[:3000]}"""
             try:
                 response = requests.post(
                     f"{OLLAMA_HOST}/api/generate",
@@ -71,43 +96,54 @@ class HackerWorker(QThread):
                 )
                 result = json.loads(response.text)["response"]
                 result = re.sub(r'<think>.*?</think>', '', result, flags=re.DOTALL)
-                full_report.append(f"📄 文件：{filepath}\n{result}\n{'━'*50}")
+                full_report.append(f"📄 文件：{filepath}\n{result}\n{'━' * 50}")
+                # 预测剩余时间：
+                pass_time = int((time.time() - start_time) / index * (len(self.files_content) - index))
+                self.button_text_update.emit(f"⌛预计剩余{seconds_utils(pass_time)}")
             except Exception as e:
                 full_report.append(f"❌ 错误：处理文件 {filepath} 时发生错误\n{str(e)}")
+                # 预测剩余时间：
+                pass_time = int((time.time() - start_time) / index * (len(self.files_content) - index))
+                self.button_text_update.emit(f"⌛预计剩余{seconds_utils(pass_time)}")
 
         self.analysis_complete.emit("\n".join(full_report))
+        self.button_text_update.emit("🚨 启动扫描协议")
+
 
 class WebshellWorker(QThread):
     detection_complete = pyqtSignal(str)
     progress_update = pyqtSignal(str)
+    button_text_update = pyqtSignal(str)
 
     def __init__(self, files_content):
         super().__init__()
         self.files_content = files_content
 
     def run(self):
+        start_time = time.time()
         detection_results = []
-        for filepath, content in self.files_content.items():
-            self.progress_update.emit(f"🕵️ 扫描 {os.path.basename(filepath)}...")
+        for index, (filepath, content) in enumerate(self.files_content.items(), start=1):
+            self.progress_update.emit(
+                f"🕵️ 扫描 {os.path.basename(filepath)} ({index}/{len(self.files_content)})...")
 
             prompt = f"""【Webshell检测指令】请严格按以下步骤分析代码：
 
-1. 检测要求：         
+    1. 检测要求：         
         请分析以下文件内容是否为WebShell或内存马。要求：
         1. 检查PHP/JSP/ASP等WebShell特征（如加密函数、执行系统命令、文件操作）
         2. 识别内存马特征（如无文件落地、进程注入、异常网络连接）
         3. 分析代码中的可疑功能（如命令执行、文件上传、信息收集）
         4. 检查混淆编码、加密手段等规避技术
 
-2. 判断规则：
+    2. 判断规则：
        - 仅当确认恶意性时报告
        - 输出格式：🔴 [高危] Webshell - 文件名:行号 - 检测到[特征1+特征2+...]
 
-3. 输出示例（严格按照此格式输出，不要有任何的补充，如果未检测到危险，则不输出，除此之外，不要有任何输出）：
+    3. 输出示例（严格按照此格式输出，不要有任何的补充，如果未检测到危险，则不输出，除此之外，不要有任何输出）：
        🔴 [高危] Webshell - malicious.php:8 - 检测到[system执行+base64解码+错误抑制]
 
-4. 待分析代码：
-{content[:3000]}"""
+    4. 待分析代码：
+    {content[:3000]}"""
 
             try:
                 response = requests.post(
@@ -120,11 +156,19 @@ class WebshellWorker(QThread):
                 )
                 result = json.loads(response.text)["response"]
                 result = re.sub(r'<think>.*?</think>', '', result, flags=re.DOTALL)
-                detection_results.append(f"📁 {filepath}\n{result}\n{'━'*50}")
+                detection_results.append(f"📁 {filepath}\n{result}\n{'━' * 50}")
+                # 预测剩余时间：
+                pass_time = int((time.time() - start_time) / index * (len(self.files_content) - index))
+                self.button_text_update.emit(f"⌛预计剩余{seconds_utils(pass_time)}")
             except Exception as e:
                 detection_results.append(f"❌ 错误：{filepath}\n{str(e)}")
+                # 预测剩余时间：
+                pass_time = int((time.time() - start_time) / index * (len(self.files_content) - index))
+                self.button_text_update.emit(f"⌛预计剩余{seconds_utils(pass_time)}")
 
         self.detection_complete.emit("\n".join(detection_results))
+        self.button_text_update.emit("🚨 启动扫描协议")
+
 
 class CyberScanner(QtWidgets.QMainWindow):
     def __init__(self):
@@ -286,18 +330,21 @@ class CyberScanner(QtWidgets.QMainWindow):
             worker = HackerWorker(self.files_content)
             init_msg = "🚀 启动深度代码分析协议..."
             complete_signal = worker.analysis_complete
+            worker.button_text_update.connect(self.update_button_text)
         else:
             worker = WebshellWorker(self.files_content)
             init_msg = "🕵️ 启动Webshell检测协议..."
             complete_signal = worker.detection_complete
+            worker.button_text_update.connect(self.update_button_text)
 
         self.scan_thread = worker
         self.scan_thread.progress_update.connect(self.update_status)
         complete_signal.connect(self.show_results)
+        self.scan_thread.button_text_update.connect(self.update_button_text)
         self.scan_thread.start()
 
         self.btn_scan.setEnabled(False)
-        self.result_display.setText(f"{init_msg}\n" + "▮"*50 + "\n")
+        self.result_display.setText(f"{init_msg}\n" + "▮" * 50 + "\n")
 
     def scan_code_files(self, directory):
         allowed_ext = ['.php', '.jsp', '.asp', '.js', '.html', '.py', '.java']
@@ -338,9 +385,14 @@ class CyberScanner(QtWidgets.QMainWindow):
 
         self.result_display.append(report)
         self.status_bar.showMessage("✅ 扫描完成")
+        self.btn_scan.setEnabled(True)
+
+    def update_button_text(self, new_text):
+        self.btn_scan.setText(new_text)
+
 
 if __name__ == "__main__":
-# 保持源文本的核心内容不变
+    # 保持源文本的核心内容不变
 
     OLLAMA_HOST = OLLAMA_API_URL.split('/api')[0]