#!/usr/bin/env python3
import os
import re
import signal
import logging
from email.utils import parseaddr
import Milter

ALLOWED_RCPT_DOMAINS_FILE = "/etc/postfix/policy/allowed_rcpt_domains.regexp"
ALLOWED_FROM_ADDR_FILE = "/etc/postfix/policy/allowed_external_senders"
BLOCKED_FROM_DOMAINS_FILE = "/etc/postfix/policy/blocked_header_from_domains.txt"
LOG_FILE = "/var/log/postfix-policy-milter/milter.log"


def load_allowed_rcpt_domain_regexes(path):
    regs = []
    if not os.path.exists(path):
        return regs
    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            s = line.strip()
            if not s or s.startswith("#"):
                continue
            m = re.match(r'^/(.*?)/\s+OK\b', s)
            if not m:
                continue
            try:
                regs.append(re.compile(m.group(1), re.IGNORECASE))
            except re.error:
                pass
    return regs


def load_allowed_from_addresses(path):
    out = set()
    if not os.path.exists(path):
        return out
    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            s = line.strip()
            if not s or s.startswith("#"):
                continue
            parts = s.split()
            if len(parts) >= 2 and parts[1].upper() == "OK":
                out.add(parts[0].lower())
    return out


def load_blocked_domains(path):
    out = set()
    if not os.path.exists(path):
        return out
    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            s = line.strip().lower()
            if not s or s.startswith("#"):
                continue
            out.add(s)
    return out


def extract_domain(addr):
    if not addr or "@" not in addr:
        return ""
    return addr.rsplit("@", 1)[1].lower()


class PolicyMilter(Milter.Base):
    allowed_rcpt_regexes = []
    allowed_from_addrs = set()
    blocked_from_domains = set()

    @classmethod
    def reload_config(cls):
        cls.allowed_rcpt_regexes = load_allowed_rcpt_domain_regexes(ALLOWED_RCPT_DOMAINS_FILE)
        cls.allowed_from_addrs = load_allowed_from_addresses(ALLOWED_FROM_ADDR_FILE)
        cls.blocked_from_domains = load_blocked_domains(BLOCKED_FROM_DOMAINS_FILE)

    def __init__(self):
        self.id = Milter.uniqueID()
        self.mail_from = ""
        self.rcpt_tos = []
        self.header_from_addr = ""
        self.header_from_domain = ""
        self.log = logging.getLogger("policy-milter")

    def envfrom(self, mailfrom, *args):
        self.mail_from = (mailfrom or "").strip("<>").lower()
        return Milter.CONTINUE

    def envrcpt(self, rcpt, *args):
        self.rcpt_tos.append((rcpt or "").strip("<>").lower())
        return Milter.CONTINUE

    def header(self, name, value):
        if name.lower() == "from" and not self.header_from_addr:
            _, addr = parseaddr(value or "")
            self.header_from_addr = (addr or "").lower()
            self.header_from_domain = extract_domain(self.header_from_addr)
        return Milter.CONTINUE

    def eom(self):
        # Rule 1: pass if any recipient domain is explicitly allowed.
        for rcpt in self.rcpt_tos:
            d = extract_domain(rcpt)
            for rgx in self.allowed_rcpt_regexes:
                if rgx.search("@" + d):
                    self.log.info("ACCEPT rule=1 rcpt=%s from_hdr=%s", rcpt, self.header_from_addr)
                    return Milter.ACCEPT

        # Rule 3: pass if header From address is explicitly allowed.
        if self.header_from_addr in self.allowed_from_addrs:
            self.log.info("ACCEPT rule=3 from_hdr=%s", self.header_from_addr)
            return Milter.ACCEPT

        # Rule 2: reject if header From domain is blocked.
        if self.header_from_domain in self.blocked_from_domains:
            self.log.warning(
                "REJECT rule=2 from_hdr=%s domain=%s mail_from=%s rcpt=%s",
                self.header_from_addr, self.header_from_domain, self.mail_from, ",".join(self.rcpt_tos)
            )
            self.setreply("550", "5.7.1", "Header From domain is not permitted.")
            return Milter.REJECT

        self.log.info("ACCEPT default from_hdr=%s rcpt=%s", self.header_from_addr, ",".join(self.rcpt_tos))
        return Milter.ACCEPT


def setup_logging():
    os.makedirs(os.path.dirname(LOG_FILE), exist_ok=True)
    logging.basicConfig(
        filename=LOG_FILE,
        level=logging.INFO,
        format="%(asctime)s %(levelname)s %(message)s"
    )


def main():
    setup_logging()
    PolicyMilter.reload_config()

    def hup(signum, frame):
        PolicyMilter.reload_config()
        logging.getLogger("policy-milter").info("Reloaded policy by SIGHUP")

    signal.signal(signal.SIGHUP, hup)

    Milter.factory = PolicyMilter
    Milter.runmilter("postfix-policy-milter", "unix:/run/postfix-policy-milter/milter.sock", 600)


if __name__ == "__main__":
    main()