import streamlit as st import os import io import json import csv # <--- 新增：用於處理 CSV import numpy as np import faiss import uuid import time import sys # === HuggingFace 模型相關套件 (替換為 InferenceClient) === try: from huggingface_hub import InferenceClient except ImportError: st.error("請檢查是否安裝了所有 Hugging Face 相關依賴：pip install huggingface-hub") # === LangChain/RAG 相關套件 (保持不變) === from langchain_community.embeddings import HuggingFaceEmbeddings from langchain_core.documents import Document from langchain_community.vectorstores import FAISS from langchain_community.vectorstores.utils import DistanceStrategy from langchain_community.docstore.in_memory import InMemoryDocstore # 嘗試匯入 pypdftry try: import pypdf except ImportError: pypdf = None # --- 頁面設定 --- st.set_page_config(page_title="Cybersecurity AI Assistant (Hugging Face RAG & Batch Analysis)", page_icon="🛡️", layout="wide") st.title("🛡️ Meta-Llama-3-8B-Instruct with FAISS RAG & Batch Analysis (Inference Client)") st.markdown("已啟用：**IndexFlatIP** + **L2 正規化** + **Hugging Face Inference Client (API)**。支援 JSON/CSV/TXT 執行批量分析。") if 'execute_batch_analysis' not in st.session_state: st.session_state.execute_batch_analysis = False if 'batch_results' not in st.session_state: st.session_state.batch_results = None if 'rag_current_file_key' not in st.session_state: st.session_state.rag_current_file_key = None if 'batch_current_file_key' not in st.session_state: # 修改變數名稱以反映多格式 st.session_state.batch_current_file_key = None if 'vector_store' not in st.session_state: st.session_state.vector_store = None if 'json_data_for_batch' not in st.session_state: # 變數名稱保留，但內容可能是轉換後的 dict st.session_state.json_data_for_batch = None # 設定模型 ID MODEL_ID = "meta-llama/Llama-4-Scout-17B-16E-Instruct" WINDOW_SIZE = 8 # --- 側邊欄設定 --- with st.sidebar: st.header("⚙️ 設定") if not os.environ.get("HF_TOKEN"): st.error("環境變數 **HF_TOKEN** 未設定。請設定後重新啟動應用程式。") st.info(f"LLM 模型：**{MODEL_ID}** (Hugging Face Inference API)") st.warning("⚠️ **注意**: 該模型使用 Inference API 呼叫，請確保您的 HF Token 具有存取權限。") st.divider() st.subheader("📂 檔案上傳") # === 1. 批量分析檔案 (修改處：支援多種格式) === batch_uploaded_file = st.file_uploader( "1️⃣ 上傳 **Log/Alert 檔案** (用於批量分析)", type=['json', 'csv', 'txt'], # <--- 修改：新增 csv 和 txt key="batch_uploader", help="支援 JSON (Array), CSV (含標題), TXT (每行一條 Log)" ) # === 2. RAG 知識庫檔案 === rag_uploaded_file = st.file_uploader( "2️⃣ 上傳 **RAG 參考知識庫** (Logs/PDF/Code 等)", type=['txt', 'py', 'log', 'csv', 'md', 'pdf'], key="rag_uploader" ) st.divider() st.subheader("💡 批量分析指令") analysis_prompt = st.text_area( "針對每個 Log/Alert 執行的指令", value="You are a security expert in charge of analyzing alerts related to Web Application Attacks and Brute Force & Reconnaissance. Respond with a clear, structured analysis using the following mandatory sections: \n\n- Priority: Provide the overall priority level. (Answer High risk, Medium risk, or Low risk only) \n- Explanation: If this alert is highly related to Web Application Attacks and Brute Force & Reconnaissance, explain the potential impact and why this specific alert requires attention. If not, **omit the explanation section**. \n- Action Plan: If this alert is highly related to Web Application Attacks and Brute Force & Reconnaissance, What should be the immediate steps to address this specific alert? If not, **omit the action plan section**. \n\nStrictly use the information in the provided Log.", height=200 ) st.markdown("此指令將對檔案中的**每一個 Log 條目**執行一次獨立分析。") if batch_uploaded_file: if st.button("🚀 執行批量分析"): if not os.environ.get("HF_TOKEN"): st.error("無法執行，環境變數 **HF_TOKEN** 未設定。") else: st.session_state.execute_batch_analysis = True else: st.info("請上傳 Log 檔案以啟用批量分析按鈕。") st.divider() st.subheader("🔍 RAG 檢索設定") similarity_threshold = st.slider("📐 Cosine Similarity 門檻", 0.0, 1.0, 0.4, 0.01) st.divider() st.subheader("模型參數") system_prompt = st.text_area("System Prompt", value="You are a Senior Security Analyst, named Ernest. You provide expert, authoritative, and concise advice on Information Security. Your analysis must be based strictly on the provided context.", height=100) max_output_tokens = st.slider("Max Output Tokens", 128, 4096, 2048, 128) temperature = st.slider("Temperature", 0.0, 1.0, 0.1, 0.1) top_p = st.slider("Top P", 0.1, 1.0, 0.95, 0.05) st.divider() if st.button("🗑️ 清除所有紀錄"): for key in list(st.session_state.keys()): del st.session_state[key] st.rerun() # --- 初始化 Hugging Face LLM Client --- @st.cache_resource def load_inference_client(model_id): if not os.environ.get("HF_TOKEN"): return None try: client = InferenceClient(model_id, token=os.environ.get("HF_TOKEN")) st.success(f"Hugging Face Inference Client **{model_id}** 載入成功。") return client except Exception as e: st.error(f"Hugging Face Inference Client 載入失敗: {e}") return None inference_client = None if os.environ.get("HF_TOKEN"): with st.spinner(f"正在連線到 Inference Client: {MODEL_ID}..."): inference_client = load_inference_client(MODEL_ID) if inference_client is None and os.environ.get("HF_TOKEN"): st.warning("Hugging Face Inference Client 無法連線。") elif not os.environ.get("HF_TOKEN"): st.error("請在環境變數中設定 HF_TOKEN。") # === Embedding 模型 (保持不變) === @st.cache_resource def load_embedding_model(): model_kwargs = {'device': 'cpu', 'trust_remote_code': True} encode_kwargs = {'normalize_embeddings': False} return HuggingFaceEmbeddings(model_name="BAAI/bge-large-zh-v1.5", model_kwargs=model_kwargs, encode_kwargs=encode_kwargs) with st.spinner("正在載入 Embedding 模型..."): embedding_model = load_embedding_model() # === 建立向量庫 / Search 函數 (保持不變) === def process_file_to_faiss(uploaded_file): text_content = "" try: if uploaded_file.type == "application/pdf": if pypdf: pdf_reader = pypdf.PdfReader(uploaded_file) for page in pdf_reader.pages: text_content += page.extract_text() + "\n" else: return None, "PDF library missing" else: stringio = io.StringIO(uploaded_file.getvalue().decode("utf-8")) text_content = stringio.read() if not text_content.strip(): return None, "File is empty" events = [line for line in text_content.splitlines() if line.strip()] docs = [Document(page_content=e) for e in events] if not docs: return None, "No documents created" embeddings = embedding_model.embed_documents([d.page_content for d in docs]) embeddings_np = np.array(embeddings).astype("float32") faiss.normalize_L2(embeddings_np) dimension = embeddings_np.shape[1] index = faiss.IndexFlatIP(dimension) index.add(embeddings_np) doc_ids = [str(uuid.uuid4()) for _ in range(len(docs))] docstore = InMemoryDocstore({_id: doc for _id, doc in zip(doc_ids, docs)}) index_to_docstore_id = {i: _id for i, _id in enumerate(doc_ids)} vector_store = FAISS(embedding_function=embedding_model, index=index, docstore=docstore, index_to_docstore_id=index_to_docstore_id, distance_strategy=DistanceStrategy.COSINE) return vector_store, f"{len(docs)} chunks created." except Exception as e: return None, f"Error: {str(e)}" def faiss_cosine_search_all(vector_store, query, threshold): q_emb = embedding_model.embed_query(query) q_emb = np.array([q_emb]).astype("float32") faiss.normalize_L2(q_emb) index = vector_store.index D, I = index.search(q_emb, k=index.ntotal) selected = [] for score, idx in zip(D[0], I[0]): if idx == -1: continue if score >= threshold: doc_id = vector_store.index_to_docstore_id[idx] doc = vector_store.docstore.search(doc_id) selected.append((doc, score)) selected.sort(key=lambda x: x[1], reverse=True) return selected # === Hugging Face 生成單一 Log 分析回答 (保持不變) === def generate_rag_response_hf_for_log(client, model_id, log_sequence_text, user_prompt, sys_prompt, vector_store, threshold, max_output_tokens, temperature, top_p): if client is None: return "ERROR: Client Error", "" context_text = "" if vector_store: selected = faiss_cosine_search_all(vector_store, log_sequence_text, threshold) if selected: retrieved_contents = [f"--- Reference Chunk (sim={score:.3f}) ---\n{doc.page_content}" for i, (doc, score) in enumerate(selected[:5])] context_text = "\n".join(retrieved_contents) rag_instruction = f"""=== RETRIEVED REFERENCE CONTEXT (Cosine ≥ {threshold}) ==={context_text if context_text else 'No relevant reference context found.'}=== END REFERENCE CONTEXT ===\nANALYSIS INSTRUCTION: {user_prompt}\nBased on the provided LOG SEQUENCE and REFERENCE CONTEXT, you must analyze the **entire sequence** to detect any continuous attack chains or evolving threats.""" log_content_section = f"""=== CURRENT LOG SEQUENCE TO ANALYZE (Window Size: {WINDOW_SIZE}) ===\n{log_sequence_text}\n=== END LOG SEQUENCE ===""" messages = [ {"role": "system", "content": sys_prompt}, {"role": "user", "content": f"{rag_instruction}\n\n{log_content_section}"} ] try: response_stream = client.chat_completion(messages, max_tokens=max_output_tokens, temperature=temperature, top_p=top_p, stream=False) if response_stream and response_stream.choices: return response_stream.choices[0].message.content.strip(), context_text else: return "Format Error", context_text except Exception as e: return f"Model Error: {str(e)}", context_text # ======================================================================= # === 檔案處理區塊 (RAG 檔案) === if rag_uploaded_file: file_key = f"vs_{rag_uploaded_file.name}_{rag_uploaded_file.size}" if st.session_state.rag_current_file_key != file_key or 'vector_store' not in st.session_state: with st.spinner(f"正在建立 RAG 參考知識庫 ({rag_uploaded_file.name})..."): vs, msg = process_file_to_faiss(rag_uploaded_file) if vs: st.session_state.vector_store = vs st.session_state.rag_current_file_key = file_key st.toast(f"RAG 參考知識庫已更新！{msg}", icon="✅") else: st.error(msg) elif 'vector_store' in st.session_state: del st.session_state.vector_store del st.session_state.rag_current_file_key st.info("RAG 檔案已移除，已清除相關知識庫。") # === 檔案處理區塊 (批量分析檔案 - 重大修改處) === # 支援 JSON, CSV, TXT 並統一轉換為 list of dicts if batch_uploaded_file: batch_file_key = f"batch_{batch_uploaded_file.name}_{batch_uploaded_file.size}" if st.session_state.batch_current_file_key != batch_file_key or 'json_data_for_batch' not in st.session_state: try: stringio = io.StringIO(batch_uploaded_file.getvalue().decode("utf-8")) parsed_data = None # --- Case 1: JSON --- if batch_uploaded_file.name.lower().endswith('.json'): parsed_data = json.load(stringio) st.toast("JSON 檔案載入成功", icon="📄") # --- Case 2: CSV --- elif batch_uploaded_file.name.lower().endswith('.csv'): # 使用 DictReader 將 CSV 轉為 List of Dicts reader = csv.DictReader(stringio) parsed_data = list(reader) st.toast("CSV 檔案已轉換為 JSON 結構", icon="📊") # --- Case 3: TXT --- else: # 預設為 TXT # 將每一行包裝成一個 JSON 物件: {"raw_content": "line text"} lines = stringio.readlines() parsed_data = [{"raw_log_entry": line.strip()} for line in lines if line.strip()] st.toast("TXT 檔案已轉換為 JSON 結構", icon="📝") # 儲存處理後的數據 st.session_state.json_data_for_batch = parsed_data st.session_state.batch_current_file_key = batch_file_key except Exception as e: st.error(f"檔案解析錯誤: {e}") if 'json_data_for_batch' in st.session_state: del st.session_state.json_data_for_batch elif 'json_data_for_batch' in st.session_state: del st.session_state.json_data_for_batch del st.session_state.batch_current_file_key if "batch_results" in st.session_state: del st.session_state.batch_results st.info("批量分析檔案已移除，已清除相關數據。") # === 執行批量分析邏輯 === if st.session_state.execute_batch_analysis and 'json_data_for_batch' in st.session_state: st.session_state.execute_batch_analysis = False start_time = time.time() st.session_state.batch_results = [] if inference_client is None: st.error("Client 未連線，無法執行。") else: data_to_process = st.session_state.json_data_for_batch logs_list = [] # 處理不同的 JSON 結構 (Dict vs List) if isinstance(data_to_process, list): logs_list = data_to_process elif isinstance(data_to_process, dict): # 嘗試尋找常見的 key if 'alerts' in data_to_process and isinstance(data_to_process['alerts'], list): logs_list = data_to_process['alerts'] elif 'logs' in data_to_process and isinstance(data_to_process['logs'], list): logs_list = data_to_process['logs'] else: logs_list = [data_to_process] else: logs_list = [data_to_process] if logs_list: vs = st.session_state.get("vector_store", None) # --- 關鍵：在這裡做 JSON String 的轉換 --- # 無論來源是 CSV(Dict) 還是 TXT(Dict)，都在這裡用 json.dumps 轉成字串 # 這保證了 Prompt 收到的永遠是 JSON 格式的文字 formatted_logs = [json.dumps(log, indent=2, ensure_ascii=False) for log in logs_list] analysis_sequences = [] for i in range(len(formatted_logs)): start_index = max(0, i - WINDOW_SIZE + 1) end_index = i + 1 current_window = formatted_logs[start_index:end_index] sequence_text = [] for j, log_str in enumerate(current_window): is_target = " <<< TARGET LOG TO ANALYZE" if j == len(current_window) - 1 else "" sequence_text.append(f"--- Log Index {i - len(current_window) + j + 1} ({len(current_window)-j} prior logs){is_target} ---\n{log_str}") analysis_sequences.append({ "sequence_text": "\n\n".join(sequence_text), "target_log_id": i + 1, "original_log_entry": logs_list[i] }) total_sequences = len(analysis_sequences) st.header(f"⚡ 批量分析執行中 (平移視窗 $N={WINDOW_SIZE}$)...") progress_bar = st.progress(0, text=f"準備處理 {total_sequences} 個序列...") results_container = st.container() full_report_chunks = ["## Cybersecurity Batch Analysis Report\n\n"] for i, seq_data in enumerate(analysis_sequences): log_id = seq_data["target_log_id"] progress_bar.progress((i + 1) / total_sequences, text=f"Processing {i + 1}/{total_sequences} (Log #{log_id})...") try: response, retrieved_ctx = generate_rag_response_hf_for_log( client=inference_client, model_id=MODEL_ID, log_sequence_text=seq_data["sequence_text"], user_prompt=analysis_prompt, sys_prompt=system_prompt, vector_store=vs, threshold=similarity_threshold, max_output_tokens=max_output_tokens, temperature=temperature, top_p=top_p ) item = { "log_id": log_id, "log_content": seq_data["original_log_entry"], "sequence_analyzed": seq_data["sequence_text"], "analysis_result": response, "context": retrieved_ctx } st.session_state.batch_results.append(item) with results_container: st.subheader(f"Log/Alert #{item['log_id']}") with st.expander("序列內容 (JSON Format)"): st.code(item["sequence_analyzed"], language='json') # 這裡顯示的會是 JSON 格式 is_high = any(x in response.lower() for x in ['high risk']) if is_high: st.error(item['analysis_result']) else: st.info(item['analysis_result']) if item['context']: with st.expander("參考 RAG 片段"): st.code(item['context']) st.markdown("---") log_content_str_for_report = json.dumps(item["log_content"], indent=2, ensure_ascii=False).replace("`", "\\`") full_report_chunks.append(f"---\n\n### Log #{item['log_id']}\n```json\n{log_content_str_for_report}\n```\nResult:\n{item['analysis_result']}\n") except Exception as e: st.error(f"Error Log {log_id}: {e}") end_time = time.time() progress_bar.empty() st.success(f"完成！耗時 {end_time - start_time:.2f} 秒。") else: st.error("無法提取有效 Log，請檢查檔案格式。") # === 顯示結果 (歷史紀錄) === if st.session_state.get("batch_results") and not st.session_state.execute_batch_analysis: st.header("⚡ 歷史分析結果") full_report_chunks = ["## Report\n\n"] for item in st.session_state.batch_results: log_content_str_for_report = json.dumps(item["log_content"], indent=2, ensure_ascii=False).replace("`", "\\`") full_report_chunks.append(f"---\n\n### Log #{item['log_id']}\n```json\n{log_content_str_for_report}\n```\n{item['analysis_result']}\n") st.download_button("📥 下載完整報告 (.md)", "\n".join(full_report_chunks), "report.md", "text/markdown")