gmail_processor.py

from __future__ import annotations

import base64
import json
import os
import re
import time
from datetime import datetime, timezone
from email.utils import parseaddr
from html import unescape
from typing import List, Tuple
from urllib.parse import urlparse

import requests

from alert_utils import send_email_alert
from db_helpers import insert_phishing_email, insert_ai_training_sample

URL_REGEX = re.compile(r"https?://[^\s<>\"]+")

DEFAULT_KEYWORD_WEIGHTS = {
    "urgent": 2,
    "immediately": 2,
    "password": 3,
    "verify": 3,
    "account": 2,
    "invoice": 3,
    "bank": 3,
    "click": 2,
    "secure": 2,
    "transfer": 3,
}

DEFAULT_AUTH_WEIGHTS = {
    "spf_fail": 2,
    "dkim_fail": 2,
    "dmarc_fail": 3,
}


def classify_with_ai(subject: str, sender: str, body: str, urls: List[str], debug: bool=False) -> dict:
    url = os.getenv('AI_CLASSIFIER_URL')
    token = os.getenv('AI_CLASSIFIER_TOKEN')
    if not url or not token:
        return {"label": "unknown", "confidence": 0.0}

    payload = {
        "subject": subject or "",
        "sender": sender or "",
        "body": body or "",
        "urls": urls or [],
    }
    headers = {
        "X-Auth-Token": token,
        "Content-Type": "application/json",
    }

    try:
        start = time.time()
        if debug:
            print(f"[DEBUG] AI request payload: {json.dumps(payload)}")
        response = requests.post(url, json=payload, headers=headers, timeout=10)
        response.raise_for_status()
        result = response.json() or {}
        result.setdefault("latency_ms", int((time.time() - start) * 1000))
        if debug:
            print(f"[DEBUG] AI response: {json.dumps(result)}")
        return result
    except Exception as exc:
        print(f"[!] AI classification error: {exc}")
        return {"label": "unknown", "confidence": 0.0}


def _decode_part(part):
    data = part.get('body', {}).get('data')
    if not data:
        return ''
    try:
        return base64.urlsafe_b64decode(data.encode('utf-8')).decode('utf-8', errors='replace')
    except Exception:
        return ''


def _collect_bodies(payload):
    texts = []
    if 'parts' in payload:
        for part in payload['parts']:
            texts.append(_collect_bodies(part))
    else:
        mime_type = payload.get('mimeType', '')
        if mime_type.startswith('text/'):
            texts.append(_decode_part(payload))
    return '\n'.join(filter(None, texts))


def _extract_urls(text):
    if not text:
        return []
    text = unescape(text)
    return URL_REGEX.findall(text)


