#!/usr/bin/env python3
"""DMARC aggregate report digester for advmfr.net. Fully local, stdlib only.

Fetches rua reports from the dmarc@advmfr.net mailbox over IMAP, parses the
XML aggregates (zip/gzip/xml attachments), and writes a human-readable weekly
Markdown summary into dmarc-reports/. Nothing leaves the machine.

Usage:
    python3 dmarc_report.py              # last 8 days (default)
    python3 dmarc_report.py --days 30
    python3 dmarc_report.py --selftest   # parse a bundled sample, no network

Credentials come from the environment or a .env file next to this script:
    DMARC_IMAP_HOST=arome.o2switch.net
    DMARC_IMAP_USER=dmarc@advmfr.net
    DMARC_IMAP_PASS=...
The .env file must never be committed or shared.
"""

from __future__ import annotations

import argparse
import email
import gzip
import imaplib
import io
import os
import socket
import sys
import xml.etree.ElementTree as ET
import zipfile
from collections import defaultdict
from datetime import date, datetime, timedelta, timezone
from pathlib import Path

HERE = Path(__file__).resolve().parent
OUT_DIR = HERE / "dmarc-reports"
ALIGN_TARGET = 98.0  # percent needed to recommend the next hardening step


def load_env() -> None:
    env_file = HERE / ".env"
    if not env_file.exists():
        return
    for line in env_file.read_text().splitlines():
        line = line.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        key, _, value = line.partition("=")
        os.environ.setdefault(key.strip(), value.strip())


def extract_xml_payloads(raw_message: bytes) -> list[bytes]:
    """Return every aggregate-report XML found in one email message."""
    payloads: list[bytes] = []
    msg = email.message_from_bytes(raw_message)
    for part in msg.walk():
        if part.get_content_maintype() == "multipart":
            continue
        data = part.get_payload(decode=True)
        if not data:
            continue
        name = (part.get_filename() or "").lower()
        ctype = part.get_content_type()
        try:
            if name.endswith(".zip") or ctype in ("application/zip", "application/x-zip-compressed"):
                with zipfile.ZipFile(io.BytesIO(data)) as zf:
                    for member in zf.namelist():
                        if member.lower().endswith(".xml"):
                            payloads.append(zf.read(member))
            elif name.endswith(".gz") or ctype in ("application/gzip", "application/x-gzip"):
                payloads.append(gzip.decompress(data))
            elif name.endswith(".xml") or ctype in ("text/xml", "application/xml"):
                payloads.append(data)
        except (zipfile.BadZipFile, gzip.BadGzipFile, OSError):
            print(f"  ! unreadable attachment skipped: {name or ctype}", file=sys.stderr)
    return payloads


def parse_aggregate(xml_bytes: bytes) -> list[dict]:
    """Flatten one aggregate report into row dicts (one per <record>)."""
    rows: list[dict] = []
    try:
        root = ET.fromstring(xml_bytes)
    except ET.ParseError:
        print("  ! XML parse error, report skipped", file=sys.stderr)
        return rows

    def text(node, path, default=""):
        found = node.find(path)
        return (found.text or default).strip() if found is not None and found.text else default

    org = text(root, "./report_metadata/org_name", "unknown-org")
    begin = text(root, "./report_metadata/date_range/begin", "0")
    policy_domain = text(root, "./policy_published/domain", "?")
    for record in root.findall("./record"):
        rows.append(
            {
                "org": org,
                "domain": policy_domain,
                "date": datetime.fromtimestamp(int(begin or 0), tz=timezone.utc).date().isoformat(),
                "source_ip": text(record, "./row/source_ip", "?"),
                "count": int(text(record, "./row/count", "0") or 0),
                "disposition": text(record, "./row/policy_evaluated/disposition", "none"),
                "spf_aligned": text(record, "./row/policy_evaluated/spf", "fail") == "pass",
                "dkim_aligned": text(record, "./row/policy_evaluated/dkim", "fail") == "pass",
                "header_from": text(record, "./identifiers/header_from", "?"),
            }
        )
    return rows


def reverse_dns(ip: str) -> str:
    try:
        return socket.gethostbyaddr(ip)[0]
    except (socket.herror, socket.gaierror, OSError):
        return "no-rdns"


def fetch_rows(days: int) -> tuple[list[dict], int]:
    host = os.environ.get("DMARC_IMAP_HOST", "arome.o2switch.net")
    user = os.environ.get("DMARC_IMAP_USER", "dmarc@advmfr.net")
    password = os.environ.get("DMARC_IMAP_PASS", "")
    if not password:
        sys.exit("DMARC_IMAP_PASS missing (environment or .env next to the script).")

    since = (date.today() - timedelta(days=days)).strftime("%d-%b-%Y")
    rows: list[dict] = []
    with imaplib.IMAP4_SSL(host) as imap:
        imap.login(user, password)
        imap.select("INBOX", readonly=True)  # PEEK only, nothing is altered
        status, data = imap.search(None, f"(SINCE {since})")
        ids = data[0].split() if status == "OK" and data and data[0] else []
        print(f"{len(ids)} report email(s) since {since}")
        for msg_id in ids:
            status, fetched = imap.fetch(msg_id, "(BODY.PEEK[])")
            if status != "OK" or not fetched or fetched[0] is None:
                continue
            for xml_bytes in extract_xml_payloads(fetched[0][1]):
                rows.extend(parse_aggregate(xml_bytes))
    return rows, len(ids)


