Update app.py

app.py

@@ -4,217 +4,123 @@ from pathlib import Path
 from processor import DataProcessor
 from llm_handler import LLMHandler
 
-class DDoSAnalyzerUI:
+class DDoSResistanceHelper:
     def __init__(self):
-        # Set up Streamlit page configuration
+        # Configure Streamlit app
         st.set_page_config(
-            page_title="DDoS Traffic Analyzer",
+            page_title="DDoS Resistance Helper URF LLM Network Analyzer",
             page_icon=":shield:",
             layout="wide",
             initial_sidebar_state="expanded"
         )
 
         # Initialize session state
-        self.setup_session_state()
+        self.initialize_session_state()
 
         # Initialize processor and LLM handler
         self.processor = DataProcessor()
         self.llm_handler = LLMHandler()
 
-        # Define themes
-        self.themes = {
-            'Light': {
-                'base': 'light',
-                'primaryColor': '#3B82F6',
-                'backgroundColor': '#FFFFFF',
-                'secondaryBackgroundColor': '#F3F4F6',
-                'textColor': '#000000'
-            },
-            'Dark': {
-                'base': 'dark',
-                'primaryColor': '#10B981',
-                'backgroundColor': '#1F2937',
-                'secondaryBackgroundColor': '#374151',
-                'textColor': '#FFFFFF'
-            },
-            'Neon': {
-                'base': 'dark',
-                'primaryColor': '#22D3EE',
-                'backgroundColor': '#000000',
-                'secondaryBackgroundColor': '#1F2937',
-                'textColor': '#22D3EE'
-            }
-        }
-
-    def setup_session_state(self):
-        """Initialize Streamlit session state variables."""
-        st.write("Setting up session state.")
+    def initialize_session_state(self):
+        """Set up Streamlit session state variables."""
         session_keys = [
-            'chat_history',
-            'current_file',
-            'analysis_complete',
-            'selected_theme',
-            'preview_data'
+            'current_file', 'preprocessed_data', 'analysis_results', 'chat_history'
         ]
-        
         for key in session_keys:
             if key not in st.session_state:
                 st.session_state[key] = None if key != 'chat_history' else []
 
-        st.write("Session state setup complete.")
-
-    def apply_theme(self):
-        """Apply selected theme to Streamlit app."""
-        theme_name = st.session_state.get('selected_theme', 'Light')
-        theme = self.themes[theme_name]
-
-        st.markdown(f"""
-        <style>
-        .stApp {{
-            background-color: {theme['backgroundColor']};
-            color: {theme['textColor']};
-        }}
-        .stTabs [data-baseweb="tab-list"] {{
-            background-color: {theme['secondaryBackgroundColor']};
-        }}
-        .stButton>button {{
-            background-color: {theme['primaryColor']};
-            color: white;
-        }}
-        </style>
-        """, unsafe_allow_html=True)
-
-    def render_theme_icons(self):
-        """Render theme selection icons on the right sidebar."""
-        st.sidebar.header("🌈 Theme Selector")
-
-        # Create buttons for theme selection with icons
-        light_icon = "☀️"
-        dark_icon = "🌙"
-        neon_icon = "💡"
-
-        col1, col2, col3 = st.sidebar.columns(3)
+    def render_top_bar(self):
+        """Render the top bar with theme and upload options."""
+        col1, col2 = st.columns([8, 2])
 
         with col1:
-            if st.button(light_icon):
-                st.session_state.selected_theme = 'Light'
-                self.apply_theme()
+            st.title("🛡️ DDoS Resistance Helper URF LLM Network Analyzer")
 
         with col2:
-            if st.button(dark_icon):
-                st.session_state.selected_theme = 'Dark'
-                self.apply_theme()
-
-        with col3:
-            if st.button(neon_icon):
-                st.session_state.selected_theme = 'Neon'
-                self.apply_theme()
+            st.markdown("### Theme Selector")
+            if st.button("Light"):
+                st.markdown("<style>.stApp { background-color: #ffffff; }</style>", unsafe_allow_html=True)
+            elif st.button("Dark"):
+                st.markdown("<style>.stApp { background-color: #1f1f1f; color: white; }</style>", unsafe_allow_html=True)
 
     def render_file_upload(self):
-        """Render file upload section in the sidebar."""
-        st.sidebar.header("📤 Upload Data")
-
-        uploaded_file = st.sidebar.file_uploader(
-            "Choose a CSV file", type=['csv'],
-            help="Upload your network traffic data for DDoS analysis"
-        )
-
+        """Render the file upload component."""
+        uploaded_file = st.file_uploader("Upload Network Traffic Data (CSV)", type=["csv"],
+                                         label_visibility="collapsed")
         if uploaded_file:
             try:
-                # Attempt to read the file
                 df = pd.read_csv(uploaded_file)
