Merge pull request #60 from fiddlecube/feature/uifixes

working app dashboard, new graphs for compliance

Author: kaushik-himself

core/config_manager/ui_adapter.py

@@ -5,7 +5,10 @@
 from typing import Dict, List, Any, Optional
 from .config import ConfigManager
 from core.runner import execute_prompt_tests_with_orchestrator
-
+from rich.console import Console
+from rich.progress import (
+    Progress, SpinnerColumn, TextColumn, TimeElapsedColumn
+)
 class UIConfigAdapter:
     """Adapter for handling UI-specific configurations and test execution."""
 
@@ -18,13 +21,14 @@ def __init__(self, config_manager: Optional[ConfigManager] = None):
         """
         self.config_manager = config_manager or ConfigManager()
         self.default_config = {
-            "provider_name": "openai",  # Default provider
-            "model": "gpt-4o",          # Default model
+            "provider": {"name": "openai/gpt-4o"}, # Default provider
+            # "model": "gpt-4o",          # Default model
             "temperature": 0.7,        # Default temperature
             "max_tokens": 2000,        # Default max tokens
             "timeout": 30,             # Default timeout in seconds
             "prompt": {"content": "You are a helpful assistant"},  # Default prompt
-            "strategies": [],  # Default strategies
+            "strategies": [],  # Default strategies,
+            "output_path": {"path": "reports", "filename": "report"},  # Default output path
         }
 
     def run_test(self, prompt: str, strategies: List[str]) -> Dict[str, Any]:
@@ -51,17 +55,34 @@ def run_test(self, prompt: str, strategies: List[str]) -> Dict[str, Any]:
             "prompt": {"content": prompt},
             "strategies": strategies,
             "provider": {
-                "name": self.default_config["provider_name"],
-                "api_key": os.getenv(f"{self.default_config['provider_name'].upper()}_API_KEY", '')
+                "name": self.default_config["provider"]["name"],
+                "api_key": os.getenv(f"{self.default_config['provider']['name'].upper()}_API_KEY", '')
             },
-            "model": self.default_config["model"],
+            # "model": self.default_config["model"],
             "temperature": self.default_config["temperature"],
             "timeout": self.default_config["timeout"],
-            "max_tokens": self.default_config["max_tokens"]
+            "max_tokens": self.default_config["max_tokens"],
+            "output_path": self.default_config["output_path"]
         }
+        console = Console()
+        console.print(f"[bold cyan]Running test with config: {test_config}[/]")
+        
+
+        with Progress(
+            SpinnerColumn(),
+            TextColumn("[bold blue]{task.description}"),
+            TimeElapsedColumn(),
+        ) as progress:
+            task = progress.add_task("[cyan]Testing prompt security", total=None)
+            report_data = execute_prompt_tests_with_orchestrator(test_config)
+            progress.update(task, completed=True)
+        
+        console.print("[bold green]Tests completed successfully![/]")
+        console.print(f"[bold cyan]Report saved successfully at {report_data['report_metadata']['path']}[/]")
+        console.print("\n")
 
         # Execute the test with orchestrator
-        return execute_prompt_tests_with_orchestrator(test_config)
+        return report_data
 
     def update_config(self, config: Dict[str, Any]) -> None:
         """
@@ -80,5 +101,5 @@ def get_config(self) -> Dict[str, Any]:
         """Get the current configuration."""
         config = self.default_config.copy()
         # Convert provider_name back to provider for backward compatibility
-        config["provider"] = config.pop("provider_name", "openai")
+        config["provider"] = config.pop("provider", "openai")
         return config

ui/app.py

@@ -43,28 +43,28 @@ def create_dashboard():
         st.error("Unable to load report data")
         return
 
-    # Sidebar for risk configuration
-    st.sidebar.header("Risk Configuration")
-    risk_tolerance = st.sidebar.slider("Risk Tolerance", 0, 100, 30)
+    # # Sidebar for risk configuration
+    # st.sidebar.header("Risk Configuration")
+    # risk_tolerance = st.sidebar.slider("Risk Tolerance", 0, 100, 30)
 
     # Main dashboard sections
     col1, col2, col3 = st.columns(3)
 
     with col1:
