import os import re import threading from threading import Event from app import audit_code, update_config, GLOBAL_CONFIG from app.utils import get_now_date from logger import Logger from PyQt6.QtGui import QColor, QGuiApplication, QTextCursor from PyQt6.QtWidgets import ( QWidget, QVBoxLayout, QHBoxLayout, QLabel, QLineEdit, QPushButton, QFileDialog, QTextEdit, QComboBox ) BACKGROUND_COLOR = '#dcdcdc' ANSI_ESCAPE = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])') ANSI_COLOR_REGEX = re.compile(r'\x1B\[(?:([0-9]+);)?([0-9]+)m') ANSI_COLOR_MAP = { '94': QColor(0, 0, 200), '92': QColor(0, 128, 0), '93': QColor(255, 127, 0), '91': QColor(220, 0, 0), '95': QColor(180, 0, 180) } def convert_ansi_to_rich_text(text): segments = [] pos = 0 for match in ANSI_COLOR_REGEX.finditer(text): start, end = match.span() if start > pos: segments.append(text[pos:start]) color_code = match.group(2) if color_code in ANSI_COLOR_MAP: color = ANSI_COLOR_MAP[color_code] html_color = color.name() segments.append(f'') else: segments.append('') pos = end segments.append(text[pos:]) segments.append('') rich_text = ''.join(segments) rich_text = ANSI_ESCAPE.sub('', rich_text) return rich_text class MainWindow(QWidget): def __init__(self): self.event = Event() self.log = Logger('ui', callback=self.process_output_callback) super().__init__() self.init_ui() def init_ui(self): main_layout = QVBoxLayout() dir_lang_layout = QHBoxLayout() # 目录选择 dir_layout = QHBoxLayout() self.dir_label = QLabel('项目目录:') self.dir_input = QLineEdit() self.dir_button = QPushButton('选择') self.dir_button.clicked.connect(self.select_directory) dir_layout.addWidget(self.dir_label) dir_layout.addWidget(self.dir_input) dir_layout.addWidget(self.dir_button) dir_lang_layout.addLayout(dir_layout) # 语言选择 languages = ['c', 'cpp', 'go', 'php', 'jsp', 'java', 'python', 'javascript'] self.lang_label = QLabel('项目语言:') self.lang_combobox = QComboBox() self.lang_combobox.addItems(languages) dir_lang_layout.addWidget(self.lang_label) dir_lang_layout.addWidget(self.lang_combobox) main_layout.addLayout(dir_lang_layout) # 配置信息 config_layout = QHBoxLayout() self.base_url_label = QLabel('接口地址:') self.base_url_input = QLineEdit() self.api_key_label = QLabel('模型密钥:') self.api_key_input = QLineEdit() self.api_key_input.setEchoMode(QLineEdit.EchoMode.Password) config_layout.addWidget(self.base_url_label) config_layout.addWidget(self.base_url_input) config_layout.addWidget(self.api_key_label) config_layout.addWidget(self.api_key_input) main_layout.addLayout(config_layout) model_layout = QHBoxLayout() self.reasoning_model_label = QLabel('推理模型:') self.reasoning_model_input = QLineEdit() self.embedding_model_label = QLabel('嵌入模型:') self.embedding_model_input = QLineEdit() model_layout.addWidget(self.reasoning_model_label) model_layout.addWidget(self.reasoning_model_input) model_layout.addWidget(self.embedding_model_label) model_layout.addWidget(self.embedding_model_input) main_layout.addLayout(model_layout) # 按钮部分 button_layout = QHBoxLayout() self.start_button = QPushButton('开始审计') self.start_button.clicked.connect(self.start_process) self.stop_button = QPushButton('终止审计') self.stop_button.clicked.connect(self.stop_process) self.update_button = QPushButton('更新配置') self.update_button.clicked.connect(self.update_config) self.clear_button = QPushButton('清空输出') self.clear_button.clicked.connect(self.clear_panel) button_layout.addWidget(self.start_button) button_layout.addWidget(self.stop_button) button_layout.addWidget(self.update_button) button_layout.addWidget(self.clear_button) main_layout.addLayout(button_layout) # 实时输出 output_layout = QVBoxLayout() # 过程输出 self.process_output_text = QTextEdit() self.process_output_text.setReadOnly(True) self.process_output_text.setStyleSheet(f'background-color: {BACKGROUND_COLOR};') output_layout.addWidget(self.process_output_text) # 结果输出 self.result_output_text = QTextEdit() self.result_output_text.setReadOnly(True) self.result_output_text.setStyleSheet(f'background-color: {BACKGROUND_COLOR};') output_layout.addWidget(self.result_output_text) output_layout.setStretch(0, 1) output_layout.setStretch(1, 2) main_layout.addLayout(output_layout) self.setLayout(main_layout) self.setWindowTitle('MollyAudit - created by yvling') screen = QGuiApplication.primaryScreen().geometry() window_width = 1000 window_height = 600 x = (screen.width() - window_width) // 2 y = (screen.height() - window_height) // 2 self.setGeometry(x, y, window_width, window_height) # 导出结果 export_button_layout = QHBoxLayout() self.export_button = QPushButton('导出结果') self.export_button.clicked.connect(self.export_result) export_button_layout.addStretch(1) # 添加伸缩项，使按钮靠右 export_button_layout.addWidget(self.export_button) main_layout.addLayout(export_button_layout) # 加载配置 self.base_url_input.setText(GLOBAL_CONFIG['base_url']) self.api_key_input.setText(GLOBAL_CONFIG['api_key']) self.reasoning_model_input.setText(GLOBAL_CONFIG['reasoning_model']) self.embedding_model_input.setText(GLOBAL_CONFIG['embedding_model']) def closeEvent(self, event): self.event.set() def clear_panel(self): self.process_output_text.clear() self.result_output_text.clear() def update_config(self): base_url = self.base_url_input.text() api_key = self.api_key_input.text() reasoning_model = self.reasoning_model_input.text() embedding_model = self.embedding_model_input.text() update_config('base_url', base_url) update_config('api_key', api_key) update_config('reasoning_model', reasoning_model) update_config('embedding_model', embedding_model) self.log.info('更新配置成功') def select_directory(self): directory = QFileDialog.getExistingDirectory(self, '选择项目目录') if directory: self.dir_input.setText(directory) def export_result(self): result_text = self.result_output_text.toPlainText() if result_text == '': self.log.warning('当前结果为空') return directory = QFileDialog.getExistingDirectory(self, '选择导出目录') if directory: file_name = f'molly-audit-{get_now_date()}.txt' file_path = os.path.join(directory, file_name).replace('\\', '/') try: with open(file_path, 'w', encoding='utf-8') as f: f.write(result_text) self.log.info(f'导出结果成功： {file_path}') except Exception as e: self.log.error(f'导出结果错误：{str(e)}') def process_output_callback(self, content): rich_text = convert_ansi_to_rich_text(content) self.process_output_text.append(rich_text) cursor = self.process_output_text.textCursor() cursor.movePosition(QTextCursor.MoveOperation.End) self.process_output_text.setTextCursor(cursor) self.process_output_text.ensureCursorVisible() def result_output_callback(self, content): self.result_output_text.append(f'{content}\n') cursor = self.result_output_text.textCursor() cursor.movePosition(QTextCursor.MoveOperation.End) self.result_output_text.setTextCursor(cursor) self.result_output_text.ensureCursorVisible() def start_process(self): selected_dir = self.dir_input.text() selected_lang = self.lang_combobox.currentText() base_url = self.base_url_input.text() api_key = self.api_key_input.text() reasoning_model = self.reasoning_model_input.text() embedding_model = self.embedding_model_input.text() if not selected_dir or not base_url or not api_key: self.log.error('请确保项目目录、接口地址和模型密钥等都已填写') return self.log.info('正在加载所需资源') try: threading.Thread( target=audit_code, args=( base_url, api_key, selected_dir, selected_lang, reasoning_model, embedding_model, self.process_output_callback, self.result_output_callback, self.event ) ).start() except Exception as e: self.log.error(f'发生异常：{str(e)}') finally: if 'OPENAI_API_BASE' in os.environ: del os.environ['OPENAI_API_BASE'] if 'OPENAI_API_KEY' in os.environ: del os.environ['OPENAI_API_KEY'] def stop_process(self): self.event.set() if 'OPENAI_API_BASE' in os.environ: del os.environ['OPENAI_API_BASE'] if 'OPENAI_API_KEY' in os.environ: del os.environ['OPENAI_API_KEY'] self.log.info('已终止代码审计流程')