File size: 4,585 Bytes
00a1906 |
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 |
import torch
import gradio as gr
import numpy as np
from transformers import AutoTokenizer, AutoModelForSequenceClassification
class ThreatAnalyzer:
def __init__(self):
# Load pretrained model for threat detection
model_name = 'distilbert-base-uncased-finetuned-sst-2-english'
self.tokenizer = AutoTokenizer.from_pretrained(model_name)
self.model = AutoModelForSequenceClassification.from_pretrained(model_name)
def analyze_threat(self, input_text):
"""
Comprehensive threat analysis with recommendations
"""
# Detect threat probability
inputs = self.tokenizer(input_text, return_tensors="pt", truncation=True, padding=True)
outputs = self.model(**inputs)
predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
threat_score = predictions[0][1].item()
# Categorize threat types
threat_patterns = {
"phishing": ["verify", "urgent", "suspend", "account", "click here"],
"social_engineering": ["tech support", "password", "credentials", "urgent"],
"malware": ["download", "patch", "update", "critical"],
"impersonation": ["admin", "support", "verify identity"]
}
# Analyze input for specific threat patterns
detected_patterns = []
for threat_type, keywords in threat_patterns.items():
if any(keyword.lower() in input_text.lower() for keyword in keywords):
detected_patterns.append(threat_type)
# Generate recommendations based on threat analysis
recommendations = []
if threat_score > 0.7:
recommendations.extend([
"🚨 High Risk: Do NOT interact with this message",
"Immediately report to IT security",
"Do not click any links or download attachments"
])
elif threat_score > 0.4:
recommendations.extend([
"⚠️ Potential Threat Detected",
"Verify the source independently",
"Contact sender through official channels"
])
# If specific threat patterns detected, add specific advice
if "phishing" in detected_patterns:
recommendations.append("Phishing Indicator: Check sender's email address carefully")
if "social_engineering" in detected_patterns:
recommendations.append("Social Engineering Alert: Never share personal credentials")
if "malware" in detected_patterns:
recommendations.append("Malware Risk: Scan with updated antivirus before opening")
return {
"threat_score": threat_score,
"threat_level": self._classify_threat_level(threat_score),
"detected_patterns": detected_patterns,
"recommendations": recommendations
}
def _classify_threat_level(self, score):
"""Classify threat level"""
if score > 0.8: return "CRITICAL THREAT"
if score > 0.6: return "HIGH THREAT"
if score > 0.4: return "MODERATE THREAT"
if score > 0.2: return "LOW THREAT"
return "NO SIGNIFICANT THREAT"
def launch_threat_analysis_demo():
"""Create Gradio interface for threat analysis"""
analyzer = ThreatAnalyzer()
def analyze_input(input_text):
try:
result = analyzer.analyze_threat(input_text)
# Format output for clear presentation
output = f"""
🔍 Threat Analysis Results:
Threat Score: {result['threat_score']:.2%}
Threat Level: {result['threat_level']}
Detected Threat Patterns:
{', '.join(result['detected_patterns']) or 'No specific patterns identified'}
Recommendations:
{chr(10).join('• ' + rec for rec in result['recommendations'])}
💡 Always verify suspicious communications
"""
return output
except Exception as e:
return f"Error in analysis: {str(e)}"
# Gradio interface
with gr.Blocks() as demo:
gr.Markdown("# Cybersecurity Threat Analysis")
input_text = gr.Textbox(label="Enter suspicious text or message")
output_text = gr.Textbox(label="Threat Analysis Results")
analyze_btn = gr.Button("Analyze Threat")
analyze_btn.click(fn=analyze_input, inputs=input_text, outputs=output_text)
return demo
# Run the demo
if __name__ == "__main__":
demo = launch_threat_analysis_demo()
demo.launch(debug=True, share=True) |