def process_gmail_messages(gmail_service, config, since_dt: datetime) -> Tuple[int, int, datetime]:
    gmail_cfg = config.get('gmail', {})
    if not gmail_cfg.get('enabled'):
        return 0, 0, since_dt

    mailbox = gmail_cfg.get('mailbox', 'me')
    max_messages = int(gmail_cfg.get('max_messages_per_poll', 50))
    include_spam = gmail_cfg.get('include_spam', False)
    query = gmail_cfg.get('query', '')
    after_ts = int(since_dt.timestamp()) if since_dt else None
    if after_ts:
        query = (query + ' ' if query else '') + f'after:{after_ts}'

    list_kwargs = {
        'userId': mailbox,
        'maxResults': max_messages,
    }
    if query:
        list_kwargs['q'] = query.strip()
    if not include_spam:
        list_kwargs['labelIds'] = ['INBOX']

    response = gmail_service.users().messages().list(**list_kwargs).execute()
    message_ids = [m['id'] for m in response.get('messages', [])]

    scanned = 0
    flagged = 0
    latest_seen = since_dt

    high_risk_names = [name.lower() for name in gmail_cfg.get('high_risk_display_names', [])]
    allowed_sender_domains = [d.lower() for d in gmail_cfg.get('allowed_sender_domains', [])]
    trusted_file_domains = [d.lower() for d in gmail_cfg.get('trusted_file_domains', [])]
    share_link_domains = [d.lower() for d in gmail_cfg.get('share_link_domains', [])]
    urgency_keywords = [u.lower() for u in gmail_cfg.get('urgency_keywords', [])]
    financial_keywords = [f.lower() for f in gmail_cfg.get('financial_keywords', [])]
    ai_min_confidence = float(os.getenv('AI_MIN_CONFIDENCE', 0.7))
    combined_threshold = float(config.get('phishing_detection', {}).get('combined_confidence_threshold', 0.75))

    keyword_weights_cfg = config.get('phishing_detection', {}).get('keyword_weights', {})
    keyword_weights = {**DEFAULT_KEYWORD_WEIGHTS, **{k: int(v) for k, v in keyword_weights_cfg.items()}}
    auth_weights_cfg = config.get('phishing_detection', {}).get('auth_weights', {})
    auth_weights = {**DEFAULT_AUTH_WEIGHTS, **{k: int(v) for k, v in auth_weights_cfg.items()}}

    for message_id in message_ids:
        msg = gmail_service.users().messages().get(userId=mailbox, id=message_id, format='full').execute()
        internal_ts = datetime.fromtimestamp(int(msg.get('internalDate', 0)) / 1000, tz=timezone.utc)
        if internal_ts <= since_dt:
            continue

        scanned += 1
        if not latest_seen or internal_ts > latest_seen:
            latest_seen = internal_ts

        headers = msg.get('payload', {}).get('headers', [])
        header_dict = {h['name'].lower(): h['value'] for h in headers}

        subject = header_dict.get('subject', '(no subject)')
        sender_header = header_dict.get('from', '')
        sender_display, sender_email = parseaddr(sender_header)
        sender_email = sender_email.lower()
        sender_domain = sender_email.split('@')[-1] if '@' in sender_email else ''
        recipients = header_dict.get('to', '')
        auth_results = header_dict.get('authentication-results', '')

        alert_prefix = config.get('alerts', {}).get('alert_subject_prefix')
        if alert_prefix and subject.startswith(alert_prefix):
            # Skip emails generated by this agent
            continue

        alert_email = os.getenv('ALERT_EMAIL', '').lower()
        smtp_username = os.getenv('SMTP_USERNAME', '').lower()
        ignore_senders = set(s.lower() for s in config.get('gmail', {}).get('ignore_senders', []))
        if sender_email in {alert_email, smtp_username} | ignore_senders:
            continue

        body_text = _collect_bodies(msg.get('payload', {}))
        snippet = msg.get('snippet', '')

        urls = _extract_urls(body_text or snippet)
        suspicious_reasons = set()
        share_links = set()
        share_link_detected = False
        high_risk_triggered = False
        lookalike_detected = False

        # External share links
        for url in urls:
            parsed = urlparse(url)
            domain = parsed.netloc.lower()
            if not domain:
                continue
            if any(domain.endswith(td) for td in trusted_file_domains):
                continue
            if any(sd in domain for sd in share_link_domains):
                share_links.add(url)
                share_link_detected = True
                if sender_domain not in allowed_sender_domains:
                    suspicious_reasons.add(f"External file share link: {domain}")

        display_lower = sender_display.lower() if sender_display else ''
        if high_risk_names and any(term in display_lower for term in high_risk_names):
            high_risk_triggered = True
            if sender_domain and sender_domain not in allowed_sender_domains:
                suspicious_reasons.add(f"Display name '{sender_display}' matches high-risk role")

        if sender_domain and sender_domain not in allowed_sender_domains:
            for allowed in allowed_sender_domains:
                if allowed in sender_domain and sender_domain != allowed:
                    suspicious_reasons.add(f"Sender domain looks similar to trusted domain: {sender_domain}")
                    lookalike_detected = True
                    break

        auth_lower = auth_results.lower()
        keyword_score = 0
        lower_content = f"{subject} {body_text}".lower()
        for keyword, weight in keyword_weights.items():
            if keyword in lower_content:
                keyword_score += weight

        if high_risk_names:
            for term in high_risk_names:
                if term and term in lower_content:
                    high_risk_triggered = True
                    if sender_domain not in allowed_sender_domains:
                        suspicious_reasons.add(f"High-risk keyword '{term}' found in message content")

        outside_org_detected = 'outside your organization' in lower_content
        urgency_detected = any(term in lower_content for term in urgency_keywords)
        financial_detected = any(term in lower_content for term in financial_keywords)
        if urgency_detected and sender_domain not in allowed_sender_domains:
            suspicious_reasons.add("Urgent language detected from external sender")
        if financial_detected and sender_domain not in allowed_sender_domains:
            suspicious_reasons.add("Financial language detected from external sender")

        spf = re.search(r"spf=(\w+)", auth_lower)
        dkim = re.search(r"dkim=(\w+)", auth_lower)
        dmarc = re.search(r"dmarc=(\w+)", auth_lower)
        spf_fail = spf and spf.group(1) in ("fail", "softfail")
        dkim_fail = dkim and dkim.group(1) == "fail"
        dmarc_fail = dmarc and dmarc.group(1) == "fail"

        rule_score = keyword_score
        if spf_fail:
            suspicious_reasons.add('SPF failure detected')
            rule_score += auth_weights.get('spf_fail', DEFAULT_AUTH_WEIGHTS['spf_fail'])
        if dkim_fail:
            suspicious_reasons.add('DKIM failure detected')
            rule_score += auth_weights.get('dkim_fail', DEFAULT_AUTH_WEIGHTS['dkim_fail'])
        if dmarc_fail:
            suspicious_reasons.add('DMARC failure detected')
            rule_score += auth_weights.get('dmarc_fail', DEFAULT_AUTH_WEIGHTS['dmarc_fail'])

        ai_result = classify_with_ai(
            subject,
            sender_email,
            body_text or snippet,
            list(share_links),
            debug=config.get('log_level', '').upper() == 'DEBUG'
        )
        if config.get('phishing_detection', {}).get('train_ai'):
            insert_ai_training_sample(
                message_id,
                subject,
                sender_email,
                sender_domain,
                body_text or snippet,
                list(share_links),
                {
                    "subject": subject or "",
                    "sender": sender_email,
                    "body": body_text or snippet or "",
                    "urls": list(share_links)
                },
                ai_result
            )
        label = ai_result.get('label', 'unknown')
        LABEL_MAP = {
            "phishing": "phishing",
            "possible_phishing": "potential_phishing",
            "potential phishing": "potential_phishing",
            "potential_phishing": "potential_phishing",
            "unsafe": "phishing",
            "uncertain": "unknown",
            "unclassified": "unknown",
            "safe": "safe",
            "marketing": "marketing"
        }
        label = LABEL_MAP.get(label.strip().lower(), "unknown")
        ai_confidence = float(ai_result.get('confidence', 0.0) or 0.0)
        ai_model = ai_result.get('model', 'unknown')
        ai_latency = ai_result.get('latency_ms', 0)

        ai_flagged = False
        reason_ai = ""
        if label in ("phishing", "potential_phishing"):
            if ai_confidence >= ai_min_confidence:
                ai_flagged = True
                reason_ai = f"AI classified as {label} ({ai_confidence:.2f})"
        elif label in ("safe", "marketing"):
            reason_ai = f"AI classified as {label} ({ai_confidence:.2f})"
            ai_confidence = 0.0
        else:
            reason_ai = f"AI returned {label} ({ai_confidence:.2f})"
            ai_confidence = 0.0

        if reason_ai:
            suspicious_reasons.add(reason_ai)

        final_confidence = round((ai_confidence * 0.7) + ((rule_score / 10) * 0.3), 2)

        if label in ("safe", "marketing"):
            ai_confidence = 0.0
            final_confidence = round(((rule_score / 10) * 0.3), 2)
            should_flag = False
        elif label in ("phishing", "potential_phishing"):
            should_flag = ai_confidence >= ai_min_confidence or final_confidence > combined_threshold
        else:
            should_flag = final_confidence > combined_threshold or spf_fail or dkim_fail or dmarc_fail

        if not should_flag:
            if config.get('log_level', '').upper() == 'DEBUG':
                print(f"[DEBUG] Skipping email {message_id} - label={label}, ai_conf={ai_confidence:.2f}, score={rule_score}")
            continue

        message_time = internal_ts
        reasons_list = sorted(suspicious_reasons)
        share_links_list = sorted(share_links)

        inserted = insert_phishing_email(
            message_id=message_id,
            subject=subject[:255],
            sender_email=sender_email,
            sender_display=sender_display[:255] if sender_display else '',
            sender_domain=sender_domain,
            recipients=recipients,
            reasons=reasons_list,
            share_links=share_links_list,
            auth_results=auth_results,
            snippet=snippet,
            message_time=message_time,
            ai_label=label,
            ai_confidence=ai_confidence,
            rule_score=rule_score,
            phishing_confidence=final_confidence
        )

        if not inserted:
            continue

        flagged += 1
        if config.get('log_level', '').upper() == 'DEBUG':
            print(
                f"[DEBUG] Phishing email stored: {subject} ({message_id})"
                f" | [AI] label={label} conf={ai_confidence:.2f} model={ai_model}"
                f" | [Rules] {rule_score} pts | SPF fail={spf_fail} | DKIM fail={dkim_fail} | DMARC fail={dmarc_fail}"
                f" | Combined={final_confidence:.2f} | Latency={ai_latency}ms"
            )

        reasons_text = '\n'.join(f"  - {reason}" for reason in reasons_list) if reasons_list else '  - None'
        links_text = '\n'.join(f"  - {link}" for link in share_links_list) if share_links_list else '  - None'

        details = (
            f"Subject: {subject}\n"
            f"From: {sender_display} <{sender_email}>\n"
            f"To: {recipients}\n"
            f"Time: {message_time.strftime('%Y-%m-%d %H:%M:%S %Z')}\n"
            f"Reasons:\n{reasons_text}\n"
            f"Links:\n{links_text}\n"
            f"AI Confidence: {ai_confidence:.2f}\n"
            f"Rule Score: {rule_score}\n"
            f"Combined Confidence: {final_confidence:.2f}\n"
        )
        send_email_alert(
            f"{config['alerts']['alert_subject_prefix']} Potential Phishing Email: {subject}",
            details,
            config
        )

    return scanned, flagged, latest_seen or datetime.now(timezone.utc)