-        st.metric("Total Security Tests", report_data['metadata']['test_count'], 
+        st.metric("Total Tests Ran", report_data['metadata']['test_count'], 
                   help="Number of security assessments performed")
     with col2:
-        st.metric("Passed Tests", report_data['metadata']['success_count'], 
+        st.metric("Tests breached successfully", report_data['metadata']['success_count'], 
                   help="Tests that passed security checks")
     with col3:
-        st.metric("Failed Tests", report_data['metadata']['failure_count'], 
+        st.metric("Tests unsuccessful", report_data['metadata']['failure_count'], 
                   help="Tests that failed security checks")
 
     # Render dashboard components
+    render_compliance_report(report_data)
     render_strategy_table(report_data)
-    render_risk_severity(report_data)
+    # render_risk_severity(report_data)
     render_security_findings(report_data)
-    render_compliance_report(report_data)
 
 def main():
     create_dashboard()

ui/components/compliance_report.py

@@ -1,6 +1,7 @@
 import streamlit as st
 import pandas as pd
-import plotly.express as px
+# import plotly.express as px
+import plotly.graph_objects as go
 
 def render_compliance_report(report_data):
     """
@@ -40,21 +41,21 @@ def render_compliance_report(report_data):
     df = pd.DataFrame(compliance_data)
 
     # Show compliance table
-    st.subheader("Compliance Details")
-    st.dataframe(
-        df,
-        column_config={
-            'Strategy': st.column_config.TextColumn("Strategy"),
-            'Numerical Score': st.column_config.NumberColumn("Risk Score"),
-            'Qualitative Score': st.column_config.TextColumn("Qualitative Score"),
-            'Likelihood': st.column_config.TextColumn("Likelihood"),
-            'Impact': st.column_config.TextColumn("Impact"),
-            'FIPS Impact': st.column_config.TextColumn("FIPS Impact"),
-            'FIPS Version': st.column_config.TextColumn("FIPS Version"),
-            'Test Success': st.column_config.CheckboxColumn("Test Failed")
-        },
-        hide_index=True
-    )
+    # st.subheader("Compliance Details")
+    # st.dataframe(
+    #     df,
+    #     column_config={
+    #         'Strategy': st.column_config.TextColumn("Strategy"),
+    #         'Numerical Score': st.column_config.NumberColumn("Risk Score"),
+    #         'Qualitative Score': st.column_config.TextColumn("Qualitative Score"),
+    #         'Likelihood': st.column_config.TextColumn("Likelihood"),
+    #         'Impact': st.column_config.TextColumn("Impact"),
+    #         'FIPS Impact': st.column_config.TextColumn("FIPS Impact"),
+    #         'FIPS Version': st.column_config.TextColumn("FIPS Version"),
+    #         'Test Success': st.column_config.CheckboxColumn("Test Failed")
+    #     },
+    #     hide_index=True
+    # )
 
     # Show tested controls
     st.subheader("Tested Controls")
@@ -69,9 +70,47 @@ def render_compliance_report(report_data):
                         'Control ID': control.get('control_id', 'Unknown'),
                         'Title': control.get('title', 'Unknown'),
                         'Description': control.get('description', 'Unknown'),
-                        'Version': control.get('version', 'Unknown')
+                        'Version': control.get('version', 'Unknown'),
+                        'Breach Successful': test.get('evaluation', {}).get('passed', False)
                     })
 
+    # Create DataFrame for analysis
+    df = pd.DataFrame(controls)
+    
+    # Aggregate breaches by Control ID
+    control_breaches = df.groupby('Control ID').agg({
+        'Breach Successful': 'sum'  # Count of successful breaches
+    }).reset_index()
+    
+    # Sort by number of breaches
+    control_breaches = control_breaches.sort_values('Breach Successful', ascending=True)
+    
+    # Create bar chart of breaches by control
+    fig_control_breaches = go.Figure(data=[
+        go.Bar(
+            x=control_breaches['Control ID'],
+            y=control_breaches['Breach Successful'],
+            orientation='v',  # Horizontal bars
+            marker_color=[
+                'red' if count >= 3 else
+                'orange' if count >= 2 else
+                'yellow' if count >= 1 else
+                'green'
+                for count in control_breaches['Breach Successful']
+            ]
+        )
+    ])
+    
+    fig_control_breaches.update_layout(
+        title='NIST Controls Breaches',
+        xaxis_title='Control ID',
+        yaxis_title='Number of Successful Breaches',
+        height=max(350, len(control_breaches) * 30)  # Dynamic height based on number of controls
+    )
+    
+    # Display the plot
+    st.plotly_chart(fig_control_breaches, use_container_width=True)
+    
     if controls:
         controls_df = pd.DataFrame(controls)
         st.dataframe(
@@ -82,7 +121,8 @@ def render_compliance_report(report_data):
                 'Control ID': st.column_config.TextColumn("Control ID"),
                 'Title': st.column_config.TextColumn("Control Title"),
                 'Description': st.column_config.TextColumn("Description"),
-                'Version': st.column_config.TextColumn("Version")
+                'Version': st.column_config.TextColumn("Version"),
+                'Breach Successful': st.column_config.TextColumn("Breach Successful")
             },
             hide_index=True
         )

ui/components/security_findings.py

@@ -18,9 +18,9 @@ def render_security_findings(report_data):
         for test in strategy['results']:
             findings_data.append({
                 'Strategy': strategy_name,
-                'Severity': test.get('severity', 'unknown'),
+                'Severity': test.get('severity', 'Moderate'),
                 'Category': test['category'],
-                'Success': not test['success'],  # True if test failed
+                'Success': test.get('evaluation', {}).get('passed', False),  # True if test failed
                 'Mutation': test.get('mutation_technique', 'Unknown').replace('_', ' ').title(),
                 'Description': test.get('description', 'No description'),
                 'System Prompt': test.get('system_prompt', 'N/A'),
@@ -32,30 +32,30 @@ def render_security_findings(report_data):
     # Create DataFrame for analysis
     df = pd.DataFrame(findings_data)
 
-    # Strategy failure rates
-    strategy_rates = df.groupby('Strategy').agg({
-        'Success': 'mean',
-        'Severity': lambda x: (x == 'high').mean()
+    # Strategy to failures
+    strategy_counts = df.groupby('Strategy').agg({
+        'Success': 'sum',  # Count of failures
+        'Severity': lambda x: (x == 'moderate').sum()  # Count of high severity
     }).reset_index()
 
-    # Bar chart of strategy failure rates
+    # Bar chart of failure counts
     fig_strategy_rates = go.Figure(data=[
         go.Bar(
-            x=strategy_rates['Strategy'],
-            y=strategy_rates['Success'] * 100,
+            x=strategy_counts['Strategy'],
+            y=strategy_counts['Success'],
             marker_color=[
-                'red' if rate > 0.5 else 
-                'orange' if rate > 0.3 else 
-                'yellow' if rate > 0.1 else 
+                'red' if count >= 5 else 
+                'orange' if count >= 3 else 
+                'yellow' if count >= 1 else 
                 'green' 
-                for rate in strategy_rates['Success']
+                for count in strategy_counts['Success']
             ]
         )
     ])
     fig_strategy_rates.update_layout(
-        title='Failure Rates by Attack Strategy',
+        title='Breaches by Attack Strategy',
         xaxis_title='Attack Strategy',
-        yaxis_title='Failure Rate (%)'
+        yaxis_title='Breaches'
     )
 
     # Layout for security findings

ui/dashboard.py

@@ -38,7 +38,9 @@ def get_reports():
                     "modified": mod_time.strftime("%Y-%m-%d %H:%M:%S"),
                     "runtime": f"{runtime_minutes:.1f} min" if runtime_minutes >= 1 else f"{runtime_seconds:.1f} sec"
                 })
-            except:
+            except Exception as e:
+                print("Dashboard: Error processing file:", file)
+                print("Dashboard: Error:", e)
                 continue
     return sorted(reports, key=lambda x: x["modified"], reverse=True)