#!/usr/bin/env python3
"""
BastionGuard Mail Security — Native Host

Protocol:
  { "type": "get-config" }
    -> { "ok": true, "config": {...} }

  {
    "type": "scan-email",
    "body": "...",
    "links": ["url1", ...],
    "attachments": [
      {
        "name": "invoice.zip",
        "content_b64": "...",
        "content_type": "application/zip",
        "size": 184
      }
    ]
  }
    -> {
         "ok": true,
         "clean": true/false,
         "threats": [...],
         "phishing_links": [...],
         "links_checked": 0,
         "attachments_checked": 0,
         "message_keys": [...]
       }
"""

import json
import os
import struct
import sys
import base64
import mimetypes
import socket
import urllib.request
import urllib.parse
import re
from concurrent.futures import ThreadPoolExecutor, as_completed

MAIL_CONFIG_PATH = os.path.expanduser("~/.config/BastionGuard/mail.json")
CLAMD_SOCKET = "/var/run/clamav/clamd.ctl"
PHISHING_API = "https://bastionguard.eu/bastionguard-security-intelligence/?q="
PHISHING_TIMEOUT = 8
MAX_LINKS = 10
MAX_WORKERS = 4
MAX_ATTACHMENTS = 10


# ============================================================
# Native Messaging protocol
# ============================================================
def read_message():
    raw = sys.stdin.buffer.read(4)
    if not raw or len(raw) < 4:
        return None
    length = struct.unpack("=I", raw)[0]
    payload = sys.stdin.buffer.read(length)
    return json.loads(payload.decode("utf-8"))


def send_message(data):
    encoded = json.dumps(data, ensure_ascii=False).encode("utf-8")
    sys.stdout.buffer.write(struct.pack("=I", len(encoded)))
    sys.stdout.buffer.write(encoded)
    sys.stdout.buffer.flush()


# ============================================================
# Logo -> data URI
# ============================================================
def logo_to_data_uri(path):
    if not path or not os.path.isfile(path):
        return ""
    try:
        mime, _ = mimetypes.guess_type(path)
        mime = mime or "image/png"
        with open(path, "rb") as f:
            data = base64.b64encode(f.read()).decode("ascii")
        return f"data:{mime};base64,{data}"
    except Exception:
        return ""


def enrich_config(config):
    for profile in config.get("profiles", []):
        sig = profile.get("signature", {})
        lp = sig.get("logo_path", "")

        if lp and not lp.startswith("data:") and not lp.startswith("http"):
            sig["logo"] = logo_to_data_uri(lp)
        elif lp.startswith("http"):
            sig["logo"] = lp
        else:
            sig.setdefault("logo", "")
    return config


def load_mail_config():
    with open(MAIL_CONFIG_PATH, "r", encoding="utf-8") as f:
        config = json.load(f)
    return enrich_config(config)


# ============================================================
# ClamAV via Unix socket
# ============================================================
def clamd_scan_bytes(data: bytes) -> dict:
    try:
        with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as s:
            s.settimeout(30)
            s.connect(CLAMD_SOCKET)

            s.sendall(b"zINSTREAM\0")

            chunk_size = 4096
            offset = 0
            while offset < len(data):
                chunk = data[offset:offset + chunk_size]
                s.sendall(struct.pack("!I", len(chunk)))
                s.sendall(chunk)
                offset += chunk_size

            s.sendall(struct.pack("!I", 0))

            response = b""
            while True:
                part = s.recv(4096)
                if not part:
                    break
                response += part
                if b"\0" in part or b"\n" in part:
                    break

        resp_str = response.decode("utf-8", errors="replace").strip().strip("\0")

        if "OK" in resp_str and "FOUND" not in resp_str:
            return {
                "clean": True,
                "threat_name": None,
                "threat_code": None,
                "note": None,
                "error": None,
            }

        if "FOUND" in resp_str:
            threat_name = resp_str.split(":")[-1].replace("FOUND", "").strip()
            return {
                "clean": False,
                "threat_name": threat_name,
                "threat_code": "clamav_threat_detected",
                "note": None,
                "error": None,
            }

        return {
            "clean": True,
            "threat_name": None,
            "threat_code": None,
            "note": resp_str,
            "error": None,
        }

    except Exception as e:
        return {
            "clean": True,
            "threat_name": None,
            "threat_code": None,
            "note": None,
            "error": str(e),
        }


