From 6fd4e6e993c2f771a8bb4261c6879957db38195f Mon Sep 17 00:00:00 2001 From: kukud_xiaoliu Date: Mon, 17 Feb 2025 17:29:41 +0800 Subject: [PATCH] =?UTF-8?q?=E5=A2=9E=E5=8A=A0=E4=BA=86=E8=BF=9B=E5=BA=A6?= =?UTF-8?q?=E5=B1=95=E7=A4=BA=E4=BB=A5=E5=8F=8A=E9=A2=84=E6=B5=8B=E5=89=A9?= =?UTF-8?q?=E4=BD=99=E6=97=B6=E9=97=B4?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- ollamaMain.py | 306 +++++++++++++++++++++++++++++--------------------- 1 file changed, 179 insertions(+), 127 deletions(-) diff --git a/ollamaMain.py b/ollamaMain.py index 0bb3c89..cfde940 100644 --- a/ollamaMain.py +++ b/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'.*?', '', 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'.*?', '', 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_()