Merge pull request #9 from s1kSec/1.1

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

ollamaMain.py

@@ -1,29 +1,52 @@
 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)
         self.setStyleSheet("""
-            QTextEdit {
-                background-color: #001a1a;
-                color: #00ff00;
-                border: 2px solid #00ffff;
-                border-radius: 5px;
-                padding: 10px;
-                font-family: 'Consolas';
-                font-size: 12pt;
-            }
-        """)
+                QTextEdit {
+                    background-color: #001a1a;
+                    color: #00ff00;
+                    border: 2px solid #00ffff;
+                    border-radius: 5px;
+                    padding: 10px;
+                    font-family: 'Consolas';
+                    font-size: 12pt;
+                }
+            """)
+
+
+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,35 +54,37 @@ 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 识别潜在风险点（SQL操作、文件操作、用户输入点、文件上传漏洞、CSRF、SSRF、XSS、RCE、OWASP top10等漏洞）
-   1.2 验证漏洞可利用性
-   1.3 按CVSS评分标准评估风险等级
 
-2. 输出规则：
-   - 仅输出确认存在的高危/中危漏洞
-   - 使用严格格式：[风险等级] 类型 - 位置:行号 - 50字内描述
-   - 禁止解释漏洞原理
-   - 禁止给出修复建议
-   - 如果有可能，给出POC（HTTP请求数据包）
+    1. 漏洞分析流程：
+       1.1 识别潜在风险点（SQL操作、文件操作、用户输入点、文件上传漏洞、CSRF、SSRF、XSS、RCE、OWASP top10等漏洞）
+       1.2 验证漏洞可利用性
+       1.3 按CVSS评分标准评估风险等级
 
-3. 输出示例（除此外不要有任何输出）：
-   [高危] SQL注入 - user_login.php:32 - 未过滤的$_GET参数直接拼接SQL查询
-   [POC]POST /login.php HTTP/1.1
-   Host: example.com
-   Content-Type: application/x-www-form-urlencoded
-   [中危] XSS - comment.jsp:15 - 未转义的userInput输出到HTML
-   [POC]POST /login.php HTTP/1.1
-   Host: example.com
-   Content-Type: application/x-www-form-urlencoded
+    2. 输出规则：
+       - 仅输出确认存在的高危/中危漏洞
+       - 使用严格格式：[风险等级] 类型 - 位置:行号 - 50字内描述
+       - 禁止解释漏洞原理
+       - 禁止给出修复建议
+       - 如果有可能，给出POC（HTTP请求数据包）
 
-4. 当前代码（仅限分析）：
-{content[:3000]}"""
+    3. 输出示例（除此外不要有任何输出）：
+       [高危] SQL注入 - user_login.php:32 - 未过滤的$_GET参数直接拼接SQL查询
+       [POC]POST /login.php HTTP/1.1
+       Host: example.com
+       Content-Type: application/x-www-form-urlencoded
+       [中危] XSS - comment.jsp:15 - 未转义的userInput输出到HTML
+       [POC]POST /login.php HTTP/1.1
+       Host: example.com
+       Content-Type: application/x-www-form-urlencoded
+
+    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. 检测要求：         
-    请分析以下文件内容是否为WebShell或内存马。要求：
-    1. 检查PHP/JSP/ASP等WebShell特征（如加密函数、执行系统命令、文件操作）
-    2. 识别内存马特征（如无文件落地、进程注入、异常网络连接）
-    3. 分析代码中的可疑功能（如命令执行、文件上传、信息收集）
-    4. 检查混淆编码、加密手段等规避技术
+    1. 检测要求：         
+        请分析以下文件内容是否为WebShell或内存马。要求：
+        1. 检查PHP/JSP/ASP等WebShell特征（如加密函数、执行系统命令、文件操作）
+        2. 识别内存马特征（如无文件落地、进程注入、异常网络连接）
+        3. 分析代码中的可疑功能（如命令执行、文件上传、信息收集）
+        4. 检查混淆编码、加密手段等规避技术
 