def scan_email_body(body: str) -> dict:
    return clamd_scan_bytes(body.encode("utf-8"))


def scan_attachment(att: dict) -> dict:
    name = att.get("name", "attachment")
    content_b64 = att.get("content_b64", "")

    try:
        if not content_b64:
            return {
                "clean": True,
                "name": name,
                "threat_name": None,
                "threat_code": None,
                "error": "empty_attachment"
            }

        raw = base64.b64decode(content_b64, validate=False)
        result = clamd_scan_bytes(raw)

        return {
            "clean": result.get("clean", True),
            "name": name,
            "threat_name": result.get("threat_name"),
            "threat_code": result.get("threat_code"),
            "error": result.get("error"),
        }

    except Exception as e:
        return {
            "clean": True,
            "name": name,
            "threat_name": None,
            "threat_code": None,
            "error": str(e),
        }


# ============================================================
# Link extraction
# ============================================================
def extract_links(body: str) -> list:
    urls = re.findall(r'https?://[^\s"\'<>]+', body)

    seen = set()
    unique = []

    for u in urls:
        u = u.rstrip(".,;)")
        if u not in seen:
            seen.add(u)
            unique.append(u)

    return unique[:MAX_LINKS]


# ============================================================
# Phishing HTML parsing
# ============================================================
def parse_phishing_html(html: str) -> dict:
    card_m = re.search(
        r'<div\s+class="phishing-card\s+([^"]+)"\s+id="bg-result"',
        html,
        re.IGNORECASE,
    )
    card_classes = card_m.group(1) if card_m else ""

    if "malicious" in card_classes:
        classification_code = "malicious"
        is_phishing = True
    elif "warning" in card_classes:
        classification_code = "warning"
        is_phishing = False
    elif "clean" in card_classes:
        classification_code = "clean"
        is_phishing = False
    else:
        badge_m = re.search(
            r'<span\s+class="phishing-badge\s+([^"]+)">',
            html,
            re.IGNORECASE,
        )
        badge_cls = badge_m.group(1) if badge_m else ""

        if "malicious" in badge_cls:
            classification_code = "malicious"
            is_phishing = True
        elif "warning" in badge_cls:
            classification_code = "warning"
            is_phishing = False
        elif "clean" in badge_cls:
            classification_code = "clean"
            is_phishing = False
        else:
            classification_code = "unknown"
            is_phishing = False

    detail_m = re.search(
        r'<span\s+class="phishing-badge[^"]*">CLASSIFICATION</span>\s*([^<]+)',
        html,
        re.IGNORECASE,
    )
    detail = detail_m.group(1).strip() if detail_m else classification_code.upper()

    sev_m = re.search(
        r'<div\s+class="phishing-sev-level">([^<]+)</div>',
        html,
        re.IGNORECASE,
    )
    severity = sev_m.group(1).strip().lower() if sev_m else None

    if not severity:
        sev_cls_m = re.search(
            r'phishing-sev-(low|medium|high|crit)',
            html,
            re.IGNORECASE,
        )
        if sev_cls_m:
            severity = sev_cls_m.group(1).lower()

    return {
        "classification_code": classification_code,
        "detail": detail,
        "severity_code": severity,
        "is_phishing": is_phishing,
    }


