admin/DeepSeekSelfTool
1
0
Fork 0
mirror of https://github.com/ChinaRan0/DeepSeekSelfTool.git synced 2025-07-10 08:23:22 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #9 from s1kSec/1.1

Browse Source 
增加了进度展示以及预测剩余时间
This commit is contained in:
宗宸·谢尔比 2025-02-18 08:15:20 +08:00 committed by GitHub
commit c25326d636
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 179 additions and 127 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

306
ollamaMain.py
View File

@ -1,29 +1,52 @@
import os import os
import re import re
import json import json
import time
import requests import requests
from PyQt5 import QtCore, QtGui, QtWidgets from PyQt5 import QtCore, QtGui, QtWidgets
from PyQt5.QtCore import QThread, pyqtSignal from PyQt5.QtCore import QThread, pyqtSignal
from config import OLLAMA_API_URL, OLLAMA_MODEL # 用户自定义配置 from config import OLLAMA_API_URL, OLLAMA_MODEL # 用户自定义配置
class CyberTextEdit(QtWidgets.QTextEdit): class CyberTextEdit(QtWidgets.QTextEdit):
def __init__(self, parent=None): def __init__(self, parent=None):
super().__init__(parent) super().__init__(parent)
self.setStyleSheet(""" self.setStyleSheet("""
QTextEdit { QTextEdit {
background-color: #001a1a; background-color: #001a1a;
color: #00ff00; color: #00ff00;
border: 2px solid #00ffff; border: 2px solid #00ffff;
border-radius: 5px; border-radius: 5px;
padding: 10px; padding: 10px;
font-family: 'Consolas'; font-family: 'Consolas';
font-size: 12pt; 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): class HackerWorker(QThread):
analysis_complete = pyqtSignal(str) analysis_complete = pyqtSignal(str)
progress_update = pyqtSignal(str) progress_update = pyqtSignal(str)
button_text_update = pyqtSignal(str)
def __init__(self, files_content): def __init__(self, files_content):
super().__init__() super().__init__()
@ -31,35 +54,37 @@ class HackerWorker(QThread):
def run(self): def run(self):
full_report = [] full_report = []
for filepath, content in self.files_content.items(): start_time = time.time()
self.progress_update.emit(f"🔍 Analyzing {os.path.basename(filepath)}...") 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，请按以下要求分析代码： prompt = f"""【强制指令】你是一个专业的安全审计AI，请按以下要求分析代码：
1. 漏洞分析流程：
1.1 识别潜在风险点（SQL操作、文件操作、用户输入点、文件上传漏洞、CSRF、SSRF、XSS、RCE、OWASP top10等漏洞）
1.2 验证漏洞可利用性
1.3 按CVSS评分标准评估风险等级
2. 输出规则： 1. 漏洞分析流程：
- 仅输出确认存在的高危/中危漏洞 1.1 识别潜在风险点（SQL操作、文件操作、用户输入点、文件上传漏洞、CSRF、SSRF、XSS、RCE、OWASP top10等漏洞）
- 使用严格格式：[风险等级] 类型 - 位置:行号 - 50字内描述 1.2 验证漏洞可利用性
- 禁止解释漏洞原理 1.3 按CVSS评分标准评估风险等级
- 禁止给出修复建议
- 如果有可能，给出POC（HTTP请求数据包）
3. 输出示例（除此外不要有任何输出）： 2. 输出规则：
[高危] SQL注入 - user_login.php:32 - 未过滤的$_GET参数直接拼接SQL查询 - 仅输出确认存在的高危/中危漏洞
[POC]POST /login.php HTTP/1.1 - 使用严格格式：[风险等级] 类型 - 位置:行号 - 50字内描述
Host: example.com - 禁止解释漏洞原理
Content-Type: application/x-www-form-urlencoded - 禁止给出修复建议
[中危] XSS - comment.jsp:15 - 未转义的userInput输出到HTML - 如果有可能，给出POC（HTTP请求数据包）
[POC]POST /login.php HTTP/1.1
Host: example.com
Content-Type: application/x-www-form-urlencoded
4. 当前代码（仅限分析）： 3. 输出示例（除此外不要有任何输出）：
{content[:3000]}""" [高危] 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: try:
response = requests.post( response = requests.post(
f"{OLLAMA_HOST}/api/generate", f"{OLLAMA_HOST}/api/generate",
@ -71,43 +96,54 @@ class HackerWorker(QThread):
) )
result = json.loads(response.text)["response"] result = json.loads(response.text)["response"]
result = re.sub(r'<think>.*?</think>', '', result, flags=re.DOTALL) 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: except Exception as e:
full_report.append(f"❌ 错误：处理文件 {filepath} 时发生错误\n{str(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.analysis_complete.emit("\n".join(full_report))
self.button_text_update.emit("🚨 启动扫描协议")
class WebshellWorker(QThread): class WebshellWorker(QThread):
detection_complete = pyqtSignal(str) detection_complete = pyqtSignal(str)
progress_update = pyqtSignal(str) progress_update = pyqtSignal(str)
button_text_update = pyqtSignal(str)
def __init__(self, files_content): def __init__(self, files_content):
super().__init__() super().__init__()
self.files_content = files_content self.files_content = files_content
def run(self): def run(self):
start_time = time.time()
detection_results = [] detection_results = []
for filepath, content in self.files_content.items(): for index, (filepath, content) in enumerate(self.files_content.items(), start=1):
self.progress_update.emit(f"🕵️ 扫描 {os.path.basename(filepath)}...") self.progress_update.emit(
f"🕵️ 扫描 {os.path.basename(filepath)} ({index}/{len(self.files_content)})...")
prompt = f"""【Webshell检测指令】请严格按以下步骤分析代码： prompt = f"""【Webshell检测指令】请严格按以下步骤分析代码：
1. 检测要求： 1. 检测要求：
请分析以下文件内容是否为WebShell或内存马。要求： 请分析以下文件内容是否为WebShell或内存马。要求：
1. 检查PHP/JSP/ASP等WebShell特征（如加密函数、执行系统命令、文件操作） 1. 检查PHP/JSP/ASP等WebShell特征（如加密函数、执行系统命令、文件操作）
2. 识别内存马特征（如无文件落地、进程注入、异常网络连接） 2. 识别内存马特征（如无文件落地、进程注入、异常网络连接）
3. 分析代码中的可疑功能（如命令执行、文件上传、信息收集） 3. 分析代码中的可疑功能（如命令执行、文件上传、信息收集）
4. 检查混淆编码、加密手段等规避技术 4. 检查混淆编码、加密手段等规避技术
2. 判断规则： 2. 判断规则：
- 仅当确认恶意性时报告 - 仅当确认恶意性时报告
- 输出格式：🔴 [高危] Webshell - 文件名:行号 - 检测到[特征1+特征2+...] - 输出格式：🔴 [高危] Webshell - 文件名:行号 - 检测到[特征1+特征2+...]
3. 输出示例（严格按照此格式输出，不要有任何的补充，如果未检测到危险，则不输出，除此之外，不要有任何输出）： 3. 输出示例（严格按照此格式输出，不要有任何的补充，如果未检测到危险，则不输出，除此之外，不要有任何输出）：
🔴 [高危] Webshell - malicious.php:8 - 检测到[system执行+base64解码+错误抑制] 🔴 [高危] Webshell - malicious.php:8 - 检测到[system执行+base64解码+错误抑制]
4. 待分析代码： 4. 待分析代码：
{content[:3000]}""" {content[:3000]}"""
try: try:
response = requests.post( response = requests.post(
@ -120,11 +156,19 @@ class WebshellWorker(QThread):
) )
result = json.loads(response.text)["response"] result = json.loads(response.text)["response"]
result = re.sub(r'<think>.*?</think>', '', result, flags=re.DOTALL) 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: except Exception as e:
detection_results.append(f"❌ 错误：{filepath}\n{str(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.detection_complete.emit("\n".join(detection_results))
self.button_text_update.emit("🚨 启动扫描协议")
class CyberScanner(QtWidgets.QMainWindow): class CyberScanner(QtWidgets.QMainWindow):
def __init__(self): def __init__(self):
@ -148,20 +192,20 @@ class CyberScanner(QtWidgets.QMainWindow):
# 目录选择按钮 # 目录选择按钮
self.btn_select = QtWidgets.QPushButton("📁 激活数据源") self.btn_select = QtWidgets.QPushButton("📁 激活数据源")
self.btn_select.setStyleSheet(""" self.btn_select.setStyleSheet("""
QPushButton { QPushButton {
background-color: #002b2b; background-color: #002b2b;
color: #00ffff; color: #00ffff;
border: 2px solid #008080; border: 2px solid #008080;
padding: 12px; padding: 12px;
font-size: 14pt; font-size: 14pt;
font-weight: bold; font-weight: bold;
border-radius: 5px; border-radius: 5px;
} }
QPushButton:hover { QPushButton:hover {
background-color: #004d4d; background-color: #004d4d;
border-color: #00ffff; border-color: #00ffff;
} }
""") """)
self.btn_select.clicked.connect(self.select_directory) self.btn_select.clicked.connect(self.select_directory)
left_layout.addWidget(self.btn_select) left_layout.addWidget(self.btn_select)
@ -171,25 +215,25 @@ class CyberScanner(QtWidgets.QMainWindow):
left_layout.addWidget(self.lbl_path) left_layout.addWidget(self.lbl_path)
# 模式选择 # 模式选择
# 模式选择 # 模式选择
mode_group = QtWidgets.QGroupBox("🔧 检测模式") mode_group = QtWidgets.QGroupBox("🔧 检测模式")
mode_group.setStyleSheet(""" mode_group.setStyleSheet("""
QGroupBox { QGroupBox {
color: #00ff00; color: #00ff00;
border: 1px solid #00ffff; border: 1px solid #00ffff;
margin-top: 10px; margin-top: 10px;
font-size: 12pt; font-size: 12pt;
} }
""") """)
mode_layout = QtWidgets.QVBoxLayout() mode_layout = QtWidgets.QVBoxLayout()
self.radio_audit = QtWidgets.QRadioButton("代码安全审计") self.radio_audit = QtWidgets.QRadioButton("代码安全审计")
self.radio_webshell = QtWidgets.QRadioButton("Webshell检测") self.radio_webshell = QtWidgets.QRadioButton("Webshell检测")
self.radio_audit.setChecked(True) self.radio_audit.setChecked(True)
for rb in [self.radio_audit, self.radio_webshell]: for rb in [self.radio_audit, self.radio_webshell]:
rb.setStyleSheet(""" rb.setStyleSheet("""
QRadioButton { color: #00ff00; padding: 8px; } QRadioButton { color: #00ff00; padding: 8px; }
QRadioButton::indicator { width: 20px; height: 20px; } QRadioButton::indicator { width: 20px; height: 20px; }
""") """)
mode_layout.addWidget(rb) mode_layout.addWidget(rb)
mode_group.setLayout(mode_layout) mode_group.setLayout(mode_layout)
left_layout.addWidget(mode_group) left_layout.addWidget(mode_group)
@ -198,9 +242,9 @@ class CyberScanner(QtWidgets.QMainWindow):
self.checkbox_audit_js = QtWidgets.QCheckBox("审计 静态 文件") self.checkbox_audit_js = QtWidgets.QCheckBox("审计 静态 文件")
self.checkbox_audit_js.setChecked(True) # 默认选中 self.checkbox_audit_js.setChecked(True) # 默认选中
self.checkbox_audit_js.setStyleSheet(""" self.checkbox_audit_js.setStyleSheet("""
QCheckBox { color: #00ff00; padding: 8px; } QCheckBox { color: #00ff00; padding: 8px; }
QCheckBox::indicator { width: 20px; height: 20px; } QCheckBox::indicator { width: 20px; height: 20px; }
""") """)
left_layout.addWidget(self.checkbox_audit_js) left_layout.addWidget(self.checkbox_audit_js)
# 文件树 # 文件树
self.file_tree = QtWidgets.QTreeView() self.file_tree = QtWidgets.QTreeView()
@ -208,35 +252,35 @@ class CyberScanner(QtWidgets.QMainWindow):
self.file_model.setRootPath("") self.file_model.setRootPath("")
self.file_tree.setModel(self.file_model) self.file_tree.setModel(self.file_model)
self.file_tree.setStyleSheet(""" self.file_tree.setStyleSheet("""
QTreeView { QTreeView {
background-color: #001a1a; background-color: #001a1a;
color: #00ff00; color: #00ff00;
border: 1px solid #008080; border: 1px solid #008080;
font-family: 'Consolas'; font-family: 'Consolas';
} }
QTreeView::item:hover { background-color: #003333; } QTreeView::item:hover { background-color: #003333; }
""") """)
left_layout.addWidget(self.file_tree) left_layout.addWidget(self.file_tree)
# 扫描按钮 # 扫描按钮
self.btn_scan = QtWidgets.QPushButton("🚨 启动扫描协议") self.btn_scan = QtWidgets.QPushButton("🚨 启动扫描协议")
self.btn_scan.setStyleSheet(""" self.btn_scan.setStyleSheet("""
QPushButton { QPushButton {
background-color: #004d4d; background-color: #004d4d;
color: #00ffff; color: #00ffff;
border: 2px solid #00ffff; border: 2px solid #00ffff;
padding: 15px; padding: 15px;
font-size: 16pt; font-size: 16pt;
font-weight: bold; font-weight: bold;
border-radius: 5px; border-radius: 5px;
} }
QPushButton:disabled { QPushButton:disabled {
background-color: #002b2b; background-color: #002b2b;
color: #008080; color: #008080;
border-color: #004d4d; border-color: #004d4d;
} }
QPushButton:hover { background-color: #006666; } QPushButton:hover { background-color: #006666; }
""") """)
self.btn_scan.clicked.connect(self.start_scan) self.btn_scan.clicked.connect(self.start_scan)
self.btn_scan.setEnabled(False) self.btn_scan.setEnabled(False)
left_layout.addWidget(self.btn_scan) left_layout.addWidget(self.btn_scan)
@ -252,17 +296,17 @@ class CyberScanner(QtWidgets.QMainWindow):
self.status_bar = QtWidgets.QStatusBar() self.status_bar = QtWidgets.QStatusBar()
self.setStatusBar(self.status_bar) self.setStatusBar(self.status_bar)
self.status_bar.setStyleSheet(""" self.status_bar.setStyleSheet("""
QStatusBar { QStatusBar {
background-color: #000d1a; background-color: #000d1a;
color: #00ff00; color: #00ff00;
border-top: 1px solid #00ffff; border-top: 1px solid #00ffff;
font-family: 'Consolas'; font-family: 'Consolas';
} }
""") """)
def select_directory(self): def select_directory(self):
directory = QtWidgets.QFileDialog.getExistingDirectory( directory = QtWidgets.QFileDialog.getExistingDirectory(
self, self,
"选择代码矩阵接入点", "选择代码矩阵接入点",
"", "",
QtWidgets.QFileDialog.ShowDirsOnly QtWidgets.QFileDialog.ShowDirsOnly
@ -281,34 +325,37 @@ class CyberScanner(QtWidgets.QMainWindow):
root_path = self.file_model.filePath(root_index) root_path = self.file_model.filePath(root_index)
self.files_content = self.scan_code_files(root_path) self.files_content = self.scan_code_files(root_path)
if self.radio_audit.isChecked(): if self.radio_audit.isChecked():
worker = HackerWorker(self.files_content) worker = HackerWorker(self.files_content)
init_msg = "🚀 启动深度代码分析协议..." init_msg = "🚀 启动深度代码分析协议..."
complete_signal = worker.analysis_complete complete_signal = worker.analysis_complete
worker.button_text_update.connect(self.update_button_text)
else: else:
worker = WebshellWorker(self.files_content) worker = WebshellWorker(self.files_content)
init_msg = "🕵️ 启动Webshell检测协议..." init_msg = "🕵️ 启动Webshell检测协议..."
complete_signal = worker.detection_complete complete_signal = worker.detection_complete
worker.button_text_update.connect(self.update_button_text)
self.scan_thread = worker self.scan_thread = worker
self.scan_thread.progress_update.connect(self.update_status) self.scan_thread.progress_update.connect(self.update_status)
complete_signal.connect(self.show_results) complete_signal.connect(self.show_results)
self.scan_thread.button_text_update.connect(self.update_button_text)
self.scan_thread.start() self.scan_thread.start()
self.btn_scan.setEnabled(False) 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): def scan_code_files(self, directory):
allowed_ext = ['.php', '.jsp', '.asp', '.js', '.html', '.py', '.java'] allowed_ext = ['.php', '.jsp', '.asp', '.js', '.html', '.py', '.java']
# 如果用户选择不审计 静态 文件，则从允许的扩展名中移除 .js # 如果用户选择不审计 静态 文件，则从允许的扩展名中移除 .js
if not self.checkbox_audit_js.isChecked(): if not self.checkbox_audit_js.isChecked():
allowed_ext.remove('.js') allowed_ext.remove('.js')
allowed_ext.remove('.html') allowed_ext.remove('.html')
code_files = {} code_files = {}
for root, _, files in os.walk(directory): for root, _, files in os.walk(directory):
for file in files: for file in files:
if os.path.splitext(file)[1].lower() in allowed_ext: if os.path.splitext(file)[1].lower() in allowed_ext:
@ -326,7 +373,7 @@ class CyberScanner(QtWidgets.QMainWindow):
def show_results(self, report): def show_results(self, report):
self.btn_scan.setEnabled(True) self.btn_scan.setEnabled(True)
if self.radio_webshell.isChecked(): if self.radio_webshell.isChecked():
self.result_display.append("\n🔍 Webshell检测完成！结果如下：\n") self.result_display.append("\n🔍 Webshell检测完成！结果如下：\n")
report = re.sub(r'🔴 \[高危\]', '🔴 [高危]', report) report = re.sub(r'🔴 \[高危\]', '🔴 [高危]', report)
@ -335,21 +382,26 @@ class CyberScanner(QtWidgets.QMainWindow):
self.result_display.append("\n🔥 代码审计完成！发现以下安全漏洞：\n") self.result_display.append("\n🔥 代码审计完成！发现以下安全漏洞：\n")
report = re.sub(r'\[高危\]', '[高危]', report) report = re.sub(r'\[高危\]', '[高危]', report)
report = re.sub(r'\[中危\]', '[中危]', report) report = re.sub(r'\[中危\]', '[中危]', report)
self.result_display.append(report) self.result_display.append(report)
self.status_bar.showMessage("✅ 扫描完成") 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__": if __name__ == "__main__":
# 保持源文本的核心内容不变 # 保持源文本的核心内容不变
OLLAMA_HOST = OLLAMA_API_URL.split('/api')[0] OLLAMA_HOST = OLLAMA_API_URL.split('/api')[0]
app = QtWidgets.QApplication([]) app = QtWidgets.QApplication([])
app.setStyle('Fusion') app.setStyle('Fusion')
font = QtGui.QFont("Consolas", 10) font = QtGui.QFont("Consolas", 10)
app.setFont(font) app.setFont(font)
window = CyberScanner() window = CyberScanner()
window.show() window.show()
app.exec_() app.exec_()