#!/usr/share/pyarc-venv/bin/python3
import Milter
import dkim
import authres
import sys
import configparser
import os
import re
import logging
from datetime import datetime
from io import BytesIO

# Erhöht den DNS-Timeout für dkimpy bei großen Keys
dkim.dns_timeout = 15

CONFIG_PATH = "/etc/pyarc/milter.conf"
config = configparser.ConfigParser()

if not os.path.exists(CONFIG_PATH):
    print(f"Fehler: Konfigurationsdatei nicht gefunden unter {CONFIG_PATH}", file=sys.stderr)
    sys.exit(1)

config.read(CONFIG_PATH)

# --- Logging Setup ---
LOG_ENABLED = config.getboolean("logging", "file_logging", fallback=False)
LOG_FILE = config.get("logging", "log_file", fallback="/var/log/pyarc/pyarc.log")

logger = logging.getLogger("pyarc_milter")
logger.setLevel(logging.INFO)

stdout_handler = logging.StreamHandler(sys.stdout)
stdout_formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
stdout_handler.setFormatter(stdout_formatter)
logger.addHandler(stdout_handler)

if LOG_ENABLED:
    log_dir = os.path.dirname(LOG_FILE)
    if not os.path.exists(log_dir):
        try:
            os.makedirs(log_dir, exist_ok=True)
        except Exception as e:
            logger.error(f"Konnte Log-Verzeichnis {log_dir} nicht erstellen: {e}")
    
    if os.path.exists(log_dir):
        try:
            file_handler = logging.FileHandler(LOG_FILE)
            file_formatter = logging.Formatter('%(asctime)s - %(levelname)s - [%(process)d] - %(message)s')
            file_handler.setFormatter(file_formatter)
            logger.addHandler(file_handler)
        except Exception as e:
            logger.error(f"Konnte Log-Datei {LOG_FILE} nicht öffnen: {e}")

# --- Allgemeine Config laden ---
AUTH_SERV_ID = config.get("general", "auth_serv_id", fallback="mx01.domain.tld")
LISTEN_SOCKET = config.get("general", "listen_socket", fallback="inet:8899@127.0.0.1")
REJECT_ON_FAIL = config.getboolean("validation", "reject_on_fail", fallback=False)

DOMAIN_REGEX = re.compile(r'@([a-zA-Z0-9.-]+\.[a-zA-Z]{2,})')


