code.py

import re
import random
from typing import Optional
import streamlit as st
import requests
import json
from urllib.parse import urlparse
from datetime import datetime

st.set_page_config(page_title="PhishCheck (Demo)", page_icon="🛡️", layout="centered")

st.title("🛡️ PhishCheck")
st.caption("A lightweight phishing-triage helper for staff. Training & awareness only (not a guarantee).")

role = st.selectbox("Role (optional)", ["General", "Finance", "HR", "IT/Engineering", "Leadership"], key="role")

# Sidebar: instructions, tips & quick reference
st.sidebar.title("PhishCheck — Tips & Instructions")
st.sidebar.markdown("**Quick steps**")
st.sidebar.write("1. Paste or upload the email text.\n2. Click Analyze.\n3. Review flags, links, and the suggested draft.\n4. Use the AI assistant to refine the message before reporting.")

with st.sidebar.expander("Red flags (quick)"):
    st.write("- Urgency/pressure language (act now, urgent)")
    st.write("- Requests for credentials or MFA codes")
    st.write("- Unexpected attachments or links")
    st.write("- Generic greetings (Dear user/customer)")
    st.write("- Requests for wire transfers/payments")

with st.sidebar.expander("Reporting & safe actions"):
    st.write("• Don’t click links or open attachments if unsure.")
    st.write("• Verify the request using a trusted channel (known phone/Slack).")
    st.write("• Capture a screenshot and copy message headers for investigators.")
    st.write("• Use your org's Report Phish button or contact IT immediately.")

with st.sidebar.expander("Tips & Tricks"):
    st.write("- Hover links to inspect the destination domain before clicking.")
    st.write("- Check sender display name vs email address (look for subtle typos).")
    st.write("- For finance: confirm any bank/account changes with a phone call to a known number.")
    st.write("- For HR: be wary of unsolicited attachments claiming to be resumes or payroll docs.")

st.sidebar.markdown("---")
st.sidebar.markdown("**Example suspicious phrases**")
st.sidebar.write("Act now, Verify your account, Final notice, Wire transfer, Click to avoid suspension")

# Optional Hugging Face Inference API settings
st.sidebar.markdown("---")
st.sidebar.subheader("Optional: Hugging Face LLM")
use_hf = st.sidebar.checkbox("Use Hugging Face Inference API for triage")
hf_token = st.sidebar.text_input("Hugging Face API token", type="password") if use_hf else None
hf_model = st.sidebar.text_input("Model name (e.g. google/flan-t5-large)", value="", placeholder="owner/model-id") if use_hf else None
hf_redact = st.sidebar.checkbox("Redact email addresses before sending", value=True) if use_hf else False
if use_hf:
    st.sidebar.write("Note: provide an API token with Inference API access. Sending email content to the model may expose sensitive data.")

# ---- Input methods ----
st.subheader("Add the email content")
uploaded = st.file_uploader("Upload a .txt file (optional)", type=["txt"], key="uploaded")

default_text = ""
if uploaded is not None:
    default_text = uploaded.read().decode("utf-8", errors="ignore")

st.text_area(
    "Paste the email (subject + body).",
    value=default_text,
    height=260,
    placeholder="Subject: ...\nFrom: ...\n\nHi, please verify your account...",
    key="email_text_input",
)

# ---- Helpers ----
def extract_urls(text: str):
    return re.findall(r"(https?://[^\s)>\"]+)", text)

def domain_of(url: str):
    try:
        return urlparse(url).netloc.lower()
    except Exception:
        return ""


def redact_emails(text: str):
    return re.sub(r"[\w\.-]+@[\w\.-]+\.[a-zA-Z]{2,}", "[REDACTED_EMAIL]", text)


