Spaces:
Running
Running
| import streamlit as st | |
| import plotly.graph_objects as go | |
| from combined_detector import CombinedCodeDetector | |
| from rule_detector import RuleBasedCodeDetector | |
| from fix_generator import FixSuggestionGenerator | |
| import pandas as pd | |
| import time | |
| import requests | |
| import json | |
| import torch | |
| from transformers import AutoModelForSequenceClassification, AutoTokenizer | |
| # 1. Page Config | |
| st.set_page_config( | |
| page_title="AI Code Security Scanner", | |
| page_icon="🔒", | |
| layout="wide" | |
| ) | |
| # 2. Optimized Component Loading for Deployment | |
| def load_tools(): | |
| # Hugging Face Model ID | |
| model_path = "mubi-613/ai-code-security-scanner" | |
| with st.spinner("🚀 Loading AI Models from Hugging Face... Please wait."): | |
| with torch.inference_mode(): | |
| # Load Tokenizer and Model from HF Hub | |
| tokenizer = AutoTokenizer.from_pretrained(model_path) | |
| model = AutoModelForSequenceClassification.from_pretrained(model_path) | |
| # Initialize detectors (these internal classes should use the loaded model) | |
| detector = CombinedCodeDetector() | |
| fix_gen = FixSuggestionGenerator() | |
| rules = RuleBasedCodeDetector() | |
| # Inject the HF model into the detector if it expects a local one | |
| if hasattr(detector, 'model'): | |
| detector.model = model | |
| detector.model.eval() | |
| if hasattr(detector, 'tokenizer'): | |
| detector.tokenizer = tokenizer | |
| if hasattr(fix_gen, 'model'): | |
| fix_gen.model.eval() | |
| return { | |
| "detector": detector, | |
| "fix_gen": fix_gen, | |
| "rules": rules, | |
| "model": model, | |
| "tokenizer": tokenizer | |
| } | |
| tools = load_tools() | |
| # --- 3. CUSTOM STYLING --- | |
| st.markdown(""" | |
| <style> | |
| .main-header { | |
| text-align: center; | |
| padding: 1.5rem; | |
| background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); | |
| color: white; | |
| border-radius: 10px; | |
| margin-bottom: 2rem; | |
| } | |
| </style> | |
| """, unsafe_allow_html=True) | |
| # --- 4. HEADER --- | |
| st.markdown('<div class="main-header"><h1>🔒 AI Code Security Scanner</h1><p>Real-time vulnerability detection and AI-powered fixes</p></div>', unsafe_allow_html=True) | |
| # 5. Sidebar | |
| with st.sidebar: | |
| st.header("⚙️ Configuration") | |
| analysis_mode = st.radio( | |
| "Analysis Mode", | |
| ["Fast (Rules Only)", "Deep (Rules + AI)", "API Mode"], | |
| help="Deep Mode uses the Hugging Face AI Model" | |
| ) | |
| show_fixes = st.checkbox("Show Fix Suggestions", value=True) | |
| api_url = st.text_input("API URL", "http://localhost:8000") | |
| st.markdown("---") | |
| st.info(""" | |
| **Features:** | |
| - Real-time code analysis | |
| - AI-powered fix suggestions | |
| - Batch processing | |
| - REST API integration | |
| """) | |
| # 4. Main Tabs | |
| tab1, tab2, tab3, tab4 = st.tabs(["🔍 Analyzer", "📁 Batch", "🌐 API", "📊 Dashboard"]) | |
| # --- TAB 1: ANALYZER --- | |
| with tab1: | |
| col1, col2 = st.columns([2, 1]) | |
| with col1: | |
| st.subheader("Code Analysis") | |
| default_code = """def get_user(user_id): | |
| # VULNERABLE: SQL Injection | |
| query = f"SELECT * FROM users WHERE id = {user_id}" | |
| api_key = "sk_live_1234567890abcdef" | |
| return execute_query(query)""" | |
| if 'code_input' not in st.session_state: | |
| st.session_state.code_input = default_code | |
| code_input = st.text_area( | |
| "Paste Python code:", | |
| value=st.session_state.code_input, | |
| height=300, | |
| key="main_editor", | |
| label_visibility="collapsed" | |
| ) | |
| with col2: | |
| st.subheader("Templates") | |
| templates = { | |
| "SQL Injection": 'query = f"SELECT * FROM users WHERE id = {user_input}"', | |
| "Hardcoded Secret": 'api_key = "sk_test_1234567890"', | |
| "XSS Vulnerability": 'return f"<div>{user_input}</div>"', | |
| "Safe Example": 'query = "SELECT * FROM users WHERE id = %s"\ncursor.execute(query, (user_id,))', | |
| } | |
| for name, template in templates.items(): | |
| if st.button(f"📋 {name}", use_container_width=True): | |
| st.session_state.code_input = template | |
| st.rerun() | |
| # Action Buttons | |
| btn_col1, btn_col2 = st.columns(2) | |
| with btn_col1: | |
| analyze_clicked = st.button("🔍 Analyze Code", type="primary", use_container_width=True) | |
| with btn_col2: | |
| if st.button("🗑️ Clear Analysis", use_container_width=True): | |
| if 'result' in st.session_state: | |
| del st.session_state.result | |
| st.rerun() | |
| if analyze_clicked: | |
| if not code_input.strip(): | |
| st.warning("Please enter some code!") | |
| else: | |
| with st.spinner("Analyzing with AI..."): | |
| start_time = time.time() | |
| try: | |
| with torch.no_grad(): | |
| if analysis_mode == "API Mode": | |
| response = requests.post(f"{api_url}/analyze", json={"code": code_input, "detailed": True}, timeout=10) | |
| result = response.json() | |
| elif analysis_mode == "Fast (Rules Only)": | |
| result = tools["rules"].analyze(code_input) | |
| else: | |
| result = tools["detector"].combined_analysis(code_input) | |
| result["analysis_time"] = time.time() - start_time | |
| st.session_state.result = result | |
| except Exception as e: | |
| st.error(f"Analysis failed: {str(e)}") | |
| # Display results | |
| if 'result' in st.session_state: | |
| result = st.session_state.result | |
| st.markdown("---") | |
| score = result.get("security_score", 0) | |
| # --- BALLOON LOGIC --- | |
| if score == 100: | |
| st.balloons() | |
| st.markdown("### **:green[✓ The code is safe! You have a 100 security score.]**") | |
| elif score >= 80: | |
| st.markdown("**:orange[⚠ Code is mostly safe, but minor issues were found.]**") | |
| else: | |
| st.markdown("**:red[✖ Critical vulnerabilities detected. Action required.]**") | |
| col_m1, col_m2, col_m3, col_m4 = st.columns(4) | |
| summary = result.get("summary", {}) | |
| critical_count = summary.get("critical", 0) if isinstance(summary, dict) else 0 | |
| with col_m1: | |
| fig = go.Figure(go.Indicator( | |
| mode="gauge+number", | |
| value=score, | |
| title={"text": "Security Score"}, | |
| gauge={ | |
| 'axis': {'range': [0, 100]}, | |
| 'bar': {'color': "darkblue"}, | |
| 'steps': [ | |
| {'range': [0, 50], 'color': "red"}, | |
| {'range': [50, 80], 'color': "orange"}, | |
| {'range': [80, 100], 'color': "green"} | |
| ] | |
| } | |
| )) | |
| fig.update_layout(height=200, margin=dict(t=50, b=0, l=10, r=10)) | |
| st.plotly_chart(fig, use_container_width=True) | |
| col_m2.metric("Issues", result.get("issue_count", 0)) | |
| col_m3.metric("Critical", critical_count) | |
| col_m4.metric("Analysis Time", f"{result.get('analysis_time', 0):.2f}s") | |
| # --- AI FIX SUGGESTIONS --- | |
| issues = result.get("issues", []) | |
| if issues: | |
| st.subheader("🚨 Detected Issues") | |
| for issue in issues: | |
| with st.expander(f"[{issue.get('severity', 'UNKNOWN')}] {issue.get('type', 'issue').replace('_', ' ').title()} - Line {issue.get('line', '??')}"): | |
| st.write(f"**Message:** {issue.get('message')}") | |
| if show_fixes: | |
| st.write("**💡 Fix Suggestions:**") | |
| st.write("**💡 AI Recommended Fixes:**") | |
| fixes = tools["fix_gen"].get_fixes(issue.get('message', ''), issue.get('type', '')) | |
| # Numbered Display | |
| for idx, fix in enumerate(fixes, 1): | |
| st.markdown(f"**Suggestion {idx}:**") | |
| st.code(fix, language="python") | |
| # --- TAB 2: BATCH --- | |
| with tab2: | |
| st.subheader("Batch Analysis") | |
| uploaded_files = st.file_uploader("Upload multiple .py files", type=['py'], accept_multiple_files=True) | |
| batch_btn1, batch_btn2 = st.columns(2) | |
| with batch_btn1: | |
| run_batch = st.button("🔍 Analyze All Files", type="primary", use_container_width=True) | |
| with batch_btn2: | |
| if st.button("🗑️ Clear Batch Results", use_container_width=True): | |
| if 'batch_df' in st.session_state: | |
| del st.session_state.batch_df | |
| st.rerun() | |
| if uploaded_files and run_batch: | |
| with st.spinner(f"Analyzing {len(uploaded_files)} files..."): | |
| results = [] | |
| for file in uploaded_files: | |
| code = file.getvalue().decode("utf-8") | |
| res = tools["detector"].combined_analysis(code) | |
| results.append({ | |
| "file": file.name, | |
| "score": res.get("security_score", 0), | |
| "issues": res.get("issue_count", 0), | |
| "critical": res.get("summary", {}).get("critical", 0) | |
| }) | |
| st.session_state.batch_df = pd.DataFrame(results) | |
| if 'batch_df' in st.session_state: | |
| df = st.session_state.batch_df | |
| st.dataframe(df, use_container_width=True) | |
| fig_b = go.Figure(data=[go.Bar(name='Score', x=df['file'], y=df['score']), go.Bar(name='Issues', x=df['file'], y=df['issues'])]) | |
| st.plotly_chart(fig_b, use_container_width=True) | |
| # --- TAB 3: API --- | |
| with tab3: | |
| st.subheader("API Integration") | |
| col_api1, col_api2 = st.columns(2) | |
| with col_api1: | |
| if st.button("Test Health Check"): | |
| try: | |
| response = requests.get(f"{api_url}/health") | |
| st.success(f"✅ API Healthy: {response.json()['status']}") | |
| except: st.error("❌ Connection Failed") | |
| with col_api2: | |
| st.markdown(f"- [OpenAPI Docs]({api_url}/docs)\n- [ReDoc]({api_url}/redoc)") | |
| # --- TAB 4: DASHBOARD (Updated per your request) --- | |
| with tab4: | |
| st.subheader("Project Dashboard") | |
| col_d1, col_d2, col_d3 = st.columns(3) | |
| with col_d1: | |
| st.metric("Vulnerability Types", "10+") | |
| st.metric("Detection Accuracy", "92%") | |
| with col_d2: | |
| st.metric("Supported Languages", "Python") | |
| st.metric("Response Time", "< 2s") | |
| with col_d3: | |
| st.metric("Integration Points", "4") | |
| st.metric("Lines Analyzed", "1000+") | |
| st.markdown("---") | |
| # 2. Centered Architecture Diagram | |
| st.write("**System Architecture:**") | |
| # Create 3 columns to center the image in the middle one | |
| # The [1, 2, 1] ratio makes the middle column 50% of the page width | |
| buf1, main_col, buf2 = st.columns([1, 2, 1]) | |
| with main_col: | |
| import os | |
| # Use the filename of the image you saved | |
| img_path = os.path.join(os.getcwd(), "architecture.jpg") | |
| if os.path.exists(img_path): | |
| st.image(img_path, width=600, caption="Project Workflow Diagram") | |
| else: | |
| st.info("💡 Please save your diagram as 'architecture.jpg' in the project folder.") | |
| st.markdown("---") | |
| st.caption("AI Code Security Scanner") |