minhan6559/Log-Analysis-MultiAgent
2
1#!/usr/bin/env python3
2"""
3Streamlit Web App for Cybersecurity Agent Pipeline
4
5A simple web interface for uploading log files and running the cybersecurity analysis pipeline
6with different LLM models.
7"""
8
9import os
10import sys
11import tempfile
12import shutil
13import time
14import streamlit as st
15from pathlib import Path
16from typing import Dict, Any, Optional
17
18# Add project root to path for agent imports
19project_root = Path(__file__).parent
20sys.path.insert(0, str(project_root))
21
22from src.full_pipeline.simple_pipeline import analyze_log_file
23
24from dotenv import load_dotenv
25from huggingface_hub import login as huggingface_login
26from huggingface_hub.utils import HfHubHTTPError
27
28load_dotenv()
29
30
31def get_model_providers() -> Dict[str, Dict[str, str]]:
32 """Get available model providers and their models."""
33 return {
34 "Google GenAI": {
35 "gemini-2.0-flash": "google_genai:gemini-2.0-flash",
36 "gemini-2.0-flash-lite": "google_genai:gemini-2.0-flash-lite",
37 "gemini-2.5-flash-lite": "google_genai:gemini-2.5-flash-lite",
38 },
39 "Groq": {
40 "openai/gpt-oss-120b": "groq:openai/gpt-oss-120b",
41 "moonshotai/kimi-k2-instruct-0905": "groq:moonshotai/kimi-k2-instruct-0905",
42 },
43 "OpenAI": {
44 "gpt-5-mini": "openai:gpt-5-mini",
45 "gpt-5": "openai:gpt-5",
46 "gpt-4.1-mini": "openai:gpt-4.1-mini",
47 },
48 }
49
50
51def get_api_key_help() -> Dict[str, str]:
52 """Get API key help information for each provider."""
53 return {
54 "Google GenAI": "https://aistudio.google.com/app/apikey",
55 "Groq": "https://console.groq.com/keys",
56 "OpenAI": "https://platform.openai.com/api-keys",
57 }
58
59
60def setup_temp_directories(temp_dir: str) -> Dict[str, str]:
61 """Setup temporary directories for the pipeline."""
62 log_files_dir = os.path.join(temp_dir, "log_files")
63 analysis_dir = os.path.join(temp_dir, "analysis")
64 final_response_dir = os.path.join(temp_dir, "final_response")
65
66 os.makedirs(log_files_dir, exist_ok=True)
67 os.makedirs(analysis_dir, exist_ok=True)
68 os.makedirs(final_response_dir, exist_ok=True)
69
70 return {
71 "log_files": log_files_dir,
72 "analysis": analysis_dir,
73 "final_response": final_response_dir,
74 }
75
76
77def save_uploaded_file(uploaded_file, temp_dir: str) -> str:
78 """Save uploaded file to temporary directory."""
79 log_files_dir = os.path.join(temp_dir, "log_files")
80 file_path = os.path.join(log_files_dir, uploaded_file.name)
81
82 with open(file_path, "wb") as f:
83 f.write(uploaded_file.getbuffer())
84
85 return file_path
86
87
88def run_analysis(
89 log_file_path: str,
90 model_name: str,
91 query: str,
92 temp_dirs: Dict[str, str],
93 api_key: str,
94 provider: str,
95 max_log_analysis_iterations: int,
96 max_retrieval_iterations: int,
97 progress_callback=None,
98) -> Dict[str, Any]:
99 """Run the cybersecurity analysis pipeline."""
100
101 # Set environment variable for API key
102 if provider == "Google GenAI":
103 os.environ["GOOGLE_API_KEY"] = api_key
104 elif provider == "Groq":
105 os.environ["GROQ_API_KEY"] = api_key
106 elif provider == "OpenAI":
107 os.environ["OPENAI_API_KEY"] = api_key
108
109 try:
110 # Run the analysis pipeline
111 result = analyze_log_file(
112 log_file=log_file_path,
113 query=query,
114 tactic=None,
115 model_name=model_name,
116 temperature=0.1,
117 max_log_analysis_iterations=max_log_analysis_iterations,
118 max_retrieval_iterations=max_retrieval_iterations,
119 log_agent_output_dir=temp_dirs["analysis"],
120 response_agent_output_dir=temp_dirs["final_response"],
121 progress_callback=progress_callback,
122 )
123 return {"success": True, "result": result}
124 except Exception as e:
125 return {"success": False, "error": str(e)}
126
127
128@st.cache_resource
129def initialize_hf_login():
130 """Initialize Hugging Face login only once."""
131 hf_token = os.getenv("HF_TOKEN")
132 if hf_token:
133 try:
134 # Check if already logged in by trying to get user info
135 from huggingface_hub import whoami
136
137 whoami()
138 return True
139 except (HfHubHTTPError, Exception):
140 # Not logged in, try to login
141 try:
142 huggingface_login(token=hf_token)
143 return True
144 except Exception as e:
145 st.warning(f"Failed to login to Hugging Face: {e}")
146 return False
147 return False
148
149
150def main():
151 """Main Streamlit app."""
152
153 # Initialize HF login (cached)
154 initialize_hf_login()
155
156 st.set_page_config(
157 page_title="Cybersecurity Agent Pipeline", page_icon="🛡️", layout="wide"
158 )
159
160 st.title("Cybersecurity Agent Pipeline")
161 st.markdown(
162 "Upload a log file and analyze it using advanced LLM-based cybersecurity agents."
163 )
164
165 # Sidebar for configuration
166 with st.sidebar:
167 st.header("Configuration")
168
169 # Model selection
170 providers = get_model_providers()
171 selected_provider = st.selectbox(
172 "Select Model Provider", list(providers.keys())
173 )
174
175 available_models = providers[selected_provider]
176 selected_model_display = st.selectbox(
177 "Select Model", list(available_models.keys())
178 )
179 selected_model = available_models[selected_model_display]
180
181 # API Key input with help
182 st.subheader("API Key")
183 api_key_help = get_api_key_help()
184
185 with st.expander("How to get API key", expanded=False):
186 st.markdown(f"**{selected_provider}**:")
187 st.markdown(f"[Get API Key]({api_key_help[selected_provider]})")
188
189 api_key = st.text_input(
190 f"Enter {selected_provider} API Key",
191 type="password",
192 help=f"Your {selected_provider} API key",
193 )
194
195 # Main content area
196 col1, col2 = st.columns([2, 1])
197
198 with col1:
199 st.header("Upload Log File")
200 uploaded_file = st.file_uploader(
201 "Choose a JSON log file",
202 type=["json"],
203 help="Upload a JSON log file from the Mordor dataset or similar security logs",
204 )
205
206 with col2:
207 st.header("Analysis Status")
208 if uploaded_file is not None:
209 st.success(f"File uploaded: {uploaded_file.name}")
210 st.info(f"Size: {uploaded_file.size:,} bytes")
211 else:
212 st.warning("Please upload a log file")
213
214 # Run analysis button
215 if st.button(
216 "Run Analysis", type="primary", disabled=not (uploaded_file and api_key)
217 ):
218 if not uploaded_file:
219 st.error("Please upload a log file first.")
220 return
221
222 if not api_key:
223 st.error("Please enter your API key.")
224 return
225
226 # Create temporary directory
227 temp_dir = tempfile.mkdtemp(prefix="cyber_agent_")
228
229 try:
230 # Setup directories
231 temp_dirs = setup_temp_directories(temp_dir)
232
233 # Save uploaded file
234 log_file_path = save_uploaded_file(uploaded_file, temp_dir)
235
236 # Show progress
237 progress_bar = st.progress(0)
238 status_text = st.empty()
239
240 status_text.text("Initializing analysis...")
241 progress_bar.progress(10)
242
243 # Start timing
244 start_time = time.time()
245
246 # Create progress callback
247 def update_progress(progress: int, message: str):
248 progress_bar.progress(progress)
249 status_text.text(message)
250
251 # Run analysis
252 analysis_result = run_analysis(
253 log_file_path=log_file_path,
254 model_name=selected_model,
255 query="",
256 temp_dirs=temp_dirs,
257 api_key=api_key,
258 provider=selected_provider,
259 max_log_analysis_iterations=2,
260 max_retrieval_iterations=2,
261 progress_callback=update_progress,
262 )
263
264 # Calculate execution time
265 end_time = time.time()
266 execution_time = end_time - start_time
267
268 progress_bar.progress(90)
269 status_text.text("Finalizing results...")
270
271 if analysis_result["success"]:
272 progress_bar.progress(100)
273 status_text.text("Analysis completed successfully!")
274
275 # Display results
276 st.header("Analysis Results")
277
278 result = analysis_result["result"]
279
280 # Show key metrics
281 col1, col2, col3 = st.columns(3)
282
283 with col1:
284 assessment = result.get("log_analysis_result", {}).get(
285 "overall_assessment", "Unknown"
286 )
287 st.metric("Overall Assessment", assessment)
288
289 with col2:
290 abnormal_events = result.get("log_analysis_result", {}).get(
291 "abnormal_events", []
292 )
293 st.metric("Abnormal Events", len(abnormal_events))
294
295 with col3:
296 st.metric("Execution Time", f"{execution_time:.2f}s")
297
298 # Show markdown report
299 markdown_report = result.get("markdown_report", "")
300 if markdown_report:
301 st.header("Detailed Report")
302 st.markdown(markdown_report)
303 else:
304 st.warning("No detailed report generated.")
305
306 else:
307 st.error(f"Analysis failed: {analysis_result['error']}")
308 st.exception(analysis_result["error"])
309
310 finally:
311 # Cleanup temporary directory
312 try:
313 shutil.rmtree(temp_dir)
314 except Exception as e:
315 st.warning(f"Could not clean up temporary directory: {e}")
316
317 # Footer
318 st.markdown("---")
319 st.markdown(
320 "**Cybersecurity Agent Pipeline** - Powered by LangGraph and LangChain | "
321 "Built for educational purposes demonstrating LLM-based multi-agent systems"
322 )
323
324
325if __name__ == "__main__":
326 main()
327