def llm_score_email_hf(email_text: str, token: str, model: str, redact: bool = True) -> dict:
    """Call Hugging Face Inference API with a prompt asking for JSON output.
    Returns a dict with keys: risk, score, reasons (list), draft_report
    """
    text = redact_emails(email_text) if redact else email_text
    prompt = (
        "You are a security assistant. Classify the following email for phishing risk and RETURN ONLY a JSON object with keys: "
        "risk (Low|Medium|High), score (0-10), reasons (list of short strings), draft_report (short text).\n\nEmail:\n"
        + '"""' + text + '"""'
    )
    headers = {"Authorization": f"Bearer {token}", "Content-Type": "application/json"}
    payload = {"inputs": prompt, "parameters": {"max_new_tokens": 300, "temperature": 0.2}}
    url = f"https://api-inference.huggingface.co/models/{model}"
    resp = requests.post(url, headers=headers, data=json.dumps(payload), timeout=30)
    resp.raise_for_status()
    out = resp.json()
    # parse different HF response shapes
    generated = ""
    if isinstance(out, list) and len(out) > 0 and isinstance(out[0], dict):
        generated = out[0].get("generated_text", "")
    elif isinstance(out, dict) and "generated_text" in out:
        generated = out.get("generated_text", "")
    else:
        generated = str(out)

    # try to parse JSON from the model output
    try:
        return json.loads(generated)
    except Exception:
        m = re.search(r"\{.*\}", generated, flags=re.S)
        if m:
            try:
                return json.loads(m.group(0))
            except Exception:
                pass
    # fallback
    return {"risk": "Unknown", "score": 0, "reasons": ["LLM output unparsable"], "draft_report": generated}

def score_email(text: str, role: str):
    t = text.lower()
    reasons = []
    score = 0

    urgency_patterns = [
        "urgent", "immediately", "act now", "within 24 hours", "expires", "final warning",
        "account will be suspended", "verify now", "limited time", "action required",
        "alert", "security alert", "high-severity", "phish delivered", "delivered due to", "tenant override",
        "view alert", "view alert details", "notification"
    ]
    if any(p in t for p in urgency_patterns):
        score += 2
        reasons.append("Urgency/pressure language (common phishing tactic).")

    credential_patterns = ["password", "verify your account", "login", "mfa", "one-time code", "security check", "sign in"]
    if any(p in t for p in credential_patterns):
        score += 2
        reasons.append("Requests login/verification action.")

    payment_patterns = ["invoice", "wire", "bank", "payment", "gift card", "transfer", "refund"]
    if any(p in t for p in payment_patterns):
        score += 2
        reasons.append("Mentions payments/financial action (high-risk category).")

    urls = extract_urls(text)
    if urls:
        score += 1
        reasons.append(f"Contains link(s): {len(urls)} detected.")
        suspicious_domains = []
        for u in urls[:8]:
            d = domain_of(u)
            # allowlist can be customized for org
            if d and not any(x in d for x in ["google.com", "microsoft.com", "okta.com", "salesforce.com"]):
                suspicious_domains.append(d)
        if suspicious_domains:
            score += 2
            reasons.append("Some link domains look unfamiliar/untrusted: " + ", ".join(sorted(set(suspicious_domains))[:3]) + ".")

    if any(p in t for p in ["dear user", "dear customer", "valued customer", "attention"]):
        score += 1
        reasons.append("Generic greeting instead of your name/team.")
    if "unsubscribe" in t and "security" in t and "@" not in t:
        score += 1
        reasons.append("Vague security context and low specificity.")

    # --- Sender / impersonation checks ---
    def extract_emails(text: str):
        return re.findall(r"[\w\.-]+@[\w\.-]+\.[a-zA-Z]{2,}", text)

    emails = extract_emails(text)
    if emails:
        sender = emails[0].lower()
        sender_domain = sender.split("@")[-1]
        # look for a display name immediately before <email>
        m = re.search(r"([\w\s\.,'-]{2,100})<" + re.escape(emails[0]) + r">", text, flags=re.IGNORECASE)
        display = m.group(1).strip() if m else ""
        known_brands = ["microsoft", "office 365", "office365", "google", "okta", "salesforce", "amazon", "paypal"]
        allowlist = ["google.com", "microsoft.com", "okta.com", "salesforce.com"]
        # if display name contains a brand but domain is not the official one, bump risk
        for b in known_brands:
            if b in display.lower() or b in t:
                if not any(sender_domain.endswith(a) for a in allowlist):
                    score += 2
                    reasons.append(f"Sender impersonation or brand misuse detected ({b} referenced, sender domain: {sender_domain}).")
                    break

    # detect specific alert-style phrases that commonly appear in phishing
    alert_phrases = ["view alert details", "phish delivered due to", "phish delivered", "high-severity alert", "phish delivered due to tenant"]
    if any(p in t for p in alert_phrases):
        score += 2
        reasons.append("Alert-style notification language present (common in delivery/notification phishing).")

    # Role-based bump
    if role == "Finance" and any(p in t for p in ["wire", "bank", "invoice", "payment", "transfer"]):
        score += 1
        reasons.append("Finance role: payment requests should be verified out-of-band.")
    if role == "HR" and any(p in t for p in ["resume", "candidate", "payroll", "w-2", "benefits"]):
        score += 1
        reasons.append("HR role: attachments/payroll requests are common lures.")
    if role == "IT/Engineering" and any(p in t for p in ["mfa", "password reset", "vpn", "admin", "security patch"]):
        score += 1
        reasons.append("IT role: MFA-fatigue and reset scams are common.")

    if score >= 6:
        risk = "High"
    elif score >= 3:
        risk = "Medium"
    else:
        risk = "Low"

    return risk, score, reasons, urls