def check_phishing_link(url: str) -> dict:
    try:
        api_url = PHISHING_API + urllib.parse.quote(url, safe="")
        req = urllib.request.Request(
            api_url,
            headers={
                "User-Agent": (
                    "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 "
                    "(KHTML, like Gecko) Chrome/124.0 Safari/537.36 "
                    "BastionGuard-MailProxy/1.0"
                )
            },
        )

        with urllib.request.urlopen(req, timeout=PHISHING_TIMEOUT) as resp:
            html = resp.read().decode("utf-8", errors="replace")

        result = parse_phishing_html(html)

        return {
            "url": url,
            "phishing": result["is_phishing"],
            "classification_code": result["classification_code"],
            "detail": result["detail"],
            "severity_code": result["severity_code"],
            "error_code": None,
            "error": None,
        }

    except Exception as e:
        return {
            "url": url,
            "phishing": False,
            "classification_code": "error",
            "detail": None,
            "severity_code": None,
            "error_code": "phishing_check_error",
            "error": str(e),
        }


def check_links_parallel(links: list) -> list:
    if not links:
        return []

    results = []
    with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
        futures = {executor.submit(check_phishing_link, url): url for url in links}
        for future in as_completed(futures):
            try:
                results.append(future.result())
            except Exception as e:
                results.append({
                    "url": futures[future],
                    "phishing": False,
                    "classification_code": "error",
                    "detail": None,
                    "severity_code": None,
                    "error_code": "phishing_check_error",
                    "error": str(e),
                })

    return results


# ============================================================
# scan-email handler
# ============================================================
def handle_scan_email(message: dict) -> dict:
    body = message.get("body", "") or ""
    links = message.get("links", []) or []
    attachments = message.get("attachments", []) or []

    threats = []
    phishing_links = []
    message_keys = []
    clean = True

    # 1. ClamAV sul body
    if body:
        clamav_result = scan_email_body(body)

        if not clamav_result.get("clean", True):
            clean = False
            if "threat_detected" not in message_keys:
                message_keys.append("threat_detected")

            threats.append({
                "engine": "clamav",
                "source": "body",
                "code": clamav_result.get("threat_code") or "clamav_threat_detected",
                "name": clamav_result.get("threat_name") or "unknown_threat",
            })

    # 2. Estrai link se non forniti
    if not links and body:
        links = extract_links(body)

    # 3. Controllo phishing link
    if links:
        link_results = check_links_parallel(links)

        for r in link_results:
            if r.get("phishing"):
                clean = False
                if "phishing_detected" not in message_keys:
                    message_keys.append("phishing_detected")

                phishing_links.append({
                    "url": r["url"],
                    "classification_code": r.get("classification_code", "malicious"),
                    "detail": r.get("detail", ""),
                    "severity_code": r.get("severity_code"),
                })

    # 4. ClamAV sugli allegati
    safe_attachments = []
    for att in attachments[:MAX_ATTACHMENTS]:
        if isinstance(att, dict):
            safe_attachments.append(att)

    for att in safe_attachments:
        result = scan_attachment(att)
        if not result.get("clean", True):
            clean = False
            if "threat_detected" not in message_keys:
                message_keys.append("threat_detected")

            threats.append({
                "engine": "clamav",
                "source": "attachment",
                "attachment_name": result.get("name", "attachment"),
                "code": result.get("threat_code") or "clamav_threat_detected",
                "name": result.get("threat_name") or "unknown_threat",
            })

    if clean:
        message_keys.append("scan_clean")

    return {
        "ok": True,
        "clean": clean,
        "threats": threats,
        "phishing_links": phishing_links,
        "links_checked": len(links),
        "attachments_checked": len(safe_attachments),
        "message_keys": message_keys,
    }


# ============================================================
# Main loop
# ============================================================
def main():
    while True:
        message = read_message()
        if message is None:
            break

        msg_type = message.get("type")

        if msg_type == "get-config":
            try:
                config = load_mail_config()
                send_message({"ok": True, "config": config})
            except Exception as e:
                send_message({
                    "ok": False,
                    "error_code": "config_load_failed",
                    "error": str(e),
                })

        elif msg_type == "scan-email":
            try:
                result = handle_scan_email(message)
                send_message(result)
            except Exception as e:
                send_message({
                    "ok": False,
                    "error_code": "scan_email_failed",
                    "error": str(e),
                })

        else:
            send_message({
                "ok": False,
                "error_code": "unknown_request",
                "error": f"unknown request: {msg_type}",
            })


if __name__ == "__main__":
    main()