File size: 5,299 Bytes
4b82ab5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 | """
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()
|