def build_report_message(risk: str, urls: list[str], snippet: str):
    now = datetime.now().strftime("%Y-%m-%d %I:%M %p")
    url_lines = "\n".join([f"- {u}" for u in urls[:5]]) if urls else "- (no links detected)"
    return f"""Subject: Possible phishing email report ({risk})

Hi IT/Security team,

I received an email that may be phishing. PhishCheck flagged it as: {risk}.
Time noted: {now}

Links found:
{url_lines}

Email snippet:
\"\"\"{snippet[:500]}\"\"\"

Please advise next steps. I have not clicked any links or opened attachments (unless otherwise noted).

Thank you,
"""

# ---- Analysis ----
st.subheader("Analyze")

if "analyzed" not in st.session_state:
    st.session_state.analyzed = False

def run_analysis():
    text = st.session_state.get("email_text_input", "")
    if not text.strip():
        st.session_state.analysis_error = "Paste some email text (or upload a .txt file) first."
        return
    role_local = st.session_state.get("role", "General")
    # If Hugging Face LLM is enabled and token/model provided, call it
    if use_hf and hf_token and hf_model:
        try:
            llm_out = llm_score_email_hf(text, hf_token, hf_model, redact=hf_redact)
            # Expecting dict with keys: risk, score, reasons, draft_report
            risk = llm_out.get("risk", "Low")
            try:
                score = int(llm_out.get("score", 0))
            except Exception:
                score = 0
            reasons = llm_out.get("reasons", []) or []
            draft = llm_out.get("draft_report", llm_out.get("draft", ""))
            urls = extract_urls(text)
            st.session_state.analysis_error = ""
            st.session_state.analyzed = True
            st.session_state.analysis = {
                "risk": risk,
                "score": score,
                "reasons": reasons,
                "urls": urls,
                "report": draft or build_report_message(risk, urls, text),
                "snippet": text[:1000],
            }
            return
        except Exception as e:
            st.warning(f"Hugging Face LLM call failed: {e}. Falling back to rule-based analysis.")

    # fallback to rule-based scoring
    risk, score, reasons, urls = score_email(text, role_local)
    st.session_state.analysis_error = ""
    st.session_state.analyzed = True
    st.session_state.analysis = {
        "risk": risk,
        "score": score,
        "reasons": reasons,
        "urls": urls,
        "report": build_report_message(risk, urls, text),
        "snippet": text[:1000],
    }

if st.button("Analyze", type="primary", on_click=run_analysis):
    st.rerun()

if st.session_state.get("analysis_error"):
    st.warning(st.session_state.analysis_error)

