app.py

"""
Streamlit UI for Vulnerability Detection
Interactive web interface
"""

import streamlit as st
import sys
from pathlib import Path

sys.path.append(str(Path(__file__).parent))

from src.inference import VulnerabilityDetector

# Page config
st.set_page_config(
    page_title="Code Vulnerability Detector",
    page_icon="🔒",
    layout="wide"
)

# Initialize detector (cache it so it loads only once)
@st.cache_resource
def load_detector():
    return VulnerabilityDetector()

# Main app
def main():
    st.title("🔒 AI-Powered Code Vulnerability Detection")
    st.markdown("### Detect security vulnerabilities in your code using fine-tuned CodeT5")
    
    # Sidebar
    with st.sidebar:
        st.header("ℹ️ About")
        st.markdown("""
        This tool uses a fine-tuned CodeT5 model to detect security vulnerabilities in source code.
        
        **Supported Languages:**
        - C/C++
        - Python
        - JavaScript
        
        **Detection Types:**
        - Buffer Overflow
        - SQL Injection
        - Command Injection
        - Format String Bugs
        - And more...
        """)
        
        st.header("📊 Model Info")
        try:
            detector = load_detector()
            st.success("Model loaded successfully!")
        except Exception as e:
            st.error(f"Error loading model: {e}")
            st.stop()
    
    # Main area
    col1, col2 = st.columns([1, 1])
    
    with col1:
        st.header("📝 Enter Code")
        
        # Example selector
        example = st.selectbox(
            "Or try an example:",
            ["Custom", "Buffer Overflow", "SQL Injection", "Safe Code"]
        )
        
        if example == "Buffer Overflow":
            default_code = '''void copy(char *input) {
    char buffer[8];
    strcpy(buffer, input);
}'''
        elif example == "SQL Injection":
            default_code = '''def get_user(user_id):
    query = "SELECT * FROM users WHERE id=" + user_id
    cursor.execute(query)
    return cursor.fetchone()'''
        elif example == "Safe Code":
            default_code = '''def add_numbers(a, b):
    return a + b'''
        else:
            default_code = ""
        
        code_input = st.text_area(
            "Paste your code here:",
            value=default_code,
            height=300,
            placeholder="Enter source code to analyze..."
        )
        
        analyze_button = st.button("🔍 Analyze Code", type="primary", use_container_width=True)
    
    with col2:
        st.header("📊 Analysis Results")
        
        if analyze_button and code_input.strip():
            with st.spinner("Analyzing code..."):
                try:
                    result = detector.predict(code_input)
                    
                    # Display result
                    if result['prediction'] == 1:
                        st.error(f"⚠️ {result['label']}")
                        st.progress(result['probabilities']['vulnerable'])
                    else:
                        st.success(f"✅ {result['label']}")
                        st.progress(result['probabilities']['safe'])
                    
                    # Confidence metrics
                    st.subheader("Confidence Breakdown")
                    col_a, col_b = st.columns(2)
                    
                    with col_a:
                        st.metric(
                            "Safe Probability",
                            f"{result['probabilities']['safe']:.1%}",
                            delta=None
                        )
                    
                    with col_b:
                        st.metric(
                            "Vulnerable Probability",
                            f"{result['probabilities']['vulnerable']:.1%}",
                            delta=None
                        )
                    
                    # Recommendations
                    if result['prediction'] == 1:
                        st.subheader("🛡️ Recommendations")
                        st.warning("""
                        **This code appears to have security vulnerabilities.**
                        
                        Common fixes:
                        - Use bounds-checked functions (strncpy instead of strcpy)
                        - Use parameterized queries for SQL
                        - Validate and sanitize all user inputs
                        - Avoid eval() and system() with user input
                        """)
                    else:
                        st.subheader("Good Practices")
                        st.info("""
                        This code appears to follow security best practices!
                        
                        Remember to:
                        - Keep dependencies updated
                        - Perform regular security audits
                        - Use static analysis tools
                        - Follow OWASP guidelines
                        """)
                
                except Exception as e:
                    st.error(f"Error during analysis: {e}")
        
        elif analyze_button:
            st.warning("Please enter some code to analyze.")

if __name__ == "__main__":
    main()