-2. 判断规则：
-   - 仅当确认恶意性时报告
-   - 输出格式：🔴 [高危] Webshell - 文件名:行号 - 检测到[特征1+特征2+...]
+    2. 判断规则：
+       - 仅当确认恶意性时报告
+       - 输出格式：🔴 [高危] Webshell - 文件名:行号 - 检测到[特征1+特征2+...]
 
-3. 输出示例（严格按照此格式输出，不要有任何的补充，如果未检测到危险，则不输出，除此之外，不要有任何输出）：
-   🔴 [高危] Webshell - malicious.php:8 - 检测到[system执行+base64解码+错误抑制]
+    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):
@@ -148,20 +192,20 @@ class CyberScanner(QtWidgets.QMainWindow):
         # 目录选择按钮
         self.btn_select = QtWidgets.QPushButton("📁 激活数据源")
         self.btn_select.setStyleSheet("""
-            QPushButton {
-                background-color: #002b2b;
-                color: #00ffff;
-                border: 2px solid #008080;
-                padding: 12px;
-                font-size: 14pt;
-                font-weight: bold;
-                border-radius: 5px;
-            }
-            QPushButton:hover {
-                background-color: #004d4d;
-                border-color: #00ffff;
-            }
-        """)
+                QPushButton {
+                    background-color: #002b2b;
+                    color: #00ffff;
+                    border: 2px solid #008080;
+                    padding: 12px;
+                    font-size: 14pt;
+                    font-weight: bold;
+                    border-radius: 5px;
+                }
+                QPushButton:hover {
+                    background-color: #004d4d;
+                    border-color: #00ffff;
+                }
+            """)
         self.btn_select.clicked.connect(self.select_directory)
         left_layout.addWidget(self.btn_select)
 
@@ -171,25 +215,25 @@ class CyberScanner(QtWidgets.QMainWindow):
         left_layout.addWidget(self.lbl_path)
 
         # 模式选择
-         # 模式选择
+        # 模式选择
         mode_group = QtWidgets.QGroupBox("🔧 检测模式")
         mode_group.setStyleSheet("""
