#!/usr/bin/env python3 """ Streamlit Web App for Cybersecurity Agent Pipeline A simple web interface for uploading log files and running the cybersecurity analysis pipeline with different LLM models. """ import os import sys import tempfile import shutil import time import streamlit as st from pathlib import Path from typing import Dict, Any, Optional # Add project root to path for agent imports project_root = Path(__file__).parent sys.path.insert(0, str(project_root)) from src.full_pipeline.simple_pipeline import analyze_log_file from dotenv import load_dotenv from huggingface_hub import login as huggingface_login from huggingface_hub.utils import HfHubHTTPError load_dotenv() def get_model_providers() -> Dict[str, Dict[str, str]]: """Get available model providers and their models.""" return { "Google GenAI": { "gemini-2.0-flash": "google_genai:gemini-2.0-flash", "gemini-2.0-flash-lite": "google_genai:gemini-2.0-flash-lite", "gemini-2.5-flash-lite": "google_genai:gemini-2.5-flash-lite", }, "Groq": { "openai/gpt-oss-120b": "groq:openai/gpt-oss-120b", "moonshotai/kimi-k2-instruct-0905": "groq:moonshotai/kimi-k2-instruct-0905", }, "OpenAI": { "gpt-5-mini": "openai:gpt-5-mini", "gpt-5": "openai:gpt-5", "gpt-4.1-mini": "openai:gpt-4.1-mini", }, } def get_api_key_help() -> Dict[str, str]: """Get API key help information for each provider.""" return { "Google GenAI": "https://aistudio.google.com/app/apikey", "Groq": "https://console.groq.com/keys", "OpenAI": "https://platform.openai.com/api-keys", } def setup_temp_directories(temp_dir: str) -> Dict[str, str]: """Setup temporary directories for the pipeline.""" log_files_dir = os.path.join(temp_dir, "log_files") analysis_dir = os.path.join(temp_dir, "analysis") final_response_dir = os.path.join(temp_dir, "final_response") os.makedirs(log_files_dir, exist_ok=True) os.makedirs(analysis_dir, exist_ok=True) os.makedirs(final_response_dir, exist_ok=True) return { "log_files": log_files_dir, "analysis": analysis_dir, "final_response": final_response_dir, } def save_uploaded_file(uploaded_file, temp_dir: str) -> str: """Save uploaded file to temporary directory.""" log_files_dir = os.path.join(temp_dir, "log_files") file_path = os.path.join(log_files_dir, uploaded_file.name) with open(file_path, "wb") as f: f.write(uploaded_file.getbuffer()) return file_path def run_analysis( log_file_path: str, model_name: str, query: str, temp_dirs: Dict[str, str], api_key: str, provider: str, max_log_analysis_iterations: int, max_retrieval_iterations: int, progress_callback=None, ) -> Dict[str, Any]: """Run the cybersecurity analysis pipeline.""" # Set environment variable for API key if provider == "Google GenAI": os.environ["GOOGLE_API_KEY"] = api_key elif provider == "Groq": os.environ["GROQ_API_KEY"] = api_key elif provider == "OpenAI": os.environ["OPENAI_API_KEY"] = api_key try: # Run the analysis pipeline result = analyze_log_file( log_file=log_file_path, query=query, tactic=None, model_name=model_name, temperature=0.1, max_log_analysis_iterations=max_log_analysis_iterations, max_retrieval_iterations=max_retrieval_iterations, log_agent_output_dir=temp_dirs["analysis"], response_agent_output_dir=temp_dirs["final_response"], progress_callback=progress_callback, ) return {"success": True, "result": result} except Exception as e: return {"success": False, "error": str(e)} @st.cache_resource def initialize_hf_login(): """Initialize Hugging Face login only once.""" hf_token = os.getenv("HF_TOKEN") if hf_token: try: # Check if already logged in by trying to get user info from huggingface_hub import whoami whoami() return True except (HfHubHTTPError, Exception): # Not logged in, try to login try: huggingface_login(token=hf_token) return True except Exception as e: st.warning(f"Failed to login to Hugging Face: {e}") return False return False def main(): """Main Streamlit app.""" # Initialize HF login (cached) initialize_hf_login() st.set_page_config( page_title="Cybersecurity Agent Pipeline", page_icon="🛡️", layout="wide" ) st.title("Cybersecurity Agent Pipeline") st.markdown( "Upload a log file and analyze it using advanced LLM-based cybersecurity agents." ) # Sidebar for configuration with st.sidebar: st.header("Configuration") # Model selection providers = get_model_providers() selected_provider = st.selectbox( "Select Model Provider", list(providers.keys()) ) available_models = providers[selected_provider] selected_model_display = st.selectbox( "Select Model", list(available_models.keys()) ) selected_model = available_models[selected_model_display] # API Key input with help st.subheader("API Key") api_key_help = get_api_key_help() with st.expander("How to get API key", expanded=False): st.markdown(f"**{selected_provider}**:") st.markdown(f"[Get API Key]({api_key_help[selected_provider]})") api_key = st.text_input( f"Enter {selected_provider} API Key", type="password", help=f"Your {selected_provider} API key", ) # Main content area col1, col2 = st.columns([2, 1]) with col1: st.header("Upload Log File") uploaded_file = st.file_uploader( "Choose a JSON log file", type=["json"], help="Upload a JSON log file from the Mordor dataset or similar security logs", ) with col2: st.header("Analysis Status") if uploaded_file is not None: st.success(f"File uploaded: {uploaded_file.name}") st.info(f"Size: {uploaded_file.size:,} bytes") else: st.warning("Please upload a log file") # Run analysis button if st.button( "Run Analysis", type="primary", disabled=not (uploaded_file and api_key) ): if not uploaded_file: st.error("Please upload a log file first.") return if not api_key: st.error("Please enter your API key.") return # Create temporary directory temp_dir = tempfile.mkdtemp(prefix="cyber_agent_") try: # Setup directories temp_dirs = setup_temp_directories(temp_dir) # Save uploaded file log_file_path = save_uploaded_file(uploaded_file, temp_dir) # Show progress progress_bar = st.progress(0) status_text = st.empty() status_text.text("Initializing analysis...") progress_bar.progress(10) # Start timing start_time = time.time() # Create progress callback def update_progress(progress: int, message: str): progress_bar.progress(progress) status_text.text(message) # Run analysis analysis_result = run_analysis( log_file_path=log_file_path, model_name=selected_model, query="", temp_dirs=temp_dirs, api_key=api_key, provider=selected_provider, max_log_analysis_iterations=2, max_retrieval_iterations=2, progress_callback=update_progress, ) # Calculate execution time end_time = time.time() execution_time = end_time - start_time progress_bar.progress(90) status_text.text("Finalizing results...") if analysis_result["success"]: progress_bar.progress(100) status_text.text("Analysis completed successfully!") # Display results st.header("Analysis Results") result = analysis_result["result"] # Show key metrics col1, col2, col3 = st.columns(3) with col1: assessment = result.get("log_analysis_result", {}).get( "overall_assessment", "Unknown" ) st.metric("Overall Assessment", assessment) with col2: abnormal_events = result.get("log_analysis_result", {}).get( "abnormal_events", [] ) st.metric("Abnormal Events", len(abnormal_events)) with col3: st.metric("Execution Time", f"{execution_time:.2f}s") # Show markdown report markdown_report = result.get("markdown_report", "") if markdown_report: st.header("Detailed Report") st.markdown(markdown_report) else: st.warning("No detailed report generated.") else: st.error(f"Analysis failed: {analysis_result['error']}") st.exception(analysis_result["error"]) finally: # Cleanup temporary directory try: shutil.rmtree(temp_dir) except Exception as e: st.warning(f"Could not clean up temporary directory: {e}") # Footer st.markdown("---") st.markdown( "**Cybersecurity Agent Pipeline** - Powered by LangGraph and LangChain | " "Built for educational purposes demonstrating LLM-based multi-agent systems" ) if __name__ == "__main__": main()