def build_report(rows: list[dict], days: int, mail_count: int) -> str:
    today = date.today().isoformat()
    lines = [
        f"# DMARC weekly digest  advmfr.net  {today}",
        "",
        f"Window: last {days} days · {mail_count} report email(s) · "
        f"{len(rows)} source rows · generated locally by dmarc_report.py",
        "",
    ]
    if not rows:
        lines += ["No aggregate reports found in the window. If the record was just",
                  "published, receivers usually start sending within 24-48 h."]
        return "\n".join(lines)

    domains = sorted({r["domain"] for r in rows})
    for domain in domains:
        drows = [r for r in rows if r["domain"] == domain]
        total = sum(r["count"] for r in drows)
        aligned = sum(r["count"] for r in drows if r["spf_aligned"] or r["dkim_aligned"])
        pct = 100.0 * aligned / total if total else 0.0

        per_ip: dict[str, dict] = defaultdict(lambda: {"count": 0, "aligned": 0, "orgs": set()})
        for r in drows:
            slot = per_ip[r["source_ip"]]
            slot["count"] += r["count"]
            slot["aligned"] += r["count"] if (r["spf_aligned"] or r["dkim_aligned"]) else 0
            slot["orgs"].add(r["org"])

        lines += [
            f"## {domain}  verdict: {pct:.1f}% aligned ({aligned}/{total} messages)",
            "",
            (f"Target reached (>= {ALIGN_TARGET}%). If the previous digest also passed, "
             "the next hardening step can be taken (see runbook step 6)."
             if pct >= ALIGN_TARGET
             else f"Below the {ALIGN_TARGET}% target: identify the failing sources below "
                  "before any hardening."),
            "",
            "### Sending sources",
            "",
            "| Source IP | rDNS | Messages | Aligned | Reported by |",
            "|---|---|---|---|---|",
        ]
        for ip, slot in sorted(per_ip.items(), key=lambda kv: -kv[1]["count"]):
            rd = reverse_dns(ip)
            flag = "" if slot["aligned"] == slot["count"] else "  ⚠️"
            lines.append(
                f"| {ip} | {rd} | {slot['count']} | "
                f"{slot['aligned']}/{slot['count']}{flag} | {', '.join(sorted(slot['orgs']))} |"
            )

        failing = {ip: s for ip, s in per_ip.items() if s["aligned"] < s["count"]}
        lines += ["", "### Sources failing alignment", ""]
        if failing:
            lines += [
                f"Each IP below sent mail claiming to be {domain} without passing "
                "SPF or DKIM alignment. Legitimate sender we forgot, or spoofing:",
                "",
            ]
            for ip, slot in sorted(failing.items(), key=lambda kv: -kv[1]["count"]):
                lines.append(f"- {ip} ({reverse_dns(ip)}): "
                             f"{slot['count'] - slot['aligned']} failing message(s)")
        else:
            lines.append("None. Every reported source aligned.")
        lines.append("")
    return "\n".join(lines)


SAMPLE_XML = b"""<?xml version=\"1.0\"?>
<feedback>
  <report_metadata><org_name>google.com</org_name>
    <date_range><begin>1755388800</begin><end>1755475200</end></date_range>
  </report_metadata>
  <policy_published><domain>advmfr.net</domain><p>none</p></policy_published>
  <record>
    <row><source_ip>109.234.162.168</source_ip><count>42</count>
      <policy_evaluated><disposition>none</disposition><dkim>pass</dkim><spf>pass</spf></policy_evaluated>
    </row>
    <identifiers><header_from>advmfr.net</header_from></identifiers>
  </record>
  <record>
    <row><source_ip>203.0.113.9</source_ip><count>7</count>
      <policy_evaluated><disposition>none</disposition><dkim>fail</dkim><spf>fail</spf></policy_evaluated>
    </row>
    <identifiers><header_from>advmfr.net</header_from></identifiers>
  </record>
</feedback>"""


def main() -> None:
    parser = argparse.ArgumentParser(description="Local DMARC digest for advmfr.net")
    parser.add_argument("--days", type=int, default=8)
    parser.add_argument("--selftest", action="store_true",
                        help="parse the bundled sample report, no network access")
    args = parser.parse_args()

    if args.selftest:
        rows, mail_count = parse_aggregate(SAMPLE_XML), 1
    else:
        load_env()
        rows, mail_count = fetch_rows(args.days)

    report = build_report(rows, args.days, mail_count)
    OUT_DIR.mkdir(exist_ok=True)
    out_file = OUT_DIR / f"dmarc-{date.today().isoformat()}.md"
    out_file.write_text(report)
    print(report)
    print(f"\nSaved: {out_file}")


if __name__ == "__main__":
    main()