File size: 2,300 Bytes
3a3a254 |
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 |
import streamlit as st
from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification
import torch
# Load Pre-trained Model (Replace with your fine-tuned model)
@st.cache_resource
def load_model():
model_name = "microsoft/codebert-base" # Replace with a fine-tuned model for vulnerability detection
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)
return tokenizer, model
tokenizer, model = load_model()
# Vulnerability Explanation Function
def analyze_code(code_snippet):
# Tokenize Input
inputs = tokenizer(code_snippet, return_tensors="pt", truncation=True, max_length=512)
outputs = model(**inputs)
predictions = torch.softmax(outputs.logits, dim=1)
vulnerability_score = predictions[0][1].item() # Assuming index 1 corresponds to "vulnerable"
# Generate Explanation
if vulnerability_score > 0.6:
explanation = (
f"The code has a high likelihood of being vulnerable. The model detected patterns "
f"indicative of potential security issues."
)
elif vulnerability_score > 0.3:
explanation = (
f"The code has some potential vulnerabilities. Review the logic carefully, especially in "
f"sensitive operations like input validation or file handling."
)
else:
explanation = (
f"The code appears to be safe based on the analysis. However, manual review is always recommended."
)
return vulnerability_score, explanation
# Streamlit UI
st.title("AI-Enhanced Code Vulnerability Scanner")
st.markdown("""
This tool uses AI to detect vulnerabilities in Python code and provides explanations for potential issues.
""")
# Input Section
code_snippet = st.text_area("Paste your Python code here:", height=200)
analyze_button = st.button("Analyze Code")
if analyze_button and code_snippet.strip():
with st.spinner("Analyzing your code..."):
score, explanation = analyze_code(code_snippet)
# Display Results
st.subheader("Analysis Results")
st.write(f"**Vulnerability Score:** {score:.2f}")
st.write(f"**Explanation:** {explanation}")
else:
st.info("Please paste Python code and click 'Analyze Code'.") |