-                st.session_state.current_file = uploaded_file
-                st.session_state.preview_data = df.head(5)  # Preview the first 5 rows
-
-                if st.sidebar.button("🚀 Process Data", use_container_width=True):
-                    self.process_uploaded_file(uploaded_file)
-
+                st.session_state.current_file = df
+                st.success("File uploaded successfully!")
             except Exception as e:
-                st.error(f"Error processing file: {e}")
-
-    def process_uploaded_file(self, file):
-        """Process the uploaded file and perform analysis."""
-        try:
-            st.write("Processing uploaded file.")
-            with st.spinner("Processing data..."):
-                # Preprocess data
-                df = self.processor.preprocess_data(file)
-
-                # Calculate statistics
-                stats = self.processor.calculate_statistics(df)
-
-                # Trigger LLM analysis
-                results = self.llm_handler.analyze_data(df)
-
-                # Update session state
-                st.session_state.analysis_complete = True
-                st.success("Analysis completed successfully!")
-                
-        except Exception as e:
-            st.error(f"Error processing file: {e}")
+                st.error(f"Error reading file: {e}")
 
-        st.write("File processing complete.")
-
-    def render_main_content(self):
-        """Render main content area with tabs."""
-        if not st.session_state.current_file:
-            st.info("Please upload a CSV file to begin analysis")
+    def render_analysis(self):
+        """Render the analysis results."""
+        if st.session_state.current_file is None:
+            st.info("Please upload a CSV file to start analysis.")
             return
 
-        # Create tabs
-        tab1, tab2 = st.tabs(["📊 Analysis", "💬 Chat"])
-
-        with tab1:
-            self.render_analysis_tab()
-
-        with tab2:
-            self.render_chat_tab()
-
-    def render_analysis_tab(self):
-        """Render analysis results tab."""
-        if st.session_state.preview_data is not None:
-            st.subheader("📋 Data Preview")
-            st.dataframe(st.session_state.preview_data)
-
-        if st.session_state.analysis_complete:
-            st.subheader("🔍 Analysis Results")
-            # Placeholder for more detailed analysis visualization
-            st.info("Full analysis results would be displayed here.")
+        # Preprocess the data
+        st.subheader("Preprocessing Data")
+        with st.spinner("Preprocessing data..."):
+            try:
+                preprocessed_data = self.processor.preprocess_data(st.session_state.current_file)
+                st.session_state.preprocessed_data = preprocessed_data
+                st.success("Data preprocessed successfully!")
+            except Exception as e:
+                st.error(f"Error during preprocessing: {e}")
 
-    def render_chat_tab(self):
-        """Render interactive chat interface."""
+        # Perform LLM analysis
+        st.subheader("Performing LLM Analysis")
+        with st.spinner("Analyzing data with LLM..."):
+            try:
+                results = self.llm_handler.analyze_data(st.session_state.preprocessed_data)
+                st.session_state.analysis_results = results
+                st.success("Analysis completed successfully!")
+            except Exception as e:
+                st.error(f"Error during LLM analysis: {e}")
+
+        # Show results
+        if st.session_state.analysis_results is not None:
+            st.subheader("Analysis Results")
+            st.dataframe(st.session_state.analysis_results)
+            csv_path = Path("~/.dataset/PROBABILITY_OF_EACH_ROW_DDOS_AND_BENGNIN.csv").expanduser()
+            st.download_button("Download Results as CSV", csv_path.read_bytes(), "analysis_results.csv")
+
+    def render_chat_interface(self):
+        """Render a chat interface for interacting with the LLM."""
+        st.sidebar.header("💬 Chat Interface")
         # Display chat history
         for message in st.session_state.chat_history:
             with st.chat_message(message['role']):
                 st.write(message['content'])
 
         # Chat input
-        if prompt := st.chat_input("Ask about the analysis..."):
+        if prompt := st.sidebar.text_input("Ask about the analysis or mitigation steps..."):
             # Add user message to chat history
             st.session_state.chat_history.append({
-                'role': 'user', 
+                'role': 'user',
                 'content': prompt
             })
 
-            # Get AI response
+            # Get LLM response
             response = self.llm_handler.get_chat_response(prompt)
 
-            # Add AI response to chat history
+            # Add LLM response to chat history
             st.session_state.chat_history.append({
-                'role': 'assistant', 
+                'role': 'assistant',
                 'content': response
             })
 
-    def main(self):
-        """Main method to run the Streamlit application."""
-        self.render_theme_icons()
-
-        # Main content header styling for better UI experience
-        st.title("🛡️ DDoS Traffic Analyzer")
+    def run(self):
+        """Run the Streamlit app."""
+        self.render_top_bar()
         self.render_file_upload()
-        self.render_main_content()
+        self.render_analysis()
+        self.render_chat_interface()
 
-# Run the application
 if __name__ == "__main__":
-    app = DDoSAnalyzerUI()
-    app.main()
+    app = DDoSResistanceHelper()
+    app.run()