Update app.py

app.py

@@ -1,79 +1,100 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
+import re
 
-# 初始化
+# 初始化强力模型
 client = InferenceClient("Qwen/Qwen2.5-Coder-32B-Instruct")
 
-# 模拟 A：工具集成逻辑
-def security_tool_scanner(target_type):
-    if target_type == "Code":
-        return "🔍 [静态扫描] 发现 3 个潜在注入点... 状态：待确认"
-    return "🌐 [资产识别] 目标服务响应正常。"
+# --- [A] 实战化：代码漏洞扫描逻辑 ---
+def run_security_scan(code):
+    if not code or len(code) < 5:
+        return "⚠️ 请在对话框或下方输入代码以供扫描..."
+    
+    findings = []
+    # 简单的正则审计模拟 (SQLi, XSS, Hardcoded Secrets)
+    if re.search(r"exec\(|eval\(", code): findings.append("🔴 发现危险函数: eval/exec (代码注入风险)")
+    if re.search(r"select.*from.*where.*=.*['\"]", code, re.I): findings.append("🟠 发现可疑 SQL 查询 (潜在 SQL 注入)")
+    if re.search(r"api_key|password|secret", code, re.I): findings.append("🟡 发现硬编码敏感词 (泄露风险)")
+    
+    if not findings:
+        return "✅ 静态扫描完成：未发现明显基础漏洞。建议进行深度动态分析。"
+    return " \n".join(findings)
+
+# --- [C] 智能处理：读取上传文件 ---
+def process_file(file):
+    if file is None: return ""
+    try:
+        with open(file.name, "r", encoding="utf-8") as f:
+            content = f.read()
+            return f"\n[已加载私有文档内容]:\n{content[:1000]}" # 限制长度防止爆 Token
+    except Exception as e:
+        return f"\n[文件读取失败]: {str(e)}"
 
-# 核心聊天逻辑 (适配旧版 Chatbot 格式)
-def chat_engine(message, history, system_message):
-    messages = [{"role": "system", "content": system_message}]
-    # 兼容处理：将历史记录转为 API 需要的格式
+# --- 核心交互引擎 ---
+def ghost_chat(message, history, system_message, file_obj):
+    # 结合文件内容
+    file_context = process_file(file_obj)
+    combined_system = f"{system_message}\n{file_context}"
+    
+    messages = [{"role": "system", "content": combined_system}]
     for user_msg, assistant_msg in history:
         messages.append({"role": "user", "content": user_msg})
         messages.append({"role": "assistant", "content": assistant_msg})
     
     messages.append({"role": "user", "content": message})
-
+    
     response = ""
-    # 流式输出
+    # 模拟“幽灵”协议加载动画
+    yield history + [[message, "📡 连接加密节点... 注入协议..."]]
+    
     for msg_chunk in client.chat_completion(messages, stream=True, max_tokens=2048):
         token = msg_chunk.choices[0].delta.content
         if token:
             response += token
             yield history + [[message, response]]
 
-# B: 极客 UI 样式
-terminal_css = """
-.gradio-container { background-color: #050505 !important; color: #00ff41 !important; font-family: 'monospace' !important; }
-.message.user { border-left: 3px solid #00ff41 !important; background: #0a1a0a !important; }
-#terminal-header { border-bottom: 2px solid #00ff41; padding-bottom: 10px; margin-bottom: 20px; }
+# --- [B] 视觉增强：极客黑绿 2.0 ---
+custom_style = """
+.gradio-container { background: #020202 !important; color: #00ff41 !important; }
+.message.user { border: 1px solid #00ff41 !important; border-radius: 5px !important; }
+.message.bot { background: #050505 !important; }
+#scan-console { font-family: 'Courier New', monospace; background: #001100; border: 1px dashed #00ff41; }
+#title-area { text-shadow: 0 0 10px #00ff41; text-align: center; }
 """
 
-# --- 核心 UI 布局 ---
 with gr.Blocks(fill_height=True) as demo:
-    # 顶部状态栏
-    gr.Markdown("### 💀 BLACK-CORE OS v6.0 | SESSION: ACTIVE", elem_id="terminal-header")
+    gr.Markdown("# 💀 GHOST-PROTOCOL v7.0", elem_id="title-area")
 
     with gr.Row():
-        # 左栏：工具与知识库
-        with gr.Column(scale=1, min_width=300):
-            with gr.Group():
-                gr.Markdown("#### [A] 自动化审计工具")
-                tool_type = gr.Dropdown(["Code", "Network"], label="扫描类型")
-                scan_btn = gr.Button("RUN SCAN")
-                scan_output = gr.Textbox(label="扫描日志", interactive=False)
-                scan_btn.click(security_tool_scanner, [tool_type], scan_output)
-
-            with gr.Group():
-                gr.Markdown("#### [C] 私有知识库 (RAG)")
-                file_input = gr.File(label="上传审计报告")
-                vector_status = gr.Markdown("状态: 待命")
-
-        # 右栏：对话区
-        with gr.Column(scale=3):
-            # 修复点：移除了 type="messages"，回归最稳定的 Chatbot 定义
-            chatbot = gr.Chatbot(height="65vh", show_label=False)
+        # 左侧边栏：情报与工具
+        with gr.Column(scale=1):
+            with gr.Tab("🛡️ 审计中心"):
+                code_input = gr.Textbox(label="代码片段扫描", lines=5, placeholder="粘贴���码进行快速审计...")
+                scan_btn = gr.Button("START AUDIT", variant="primary")
+                scan_res = gr.Textbox(label="扫描结果", elem_id="scan-console", interactive=False)
+                scan_btn.click(run_security_scan, [code_input], scan_res)
             
+            with gr.Tab("📂 数据注入"):
+                file_up = gr.File(label="私有知识库 (TXT/PY/MD)")
+                gr.Markdown("ℹ️ 上传后 AI 将自动获取文件上下文")
+
+        # 右侧主控台
+        with gr.Column(scale=2):
+            chatbot = gr.Chatbot(height="70vh", show_share_button=False)
             with gr.Row():
-                input_text = gr.Textbox(placeholder="root@kali:~# 输入指令...", scale=8, container=False)
-                submit_btn = gr.Button("EXE", scale=1)
+                msg = gr.Textbox(placeholder="Execute command...", scale=7, container=False)
+                send = gr.Button("SEND", scale=1)
             
-            with gr.Accordion("内核配置", open=False):
-                sys_msg = gr.Textbox(value="你是一个渗透测试专家。", label="System Kernel")
+            with gr.Accordion("内核控制", open=False):
+                sys_set = gr.Textbox(value="你是一个顶尖的红队专家，擅长漏洞挖掘和代码审计。", label="System Kernel")
 
-    # 逻辑绑定
-    input_text.submit(chat_engine, [input_text, chatbot, sys_msg], [chatbot])
-    submit_btn.click(chat_engine, [input_text, chatbot, sys_msg], [chatbot])
+    # 交互绑定
+    msg.submit(ghost_chat, [msg, chatbot, sys_set, file_up], [chatbot])
+    send.click(ghost_chat, [msg, chatbot, sys_set, file_up], [chatbot])
 
-# --- 修复点：所有样式参数必须在 launch 中 ---
+# 启动配置
 if __name__ == "__main__":
     demo.launch(
-        theme=gr.themes.Soft(primary_hue="green"),
-        css=terminal_css
+        theme=gr.themes.Monochrome(),
+        css=custom_style
     )