if st.session_state.analyzed:
    analysis = st.session_state.analysis
    risk = analysis["risk"]
    score = analysis["score"]
    reasons = analysis["reasons"]
    urls = analysis["urls"]

    st.markdown("---")
    cols = st.columns([1, 2])
    with cols[0]:
        if risk == "High":
            st.error(f"Risk: {risk}")
        elif risk == "Medium":
            st.warning(f"Risk: {risk}")
        else:
            st.success(f"Risk: {risk}")
        st.metric("Score", score)
    with cols[1]:
        st.markdown("**Summary**")
        st.write(f"Detected {len(urls)} link(s).")
        st.write("Role-aware checks applied.")

    with st.expander("Why it was flagged (details)"):
        for r in reasons:
            st.write("• " + r)

    with st.expander("What to do next (safe guidance)"):
        st.write("• Don’t click links or open attachments if unsure.")
        st.write("• Verify the request using a trusted channel (call/Slack).")
        st.write("• Use your organization’s **Report Phish** process if available.")
        st.write("• If you already clicked, notify IT immediately.")

    if urls:
        with st.expander("Links found"):
            for u in urls[:20]:
                st.code(u)

    st.markdown("### Report / Draft (copy/paste)")
    st.text_area("Report message", value=analysis["report"], height=220, key="report_area")
    # allow downloading the draft/report
    try:
        report_bytes = analysis["report"].encode("utf-8")
        st.download_button("Download report (.txt)", data=report_bytes, file_name="phish_report.txt", mime="text/plain")
    except Exception:
        pass

    # # Single draft generator button (keeps only one report/draft area)
    # if st.button("Create draft from analysis"):
    #     analysis_local = st.session_state.analysis if st.session_state.get("analyzed") else None
    #     if analysis_local is None:
    #         st.info("Analyze an email first to generate a context-aware draft.")
    #     else:
    #         draft_templates = [
    #             "I received an email that appears suspicious. It contained {nlinks} link(s) and the following concerns: {concerns}. Please advise.",
    #             "Possible phishing: {concerns}. Links detected: {nlinks}. I haven't clicked any links. Please review.",
    #             "Alert: automated triage marked this {risk} risk. Summary: {concerns}. Links: {nlinks}. Requesting next steps.",
    #         ]
    #         concerns = ", ".join(analysis_local["reasons"][:4]) if analysis_local["reasons"] else "No clear flags"
    #         draft = random.choice(draft_templates).format(nlinks=len(analysis_local["urls"]), concerns=concerns, risk=analysis_local["risk"])
    #         full_draft = draft + "\n\n" + analysis_local["report"]
    #         st.session_state.analysis["report"] = full_draft
    #         st.success("Draft generated and applied to the Report area.")
    #         st.rerun()

# ---- Chat-style follow-up (rule-based assistant) ----
st.markdown("---")
st.subheader("AI Assistant")
if "chat" not in st.session_state:
    st.session_state.chat = []

# Simple template-based assistant to avoid identical answers every time
def generate_ai_response(question: str, analysis: Optional[dict]):
    ql = question.lower()
    templates = []
    if analysis is None:
        templates = [
            "I can help, but first please analyze the email using the Analyze button.",
            "Please run the analysis so I can provide detailed guidance based on detected flags.",
        ]
        return random.choice(templates)

    risk = analysis["risk"]
    reasons = analysis["reasons"]
    # varied phrasing
    intro = random.choice([
        f"Based on the analysis, this looks like a {risk} risk.",
        f"The automated triage rates this as {risk}.",
        f"Triage result: {risk} risk detected.",
    ])

    if "red flag" in ql or "why" in ql:
        body = "Top red flags:\n- " + "\n- ".join(reasons[:6]) if reasons else "No clear red flags found."
    elif "safe" in ql or "click" in ql:
        if risk == "High":
            body = random.choice([
                "Do not click links. Verify the sender by phone and report to IT.",
                "Avoid interacting with the message. Treat it as malicious until proven otherwise.",
            ])
        elif risk == "Medium":
            body = "Treat as suspicious — verify independently before clicking any links."
        else:
            body = "Lower risk, but still validate sender and hover links before clicking."
    elif "finance" in ql:
        body = "Finance checklist: confirm changes via phone, never follow emailed bank-change instructions, and verify invoices with originators."
    else:
        body = f"If you tell me the action you're considering, I can suggest a safe next step. Current score: {analysis['score']}."

    closing = random.choice([
        "If in doubt, escalate to your security team.",
        "When uncertain, assume it's suspicious and verify out-of-band.",
    ])

    return f"{intro}\n\n{body}\n\n{closing}"

enable_ai = st.checkbox("Enable AI assistant (simple, local)")
if enable_ai:
    user_q = st.text_input("Ask the assistant a question (e.g. 'What are the red flags?')", key="user_q")
    if st.button("Ask AI"):
        analysis = st.session_state.analysis if st.session_state.get("analyzed") else None
        if not user_q.strip():
            st.info("Type a question first.")
        else:
            ans = generate_ai_response(user_q.strip(), analysis)
            st.session_state.chat.append(("You", user_q.strip()))
            st.session_state.chat.append(("AI Assistant", ans))

    for speaker, msg in st.session_state.chat[-10:]:
        st.markdown(f"**{speaker}:** {msg}")

    # Draft generator (simple) — a small step toward Google Docs integration
    # (Draft generator moved to analysis section to keep a single Report/Draft area)