class ArcMilter(Milter.Base):
    def __init__(self):
        self.id = Milter.uniqueID()
        self.headers = []
        self.body_buffer = BytesIO()
        self.from_domain = None

    def header(self, name, hval):
        self.headers.append((name.encode('utf-8'), hval.encode('utf-8')))
        if name.lower() == 'from':
            match = DOMAIN_REGEX.search(hval)
            if match:
                self.from_domain = match.group(1).lower().strip()
        return Milter.CONTINUE

    def body(self, chunk):
        self.body_buffer.write(chunk)
        return Milter.CONTINUE

    def eom(self):
        self.body_buffer.seek(0)
        full_mail = b"".join([b"%s: %s\r\n" % (k, v) for k, v in self.headers]) + b"\r\n" + self.body_buffer.read()
        
        try:
            # 1. Vorhandene X-ARC Header aus früheren Stationen entfernen
            idx = 1
            while True:
                try:
                    self.changeheader('X-ARC', idx, '')
                    idx += 1
                except:
                    break

            # --- 1. Eingehend: ARC Validierung ---
            arc_verifier = dkim.ARC(full_mail)
            cv, results, comment = arc_verifier.verify()
            cv_str = cv.decode('utf-8') if isinstance(cv, bytes) else str(cv)
            
            logger.info(f"[{self.id}] Validierung für Domain '{self.from_domain}': CV={cv_str} ({comment})")
            self.addheader('Authentication-Results', f"{AUTH_SERV_ID}; arc={cv_str}")
            
            is_arc_valid = (cv_str in ["pass", "none"] and "fail" not in comment.lower())

            if cv_str == "fail" and REJECT_ON_FAIL:
                logger.warning(f"[{self.id}] Mail abgewiesen: ARC Validation fehlgeschlagen.")
                self.setreply('550', '5.7.1', 'ARC validation failed')
                return Milter.REJECT

            # --- 2. Ausgehend / First-Hop: Dynamische ARC Signierung ---
            sign_domain = None
            selector = None
            key_path = None

            if self.from_domain and config.has_section(self.from_domain):
                sign_domain = self.from_domain
                selector = config.get(self.from_domain, "selector")
                key_path = config.get(self.from_domain, "private_key_path")
            
            elif cv_str == "none":
                main_domain = ".".join(AUTH_SERV_ID.split(".")[-2:])
                if config.has_section(main_domain):
                    sign_domain = main_domain
                    selector = config.get(main_domain, "selector")
                    key_path = config.get(main_domain, "private_key_path")
                    logger.info(f"[{self.id}] Keine ARC-Kette vorhanden (CV=none). Starte neue Kette für lokale Hauptdomain {main_domain}.")

            if sign_domain and key_path and os.path.exists(key_path):
                with open(key_path, "rb") as f:
                    private_key = f.read()

                sig_headers = dkim.arc_sign(
                    full_mail,
                    selector=selector.encode('utf-8'),
                    domain=sign_domain.encode('utf-8'),
                    privkey=private_key,
                    srv_id=AUTH_SERV_ID.encode('utf-8')
                )

                for header_line in sig_headers:
                    if b':' in header_line:
                        # 1. Header in Name und Wert splitten
                        name_bytes, val_bytes = header_line.split(b':', 1)
                        name = name_bytes.decode('utf-8').strip()
                        
                        # 2. Den Wert komplett von allen Umbrüchen und doppelten Leerzeichen befreien (linearisieren)
                        val_clean = " ".join(val_bytes.decode('utf-8').split())
                        
                        # 3. Das Folding selbst bauen: Alle 72 Zeichen sauber mit Umbruch + Tab trennen (Rspamd-Style)
                        # Wir splitten den langen String in verdauliche Häppchen
                        chunks = []
                        # Der erste Chunk darf etwas länger sein, da der Headername schon links steht
                        remaining = val_clean
                        
                        # Wir brechen den Text sauber in Zeilen auf
                        while len(remaining) > 72:
                            # Suche das letzte Leerzeichen innerhalb der ersten 72 Zeichen für einen schönen Umbruch
                            split_idx = remaining.rfind(' ', 0, 72)
                            if split_idx == -1:  # Kein Leerzeichen gefunden (z.B. mitten im Base64-Blob)
                                split_idx = 72
                            
                            chunks.append(remaining[:split_idx].strip())
                            remaining = remaining[split_idx:].strip()
                        if remaining:
                            chunks.append(remaining)
                        
                        # Die Zeilen mit einem harten Umbruch und einem führenden TAB (\n\t) für Postfix verbinden
                        folded_value = "\n\t".join(chunks)
                        
                        # Übergabe an Postfix libmilter
                        self.addheader(name, folded_value)

                logger.info(f"[{self.id}] Mail erfolgreich ARC-versiegelt (Domain: {sign_domain}, Selector: {selector}).")
                is_arc_valid = True
            elif sign_domain:
                logger.error(f"[{self.id}] Key für {sign_domain} nicht gefunden unter {key_path}")

            # --- 3. Einmaliges Setzen des X-ARC Headers ---
            already_added = any(h[0].lower() == b'x-arc' for h in self.headers)
            if not already_added:
                self.addheader('X-ARC', 'TRUE' if is_arc_valid else 'FALSE')

        except Exception as e:
            logger.error(f"[{self.id}] Fehler beim ARC-Processing: {e}", exc_info=True)

        return Milter.CONTINUE

def main():
    Milter.factory = ArcMilter
    logger.info(f"Multidomain ARC Milter startet auf {LISTEN_SOCKET}...")
    Milter.runmilter("ArcMilter", LISTEN_SOCKET, timeout=30)

