Spaces:
Sleeping
Sleeping
File size: 3,849 Bytes
fdff15a |
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 |
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import gradio as gr
from utils.heuristics import load_rules, explain_email, extract_keywords
from langdetect import detect
from pathlib import Path
import re
from utils.virustotal import check_url_virustotal
import extract_msg
# Modelo
model_name = "ElSlay/BERT-Phishing-Email-Model"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)
# Heurísticas
rules = load_rules("rules.yaml")
# Classificação com BERT
def classify_email(email_text):
inputs = tokenizer(email_text, return_tensors="pt", truncation=True, padding=True)
outputs = model(**inputs)
probs = torch.softmax(outputs.logits, dim=1).detach().numpy()[0]
labels = ["Legitimate", "Phishing"]
prediction = labels[probs.argmax()]
confidence = probs.max()
return prediction, confidence
# Análise completa do email
def analyze_email(file_input=None, text_input=None):
email_text = None
if text_input:
email_text = text_input
elif file_input:
path = file_input.name if hasattr(file_input, "name") else file_input
if path.endswith(".txt"):
with open(path, "r", encoding="utf-8") as f:
email_text = f.read()
elif path.endswith(".msg"):
msg = extract_msg.Message(path)
email_text = f"{msg.subject or ''}\n{msg.body or ''}"
else:
return "Unsupported file type."
else:
return "No input provided."
if not email_text:
return "Could not extract text from file."
# Classificação
label, confidence = classify_email(email_text)
# Heurística
explanations, score = explain_email(email_text, rules)
# Keywords
lang = detect(email_text)
keywords = extract_keywords(email_text, lang)
keywords_text = "Top keywords: " + ", ".join(keywords)
# Explicação heurística
explanation_text = "📌 Explanation:\n• " + "\n• ".join(explanations)
# Verificação VirusTotal
urls = re.findall(r"http[s]?://\S+", email_text)
vt_results = []
for url in urls:
stats = check_url_virustotal(url)
if "error" in stats:
vt_results.append(f"URL: {url} | VT: {stats['error']}")
else:
vt_results.append(f"URL: {url} | Malicious: {stats.get('malicious', 0)}, Suspicious: {stats.get('suspicious', 0)}, Harmless: {stats.get('harmless', 0)}")
vt_text = "\n".join(vt_results) if vt_results else "No URLs found."
return f"Classification: {label} ({confidence:.2%})\n\n{explanation_text}\n\nScore: {score}\n\n{keywords_text}\n\nVirusTotal Results:\n{vt_text}"
# Carregar exemplos
def update_text_from_example(example_name):
return example_emails[example_name]
def load_example_files():
examples_path = Path("examples")
files = sorted(examples_path.glob("*.txt"))
return {file.name: file.read_text(encoding="utf-8") for file in files}
example_emails = load_example_files()
# Interface Gradio
with gr.Blocks() as demo:
gr.Markdown("## 🛡️ PhishHunter – Email Phishing Detector")
with gr.Row():
dropdown = gr.Dropdown(
choices=list(example_emails.keys()),
label="Load Example Email",
info="Select a sample email to test",
)
textbox = gr.Textbox(lines=15, label="Paste or load email content")
filebox = gr.File(label="Upload email (.txt or .msg)", file_types=[".txt", ".msg"])
dropdown.change(fn=update_text_from_example, inputs=dropdown, outputs=textbox)
output = gr.Textbox(label="Classification & Explanation")
btn = gr.Button("Analyze")
btn.click(fn=analyze_email, inputs=[filebox, textbox], outputs=output)
if __name__ == "__main__":
demo.launch()
|