-            QGroupBox {
-                color: #00ff00;
-                border: 1px solid #00ffff;
-                margin-top: 10px;
-                font-size: 12pt;
-            }
-        """)
+                QGroupBox {
+                    color: #00ff00;
+                    border: 1px solid #00ffff;
+                    margin-top: 10px;
+                    font-size: 12pt;
+                }
+            """)
         mode_layout = QtWidgets.QVBoxLayout()
         self.radio_audit = QtWidgets.QRadioButton("代码安全审计")
         self.radio_webshell = QtWidgets.QRadioButton("Webshell检测")
         self.radio_audit.setChecked(True)
         for rb in [self.radio_audit, self.radio_webshell]:
             rb.setStyleSheet("""
-                QRadioButton { color: #00ff00; padding: 8px; }
-                QRadioButton::indicator { width: 20px; height: 20px; }
-            """)
+                    QRadioButton { color: #00ff00; padding: 8px; }
+                    QRadioButton::indicator { width: 20px; height: 20px; }
+                """)
             mode_layout.addWidget(rb)
         mode_group.setLayout(mode_layout)
         left_layout.addWidget(mode_group)
@@ -198,9 +242,9 @@ class CyberScanner(QtWidgets.QMainWindow):
         self.checkbox_audit_js = QtWidgets.QCheckBox("审计 静态 文件")
         self.checkbox_audit_js.setChecked(True)  # 默认选中
         self.checkbox_audit_js.setStyleSheet("""
-            QCheckBox { color: #00ff00; padding: 8px; }
-            QCheckBox::indicator { width: 20px; height: 20px; }
-        """)
+                QCheckBox { color: #00ff00; padding: 8px; }
+                QCheckBox::indicator { width: 20px; height: 20px; }
+            """)
         left_layout.addWidget(self.checkbox_audit_js)
         # 文件树
         self.file_tree = QtWidgets.QTreeView()
@@ -208,35 +252,35 @@ class CyberScanner(QtWidgets.QMainWindow):
         self.file_model.setRootPath("")
         self.file_tree.setModel(self.file_model)
         self.file_tree.setStyleSheet("""
-            QTreeView {
-                background-color: #001a1a;
-                color: #00ff00;
-                border: 1px solid #008080;
-                font-family: 'Consolas';
-            }
-            QTreeView::item:hover { background-color: #003333; }
-        """)
+                QTreeView {
+                    background-color: #001a1a;
+                    color: #00ff00;
+                    border: 1px solid #008080;
+                    font-family: 'Consolas';
+                }
+                QTreeView::item:hover { background-color: #003333; }
+            """)
         left_layout.addWidget(self.file_tree)
 
         # 扫描按钮
         self.btn_scan = QtWidgets.QPushButton("🚨 启动扫描协议")
         self.btn_scan.setStyleSheet("""
-            QPushButton {
-                background-color: #004d4d;
-                color: #00ffff;
-                border: 2px solid #00ffff;
-                padding: 15px;
-                font-size: 16pt;
-                font-weight: bold;
-                border-radius: 5px;
-            }
-            QPushButton:disabled { 
-                background-color: #002b2b;
-                color: #008080;
-                border-color: #004d4d;
-            }
-            QPushButton:hover { background-color: #006666; }
-        """)
+                QPushButton {
+                    background-color: #004d4d;
+                    color: #00ffff;
+                    border: 2px solid #00ffff;
+                    padding: 15px;
+                    font-size: 16pt;
+                    font-weight: bold;
+                    border-radius: 5px;
+                }
+                QPushButton:disabled { 
+                    background-color: #002b2b;
+                    color: #008080;
+                    border-color: #004d4d;
+                }
+                QPushButton:hover { background-color: #006666; }
+            """)
         self.btn_scan.clicked.connect(self.start_scan)
         self.btn_scan.setEnabled(False)
         left_layout.addWidget(self.btn_scan)
@@ -252,17 +296,17 @@ class CyberScanner(QtWidgets.QMainWindow):
         self.status_bar = QtWidgets.QStatusBar()
         self.setStatusBar(self.status_bar)
         self.status_bar.setStyleSheet("""
-            QStatusBar {
-                background-color: #000d1a;
-                color: #00ff00;
-                border-top: 1px solid #00ffff;
-                font-family: 'Consolas';
-            }
-        """)
+                QStatusBar {
+                    background-color: #000d1a;
+                    color: #00ff00;
+                    border-top: 1px solid #00ffff;
+                    font-family: 'Consolas';
+                }
+            """)
 
     def select_directory(self):
         directory = QtWidgets.QFileDialog.getExistingDirectory(
-            self, 
+            self,
             "选择代码矩阵接入点",
             "",
             QtWidgets.QFileDialog.ShowDirsOnly
@@ -281,34 +325,37 @@ class CyberScanner(QtWidgets.QMainWindow):
 
         root_path = self.file_model.filePath(root_index)
         self.files_content = self.scan_code_files(root_path)
-        
+
         if self.radio_audit.isChecked():
             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']
-        
+
         # 如果用户选择不审计 静态 文件，则从允许的扩展名中移除 .js
         if not self.checkbox_audit_js.isChecked():
             allowed_ext.remove('.js')
             allowed_ext.remove('.html')
-        
+
         code_files = {}
-        
+
         for root, _, files in os.walk(directory):
             for file in files:
                 if os.path.splitext(file)[1].lower() in allowed_ext:
@@ -326,7 +373,7 @@ class CyberScanner(QtWidgets.QMainWindow):
 
     def show_results(self, report):
         self.btn_scan.setEnabled(True)
-        
+
         if self.radio_webshell.isChecked():
             self.result_display.append("\n🔍 Webshell检测完成！结果如下：\n")
             report = re.sub(r'🔴 \[高危\]', '🔴 [高危]', report)
@@ -335,21 +382,26 @@ class CyberScanner(QtWidgets.QMainWindow):
             self.result_display.append("\n🔥 代码审计完成！发现以下安全漏洞：\n")
             report = re.sub(r'\[高危\]', '[高危]', report)
             report = re.sub(r'\[中危\]', '[中危]', report)
-        
+
         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]
 
     app = QtWidgets.QApplication([])
     app.setStyle('Fusion')
-    
+
     font = QtGui.QFont("Consolas", 10)
     app.setFont(font)
-    
+
     window = CyberScanner()
     window.show()
     app.exec_()