if __name__ == "__main__":
    main()
root@mail:~# nano /usr/local/bin/pyarc-milter 
root@mail:~# syste^C
root@mail:~# systemctl restart pyarc-milter.service 
root@mail:~# cat /usr/local/bin/pyarc-milter 
#!/usr/share/pyarc-venv/bin/python3
import Milter
import dkim
import authres
import sys
import configparser
import os
import re
import logging
from datetime import datetime
from io import BytesIO

# Erhöht den DNS-Timeout für dkimpy bei großen Keys
dkim.dns_timeout = 15

CONFIG_PATH = "/etc/pyarc/milter.conf"
config = configparser.ConfigParser()

if not os.path.exists(CONFIG_PATH):
    print(f"Fehler: Konfigurationsdatei nicht gefunden unter {CONFIG_PATH}", file=sys.stderr)
    sys.exit(1)

config.read(CONFIG_PATH)

# --- Logging Setup ---
LOG_ENABLED = config.getboolean("logging", "file_logging", fallback=False)
LOG_FILE = config.get("logging", "log_file", fallback="/var/log/pyarc/pyarc.log")

logger = logging.getLogger("pyarc_milter")
logger.setLevel(logging.INFO)

stdout_handler = logging.StreamHandler(sys.stdout)
stdout_formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
stdout_handler.setFormatter(stdout_formatter)
logger.addHandler(stdout_handler)

if LOG_ENABLED:
    log_dir = os.path.dirname(LOG_FILE)
    if not os.path.exists(log_dir):
        try:
            os.makedirs(log_dir, exist_ok=True)
        except Exception as e:
            logger.error(f"Konnte Log-Verzeichnis {log_dir} nicht erstellen: {e}")
    
    if os.path.exists(log_dir):
        try:
            file_handler = logging.FileHandler(LOG_FILE)
            file_formatter = logging.Formatter('%(asctime)s - %(levelname)s - [%(process)d] - %(message)s')
            file_handler.setFormatter(file_formatter)
            logger.addHandler(file_handler)
        except Exception as e:
            logger.error(f"Konnte Log-Datei {LOG_FILE} nicht öffnen: {e}")

# --- Allgemeine Config laden ---
AUTH_SERV_ID = config.get("general", "auth_serv_id", fallback="mx01.domain.tld")
LISTEN_SOCKET = config.get("general", "listen_socket", fallback="inet:8899@127.0.0.1")
REJECT_ON_FAIL = config.getboolean("validation", "reject_on_fail", fallback=False)

DOMAIN_REGEX = re.compile(r'@([a-zA-Z0-9.-]+\.[a-zA-Z]{2,})')


