ashwinkumarbv commited on
Commit
3a3a254
1 Parent(s): c203e1d

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +59 -0
app.py ADDED
@@ -0,0 +1,59 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification
3
+ import torch
4
+
5
+ # Load Pre-trained Model (Replace with your fine-tuned model)
6
+ @st.cache_resource
7
+ def load_model():
8
+ model_name = "microsoft/codebert-base" # Replace with a fine-tuned model for vulnerability detection
9
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
10
+ model = AutoModelForSequenceClassification.from_pretrained(model_name)
11
+ return tokenizer, model
12
+
13
+ tokenizer, model = load_model()
14
+
15
+ # Vulnerability Explanation Function
16
+ def analyze_code(code_snippet):
17
+ # Tokenize Input
18
+ inputs = tokenizer(code_snippet, return_tensors="pt", truncation=True, max_length=512)
19
+ outputs = model(**inputs)
20
+ predictions = torch.softmax(outputs.logits, dim=1)
21
+ vulnerability_score = predictions[0][1].item() # Assuming index 1 corresponds to "vulnerable"
22
+
23
+ # Generate Explanation
24
+ if vulnerability_score > 0.6:
25
+ explanation = (
26
+ f"The code has a high likelihood of being vulnerable. The model detected patterns "
27
+ f"indicative of potential security issues."
28
+ )
29
+ elif vulnerability_score > 0.3:
30
+ explanation = (
31
+ f"The code has some potential vulnerabilities. Review the logic carefully, especially in "
32
+ f"sensitive operations like input validation or file handling."
33
+ )
34
+ else:
35
+ explanation = (
36
+ f"The code appears to be safe based on the analysis. However, manual review is always recommended."
37
+ )
38
+ return vulnerability_score, explanation
39
+
40
+ # Streamlit UI
41
+ st.title("AI-Enhanced Code Vulnerability Scanner")
42
+ st.markdown("""
43
+ This tool uses AI to detect vulnerabilities in Python code and provides explanations for potential issues.
44
+ """)
45
+
46
+ # Input Section
47
+ code_snippet = st.text_area("Paste your Python code here:", height=200)
48
+ analyze_button = st.button("Analyze Code")
49
+
50
+ if analyze_button and code_snippet.strip():
51
+ with st.spinner("Analyzing your code..."):
52
+ score, explanation = analyze_code(code_snippet)
53
+
54
+ # Display Results
55
+ st.subheader("Analysis Results")
56
+ st.write(f"**Vulnerability Score:** {score:.2f}")
57
+ st.write(f"**Explanation:** {explanation}")
58
+ else:
59
+ st.info("Please paste Python code and click 'Analyze Code'.")