Spaces:
Sleeping
Sleeping
| """ | |
| Cyber Defense SOC Email Threat Classifier | |
| Production Web Application powered by Gradio, Scikit-Learn, and MLflow. | |
| Hugging Face Spaces Deployable. | |
| """ | |
| import os | |
| import sys | |
| import html | |
| import gradio as gr | |
| # ZeroGPU Spaces require at least one GPU-decorated callable at startup. The | |
| # classifier itself is CPU-friendly, but the decorator lets this app run on | |
| # the existing ZeroGPU Space configuration. Keep a local fallback so tests | |
| # and ordinary CPU launches do not require the Spaces runtime package. | |
| try: | |
| import spaces | |
| except ImportError: | |
| class _LocalSpaces: | |
| GPU = staticmethod(lambda function: function) | |
| spaces = _LocalSpaces() | |
| # Ensure src modules are resolvable | |
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) | |
| from src.predictor import EmailThreatPredictor | |
| from src.data_loader import CLASS_NAMES | |
| # Initialize global inference engine | |
| MODEL_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "models", "artifacts", "best_threat_model.joblib") | |
| predictor = None | |
| def get_predictor(): | |
| global predictor | |
| if predictor is None: | |
| if not os.path.exists(MODEL_PATH): | |
| raise FileNotFoundError( | |
| f"Model file not found at '{MODEL_PATH}'. " | |
| "Ensure train.py has been run and model artifacts are preserved." | |
| ) | |
| predictor = EmailThreatPredictor(MODEL_PATH) | |
| return predictor | |
| # Curated sample emails for quick testing | |
| SAMPLE_EMAILS = { | |
| "Legitimate: Project Update": { | |
| "subject": "Quarterly Q3 Financial Review Meeting Notes", | |
| "body": "Hi Team,\n\nPlease find the summary of our Q3 financial review meeting attached below. Key highlights include an 18% increase in operational efficiency across cloud services and steady revenue growth in APAC.\n\nAction items for next sprint:\n1. Dave to finalize infrastructure budget by Thursday.\n2. Sarah to review AWS reserved instance utilization.\n\nLet me know if anyone has questions before our next sync.\n\nBest regards,\nMarcus Vance\nDirector of Operations" | |
| }, | |
| "Spam: SEO Backlink Pitch": { | |
| "subject": "EXCLUSIVE OFFER: Boost Your Website Traffic by 500% Overnight!", | |
| "body": "Struggling to rank #1 on Google? Look no further!\n\nOur automated SEO backlink blast will drive 50,000 guaranteed unique visitors to your website in just 48 hours. We use proprietary high-authority PBN links to outrank all your competitors.\n\nSpecial introductory price: Only $49.99! Click here to claim your discount coupon now: http://cheap-seo-rankings-fast.biz/promo\n\nAct now! Offer expires at midnight!" | |
| }, | |
| "Suspicious: Microsoft 365 Phishing": { | |
| "subject": "URGENT: Microsoft 365 Account Suspension Notice - Action Required", | |
| "body": "SECURITY ALERT: We detected unauthorized login attempts to your Office 365 mailbox from IP 185.220.101.5 (Moscow, Russia).\n\nAs a security precaution, your email access and OneDrive cloud sync will be permanently TERMINATED within 24 hours unless you re-verify your employee identity.\n\nPlease click here immediately to verify your credentials and restore access: http://192.168.1.105/auth/m365/login?user=account_verification\n\nFailure to authenticate will result in account lockout by your IT administrator.\n\nOffice 365 Security Operations Desk" | |
| }, | |
| "Suspicious: Urgent CEO Wire Fraud (BEC)": { | |
| "subject": "URGENT: Confidential Wire Transfer Request - CEO Authorization", | |
| "body": "Abdul,\n\nAre you currently at your desk? I am currently in an urgent confidential acquisition meeting with external legal counsel and cannot take calls.\n\nWe need to process an immediate wire transfer of $48,500 to our acquisition escrow account before 3:00 PM EST to finalize the agreement. Please confirm your availability to execute this payment immediately and I will forward the beneficiary routing and SWIFT details.\n\nDo not discuss this with anyone on the floor as it is strictly subject to NDA.\n\nSent from my iPhone\nArthur Pendelton\nChief Executive Officer" | |
| } | |
| } | |
| def classify_email(subject: str, body: str): | |
| """ | |
| Primary Gradio callback function for classifying email threat levels. | |
| """ | |
| try: | |
| pred_engine = get_predictor() | |
| res = pred_engine.predict(subject, body) | |
| except Exception as e: | |
| error_html = f""" | |
| <div style="padding: 16px; background: #2d1515; border: 1px solid #e53e3e; border-radius: 8px; color: #ff8b8b;"> | |
| <strong>Inference Notice:</strong> {html.escape(str(e))} | |
| </div> | |
| """ | |
| return error_html, {} | |
| pred_class = res["predicted_class"] | |
| confidence = res["confidence"] | |
| severity = res["severity"] | |
| badge_color = res["badge_color"] | |
| risk_score = res["risk_score_out_of_10"] | |
| indicators = res["indicators"] | |
| recs = res["recommendations"] | |
| # Formatting detected indicators | |
| detected_tags = [] | |
| if indicators["urgency_count"] > 0: | |
| detected_tags.append(f"<span style='background:#742a2a; color:#feb2b2; padding:3px 8px; border-radius:4px; font-size:12px; margin-right:6px;'>⚠️ Urgency ({indicators['urgency_count']} hits: {', '.join(indicators['urgency_hits'][:3])})</span>") | |
| if indicators["credential_count"] > 0: | |
| detected_tags.append(f"<span style='background:#742a2a; color:#feb2b2; padding:3px 8px; border-radius:4px; font-size:12px; margin-right:6px;'>🔑 Credential Harvesting ({indicators['credential_count']} hits)</span>") | |
| if indicators["financial_count"] > 0: | |
| detected_tags.append(f"<span style='background:#742a2a; color:#feb2b2; padding:3px 8px; border-radius:4px; font-size:12px; margin-right:6px;'>💳 Financial / Wire Trigger ({indicators['financial_count']} hits)</span>") | |
| if indicators["impersonation_count"] > 0: | |
| detected_tags.append(f"<span style='background:#553c9a; color:#d6bcfa; padding:3px 8px; border-radius:4px; font-size:12px; margin-right:6px;'>👔 Impersonation Cues ({indicators['impersonation_count']} hits)</span>") | |
| if indicators["ip_url_count"] > 0: | |
| detected_tags.append(f"<span style='background:#9b2c2c; color:#fff; padding:3px 8px; border-radius:4px; font-size:12px; margin-right:6px;'>🌐 Direct IP URL Detected</span>") | |
| if indicators["shortened_url_count"] > 0: | |
| detected_tags.append(f"<span style='background:#7b341e; color:#fbd38d; padding:3px 8px; border-radius:4px; font-size:12px; margin-right:6px;'>🔗 URL Shortener ({indicators['shortened_url_count']})</span>") | |
| if indicators["url_count"] > 0 and indicators["ip_url_count"] == 0 and indicators["shortened_url_count"] == 0: | |
| detected_tags.append(f"<span style='background:#285e61; color:#b2f5ea; padding:3px 8px; border-radius:4px; font-size:12px; margin-right:6px;'>🔗 Hyperlinks Found ({indicators['url_count']})</span>") | |
| if indicators["caps_ratio"] > 0.20: | |
| detected_tags.append(f"<span style='background:#744210; color:#fefcbf; padding:3px 8px; border-radius:4px; font-size:12px; margin-right:6px;'>🔠 Excessive Uppercase ({indicators['caps_ratio']*100:.1f}%)</span>") | |
| indicators_html = " ".join(detected_tags) if detected_tags else "<span style='color: #a0aec0; font-size: 13px;'>No anomalous behavioral cues flagged.</span>" | |
| # Actionable guidance checklist | |
| rec_items = "".join([f"<li style='margin-bottom: 6px;'>{html.escape(r)}</li>" for r in recs]) | |
| # Result card HTML | |
| result_html = f""" | |
| <div style="background: #1a202c; border: 1px solid #2d3748; border-radius: 12px; padding: 20px; color: #edf2f7; font-family: system-ui, sans-serif; box-shadow: 0 4px 6px rgba(0,0,0,0.3);"> | |
| <div style="display: flex; justify-content: space-between; align-items: center; border-bottom: 1px solid #2d3748; padding-bottom: 14px; margin-bottom: 16px;"> | |
| <div> | |
| <span style="font-size: 12px; text-transform: uppercase; letter-spacing: 1.5px; color: #a0aec0; font-weight: 600;">SOC Threat Assessment</span> | |
| <div style="font-size: 26px; font-weight: 800; color: {badge_color}; margin-top: 4px;"> | |
| {pred_class.upper()} | |
| </div> | |
| </div> | |
| <div style="text-align: right;"> | |
| <span style="display: inline-block; padding: 6px 14px; background: {badge_color}22; border: 1.5px solid {badge_color}; color: {badge_color}; border-radius: 20px; font-weight: 700; font-size: 13px;"> | |
| {severity} (Risk: {risk_score}/10) | |
| </span> | |
| <div style="font-size: 13px; color: #a0aec0; margin-top: 6px;"> | |
| Model Confidence: <strong>{confidence*100:.1f}%</strong> | |
| </div> | |
| </div> | |
| </div> | |
| <div style="margin-bottom: 18px;"> | |
| <div style="font-size: 13px; font-weight: 600; color: #cbd5e0; margin-bottom: 6px;">Behavioral Threat Indicators:</div> | |
| <div style="line-height: 2.0;"> | |
| {indicators_html} | |
| </div> | |
| </div> | |
| <div style="background: #232d3f; border-left: 4px solid {badge_color}; border-radius: 6px; padding: 14px; margin-bottom: 14px;"> | |
| <div style="font-size: 13px; font-weight: 700; color: #e2e8f0; margin-bottom: 8px;">Recommended Incident Response Action:</div> | |
| <ul style="margin: 0; padding-left: 20px; font-size: 13px; color: #cbd5e0;"> | |
| {rec_items} | |
| </ul> | |
| </div> | |
| <div style="font-size: 11px; color: #718096; line-height: 1.4; border-top: 1px solid #2d3748; padding-top: 10px;"> | |
| ⚠️ <strong>Security Advisory:</strong> This classifier provides heuristic assistance and machine learning probability metrics. No automated classifier guarantees 100% email security. Always adhere to organizational cybersecurity standards and report anomalies to your CSIRT. | |
| </div> | |
| </div> | |
| """ | |
| return result_html, res["probabilities"] | |
| # Custom cybersecurity CSS theme | |
| CUSTOM_CSS = """ | |
| body, .gradio-container { | |
| background-color: #0d1117 !important; | |
| font-family: 'Segoe UI', Roboto, Helvetica, Arial, sans-serif; | |
| } | |
| .soc-header { | |
| background: linear-gradient(135deg, #161b22 0%, #0d1117 100%); | |
| border: 1px solid #30363d; | |
| border-radius: 12px; | |
| padding: 24px; | |
| margin-bottom: 20px; | |
| text-align: center; | |
| } | |
| .soc-badge { | |
| display: inline-block; | |
| background: #1f6feb22; | |
| border: 1px solid #1f6feb; | |
| color: #58a6ff; | |
| padding: 4px 12px; | |
| border-radius: 20px; | |
| font-size: 12px; | |
| font-weight: 600; | |
| letter-spacing: 1px; | |
| margin-bottom: 8px; | |
| } | |
| """ | |
| def create_ui(): | |
| """ | |
| Constructs the responsive cybersecurity Gradio UI blocks. | |
| """ | |
| with gr.Blocks(title="Email Threat Intelligence Classifier") as demo: | |
| with gr.Column(elem_classes=["soc-header"]): | |
| gr.HTML(""" | |
| <div style="text-align: center;"> | |
| <div class="soc-badge">SHIELD CYBER DEFENSE LABS • MLFLOW & TF-IDF ENGINE</div> | |
| <h1 style="color: #f0f6fc; margin: 8px 0; font-size: 28px; font-weight: 700;"> | |
| Email Spam & Business Email Threat Classifier | |
| </h1> | |
| <p style="color: #8b949e; margin: 0; font-size: 14px; max-width: 780px; margin: 0 auto;"> | |
| Real-time tripartite email security analysis for <strong>Legitimate</strong> corporate communications, | |
| unsolicited <strong>Spam</strong>, and high-risk <strong>Suspicious</strong> threats (Phishing & BEC wire fraud). | |
| </p> | |
| </div> | |
| """) | |
| with gr.Row(): | |
| # Left Column: Inputs & Samples | |
| with gr.Column(scale=5): | |
| gr.Markdown("### 📨 Email Message Ingestion") | |
| subject_input = gr.Textbox( | |
| label="Email Subject Line", | |
| placeholder="e.g. URGENT: Microsoft 365 Account Suspension", | |
| lines=1 | |
| ) | |
| body_input = gr.Textbox( | |
| label="Email Body Content", | |
| placeholder="Paste full raw email body text including headers, links, and greetings...", | |
| lines=8 | |
| ) | |
| with gr.Row(): | |
| classify_btn = gr.Button("🛡️ Inspect & Classify Threat", variant="primary", scale=2) | |
| clear_btn = gr.Button("🗑️ Clear Inputs", variant="secondary", scale=1) | |
| gr.Markdown("#### ⚡ Quick Telemetry Test Samples") | |
| with gr.Row(): | |
| btn_legit = gr.Button("📄 Legitimate Sample", size="sm") | |
| btn_spam = gr.Button("📢 Spam Sample", size="sm") | |
| with gr.Row(): | |
| btn_m365 = gr.Button("🎣 Phishing Sample", size="sm") | |
| btn_bec = gr.Button("🚨 CEO BEC Wire Fraud", size="sm") | |
| # Right Column: Analysis & Metrics | |
| with gr.Column(scale=5): | |
| gr.Markdown("### 🔍 Threat Telemetry & Risk Assessment") | |
| result_display = gr.HTML(""" | |
| <div style="background: #161b22; border: 1px dashed #30363d; border-radius: 12px; padding: 40px; text-align: center; color: #8b949e;"> | |
| <div style="font-size: 32px; margin-bottom: 12px;">🛡️</div> | |
| <div style="font-weight: 600; font-size: 15px; color: #c9d1d9;">Awaiting Message Ingestion</div> | |
| <div style="font-size: 13px; margin-top: 4px;">Enter email contents or select a test sample to evaluate risk level.</div> | |
| </div> | |
| """) | |
| prob_chart = gr.Label( | |
| label="Model Confidence & Class Probability Distribution", | |
| num_top_classes=3 | |
| ) | |
| # Wire sample button handlers | |
| def load_sample(sample_key): | |
| sample = SAMPLE_EMAILS[sample_key] | |
| return sample["subject"], sample["body"] | |
| btn_legit.click(lambda: load_sample("Legitimate: Project Update"), outputs=[subject_input, body_input]) | |
| btn_spam.click(lambda: load_sample("Spam: SEO Backlink Pitch"), outputs=[subject_input, body_input]) | |
| btn_m365.click(lambda: load_sample("Suspicious: Microsoft 365 Phishing"), outputs=[subject_input, body_input]) | |
| btn_bec.click(lambda: load_sample("Suspicious: Urgent CEO Wire Fraud (BEC)"), outputs=[subject_input, body_input]) | |
| # Action handlers | |
| classify_btn.click( | |
| fn=classify_email, | |
| inputs=[subject_input, body_input], | |
| outputs=[result_display, prob_chart] | |
| ) | |
| def clear_all(): | |
| empty_html = """ | |
| <div style="background: #161b22; border: 1px dashed #30363d; border-radius: 12px; padding: 40px; text-align: center; color: #8b949e;"> | |
| <div style="font-size: 32px; margin-bottom: 12px;">🛡️</div> | |
| <div style="font-weight: 600; font-size: 15px; color: #c9d1d9;">Awaiting Message Ingestion</div> | |
| <div style="font-size: 13px; margin-top: 4px;">Enter email contents or select a test sample to evaluate risk level.</div> | |
| </div> | |
| """ | |
| return "", "", empty_html, {} | |
| clear_btn.click(fn=clear_all, outputs=[subject_input, body_input, result_display, prob_chart]) | |
| # Footer notes | |
| gr.Markdown(""" | |
| --- | |
| <div style="display: flex; justify-content: space-between; font-size: 12px; color: #8b949e;"> | |
| <div>Model: <strong>TF-IDF + Scikit-Learn / XGBoost Ensemble</strong> | Tracked via <strong>MLflow 3.17</strong></div> | |
| <div>Deployment: <strong>Hugging Face Spaces</strong> | Methodology: <strong>Jira Agile Scrum</strong></div> | |
| </div> | |
| """) | |
| return demo | |
| # App launch entrypoint | |
| demo = create_ui() | |
| if __name__ == "__main__": | |
| demo.launch( | |
| server_name="0.0.0.0", | |
| server_port=7860, | |
| share=False, | |
| theme=gr.themes.Base(), | |
| css=CUSTOM_CSS | |
| ) | |