class ArcMilter(Milter.Base):
    def __init__(self):
        self.id = Milter.uniqueID()
        self.headers = []
        self.body_buffer = BytesIO()
        self.from_domain = None

    def header(self, name, hval):
        self.headers.append((name.encode('utf-8'), hval.encode('utf-8')))
        if name.lower() == 'from':
            match = DOMAIN_REGEX.search(hval)
            if match:
                self.from_domain = match.group(1).lower().strip()
        return Milter.CONTINUE

    def body(self, chunk):
        self.body_buffer.write(chunk)
        return Milter.CONTINUE

    def eom(self):
        self.body_buffer.seek(0)
        full_mail = b"".join([b"%s: %s\r\n" % (k, v) for k, v in self.headers]) + b"\r\n" + self.body_buffer.read()
        
        try:
            # 1. Vorhandene X-ARC Header aus früheren Stationen entfernen
            idx = 1
            while True:
                try:
                    self.changeheader('X-ARC', idx, '')
                    idx += 1
                except:
                    break

            # --- 1. Eingehend: ARC Validierung ---
            arc_verifier = dkim.ARC(full_mail)
            cv, results, comment = arc_verifier.verify()
            cv_str = cv.decode('utf-8') if isinstance(cv, bytes) else str(cv)
            
            logger.info(f"[{self.id}] Validierung für Domain '{self.from_domain}': CV={cv_str} ({comment})")
            self.addheader('Authentication-Results', f"{AUTH_SERV_ID}; arc={cv_str}")
            
            is_arc_valid = (cv_str in ["pass", "none"] and "fail" not in comment.lower())

            if cv_str == "fail" and REJECT_ON_FAIL:
                logger.warning(f"[{self.id}] Mail abgewiesen: ARC Validation fehlgeschlagen.")
                self.setreply('550', '5.7.1', 'ARC validation failed')
                return Milter.REJECT

            # --- 2. Ausgehend / Inbound-Forwarder: Dynamische ARC Signierung ---
            sign_domain = None
            selector = None
            key_path = None

            # Fall A: Die Absender-Domain gehört uns (Echter lokaler Ausgang)
            if self.from_domain and config.has_section(self.from_domain):
                sign_domain = self.from_domain
                selector = config.get(self.from_domain, "selector")
                key_path = config.get(self.from_domain, "private_key_path")
            
            # Fall B: Externe Mail (Egal ob CV=none oder CV=pass). Wir stempeln unsere Instanz drauf!
            else:
                main_domain = ".".join(AUTH_SERV_ID.split(".")[-2:])
                if config.has_section(main_domain):
                    sign_domain = main_domain
                    selector = config.get(main_domain, "selector")
                    key_path = config.get(main_domain, "private_key_path")
                    if cv_str == "none":
                        logger.info(f"[{self.id}] Keine ARC-Kette vorhanden (CV=none). Starte neue Kette für lokale Hauptdomain {main_domain}.")
                    else:
                        logger.info(f"[{self.id}] Bestehende ARC-Kette gefunden (CV={cv_str}). Erhöhe Kette um lokale Hauptdomain {main_domain}.")

            if sign_domain and key_path and os.path.exists(key_path):
                with open(key_path, "rb") as f:
                    private_key = f.read()

                sig_headers = dkim.arc_sign(
                    full_mail,
                    selector=selector.encode('utf-8'),
                    domain=sign_domain.encode('utf-8'),
                    privkey=private_key,
                    srv_id=AUTH_SERV_ID.encode('utf-8')
                )

                for header_line in sig_headers:
                    if b':' in header_line:
                        name_bytes, val_bytes = header_line.split(b':', 1)
                        name = name_bytes.decode('utf-8').strip()
                        
                        val_clean = " ".join(val_bytes.decode('utf-8').split())
                        
                        chunks = []
                        remaining = val_clean
                        
                        while len(remaining) > 72:
                            split_idx = remaining.rfind(' ', 0, 72)
                            if split_idx == -1:
                                split_idx = 72
                            
                            chunks.append(remaining[:split_idx].strip())
                            remaining = remaining[split_idx:].strip()
                        if remaining:
                            chunks.append(remaining)
                        
                        folded_value = "\n\t".join(chunks)
                        self.addheader(name, folded_value)

                logger.info(f"[{self.id}] Mail erfolgreich ARC-versiegelt (Domain: {sign_domain}, Selector: {selector}).")
                is_arc_valid = True
            elif sign_domain:
                logger.error(f"[{self.id}] Key für {sign_domain} nicht gefunden unter {key_path}")

            # --- 3. Einmaliges Setzen des X-ARC Headers ---
            already_added = any(h[0].lower() == b'x-arc' for h in self.headers)
            if not already_added:
                self.addheader('X-ARC', 'TRUE' if is_arc_valid else 'FALSE')

        except Exception as e:
            logger.error(f"[{self.id}] Fehler beim ARC-Processing: {e}", exc_info=True)

        return Milter.CONTINUE

def main():
    Milter.factory = ArcMilter
    logger.info(f"Multidomain ARC Milter startet auf {LISTEN_SOCKET}...")
    Milter.runmilter("ArcMilter", LISTEN_SOCKET, timeout=30)

if __name__ == "__main__":
    main()