inference/bengu_view.py

"""Bengu-1 input view: raw RFC 822 mail bytes -> the exact text the Bengu-1 classifier reads.

Standalone (Python standard library only). The model was trained on this view: a header-facts line,
sender, Reply-To, subject, then the decoded body text (HTML reduced to text), at most about 12,000
characters of a JSON-serialised data block. Mail larger than 256 KiB is read from its headers only.

Mail the user sent or accepted is not marked: no acceptance or provenance facts reach the view (they
are fixed to "unknown" here, exactly as in training). Nothing is fetched; links are never followed.

    from bengu_view import view
    text = view(open("message.eml", "rb").read())
"""

import json
import re
import unicodedata
from email import policy
from email.parser import BytesParser
from email.utils import getaddresses
from html.parser import HTMLParser
import ipaddress

LIMIT = 8000                 # body and header characters kept from the parse
MAX_MESSAGE = 262144         # larger mail: headers only, body marked truncated
DATA_LIMIT = 12000           # characters of the serialised data block
FIELD_LIMIT = 300            # per header-summary string
DOMAIN = re.compile(r"[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?(?:\.[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?)+\Z")
LOCAL = re.compile(r"[a-z0-9.!#$%&'*+/=?^_`{|}~-]{1,64}\Z")
FREEMAIL = {"gmail.com", "googlemail.com", "yahoo.com", "hotmail.com", "outlook.com", "live.com", "icloud.com",
            "me.com", "mac.com", "aol.com", "yandex.com", "yandex.com.tr", "proton.me", "protonmail.com",
            "msn.com", "mail.com", "gmx.com", "hotmail.com.tr", "windowslive.com"}
NEUTRAL_AUTH = {"dkim_aligned_pass": False, "spf": "unknown", "dmarc": "unknown", "gmail_forwarded": False}


class _Unsafe(ValueError):
    pass


def clean(text, limit=LIMIT):
    text = unicodedata.normalize("NFKC", str(text))
    return "".join(c for c in text if c in "\n\t" or not unicodedata.category(c).startswith("C"))[:limit]


def _domain(value):
    value = value.lower()
    if not value.isascii() or "xn--" in value or len(value) > 253 or not DOMAIN.fullmatch(value):
        raise _Unsafe("domain")
    try:
        ipaddress.ip_address(value)
    except ValueError:
        return value
    raise _Unsafe("domain")


def _address(value):
    if not isinstance(value, str) or value.count("@") != 1 or any(c in value for c in "\r\n\x00"):
        raise _Unsafe("address")
    local, host = value.lower().split("@")
    if not LOCAL.fullmatch(local):
        raise _Unsafe("address")
    return local + "@" + _domain(host)


def _single_address(headers, name):
    values = headers.get_all(name, [])
    if len(values) != 1 or any(c in str(values[0]) for c in "\r\n\x00"):
        return None
    parsed = getaddresses([str(values[0])])
    if len(parsed) != 1:
        return None
    try:
        return _address(parsed[0][1])
    except _Unsafe:
        return None


class _TextHTML(HTMLParser):
    def __init__(self):
        super().__init__(convert_charrefs=True)
        self.hidden, self.parts, self.length = 0, [], 0

    def handle_starttag(self, tag, attrs):
        if tag in {"script", "style", "head", "template"}:
            self.hidden += 1
        elif tag in {"p", "br", "div", "li", "tr"}:
            self.handle_data("\n")

    def handle_endtag(self, tag):
        if tag in {"script", "style", "head", "template"} and self.hidden:
            self.hidden -= 1

    def handle_data(self, data):
        if not self.hidden and self.length < LIMIT:
            chunk = data[:LIMIT - self.length]
            self.parts.append(chunk)
            self.length += len(chunk)


def _header_bytes(raw):
    for sep in (b"\r\n\r\n", b"\n\n"):
        cut = raw.find(sep)
        if cut >= 0:
            return raw[:cut + len(sep)]
    return raw


def parse(raw: bytes):
    """(headers, sender, subject, text, bounded)"""
    oversized = len(raw) > MAX_MESSAGE
    if oversized:
        raw = _header_bytes(raw)
        if len(raw) > MAX_MESSAGE:
            raise ValueError("header section larger than 256 KiB")
    msg = BytesParser(policy=policy.default).parsebytes(raw)
    chunks, size, complete = [], 0, True
    for part in msg.walk():
        if part.is_multipart() or part.get_content_disposition() == "attachment" or part.get_filename():
            if not part.is_multipart():
                complete = False
            continue
        if part.get_content_type() not in {"text/plain", "text/html"}:
            complete = False
            continue
        payload = part.get_payload(decode=True) or b""
        try:
            text = payload.decode(part.get_content_charset() or "utf-8", errors="replace")
        except LookupError:
            text = payload.decode("utf-8", errors="replace")
        if part.get_content_type() == "text/html":
            parser = _TextHTML()
            parser.feed(text)
            text = "".join(parser.parts)
            if parser.length >= LIMIT:
                complete = False
        if len(text) > LIMIT - size:
            complete = False
        chunk = clean(text, LIMIT - size)
        chunks.append(chunk)
        size += len(chunk)
        if size >= LIMIT:
            complete = False
            break
    all_headers = "\n".join(f"{k}: {v}" for k, v in msg.items())
    if len(all_headers) > LIMIT or len(msg.get_all("Subject", [])) > 1:
        complete = False
    return (msg, _single_address(msg, "From"), clean(msg.get("Subject", ""), 500), clean("\n".join(chunks)),
            complete and not oversized and not msg.defects)


def _one_line(value, limit=FIELD_LIMIT):
    return " ".join(clean(value, limit * 2).split())[:limit]


def _single(headers, name):
    values = headers.get_all(name, [])
    return str(values[0]) if len(values) == 1 else ""


def _display_name(headers):
    values = headers.get_all("From", [])
    if len(values) != 1:
        return ""
    parsed = getaddresses([str(values[0])])
    return parsed[0][0] if len(parsed) == 1 else ""


def _addresses(headers, name, limit=10):
    found = []
    for value in headers.get_all(name, [])[:2]:
        for _, item in getaddresses([str(value)]):
            try:
                found.append(_address(item))
            except _Unsafe:
                continue
    return found[:limit]


def _attachments(headers):
    return sum(1 for part in headers.walk()
               if not part.is_multipart() and (part.get_content_disposition() == "attachment" or part.get_filename()))


def data_block(raw: bytes) -> str:
    """The serialised data block (header summary + body), at most DATA_LIMIT characters."""
    headers, sender, subject, body, bounded = parse(raw)
    name = _display_name(headers)
    sender = sender or ""
    summary = {
        "from": _one_line((name + " <" + sender + ">") if name else sender or _single(headers, "From")),
        "reply_to": _one_line(", ".join(_addresses(headers, "Reply-To"))),
        "to": _one_line(", ".join(_addresses(headers, "To"))),
        "cc": _one_line(", ".join(_addresses(headers, "Cc"))),
        "subject": _one_line(subject),
        "date": _one_line(_single(headers, "Date"), 80),
        "list_unsubscribe_present": bool(headers.get_all("List-Unsubscribe", [])),
        "list_id": _one_line(_single(headers, "List-ID"), 200),
        "precedence": _one_line(_single(headers, "Precedence"), 40),
        "in_reply_to_present": bool(headers.get_all("In-Reply-To", []) or headers.get_all("References", [])),
        "authentication_verified_by_icloud": dict(NEUTRAL_AUTH),
        "attachments": _attachments(headers),
    }

    def serialize(text, truncated):
        value = json.dumps({"headers": summary, "body": text, "body_truncated": truncated},
                           ensure_ascii=False, separators=(",", ":"), sort_keys=False)
        return value.replace("<", "\\u003c").replace(">", "\\u003e")
    data = serialize(body, not bounded)
    if len(data) > DATA_LIMIT:
        low, high = 0, len(body)
        while low < high:
            middle = (low + high + 1) // 2
            if len(serialize(body[:middle], True)) <= DATA_LIMIT:
                low = middle
            else:
                high = middle - 1
        data = serialize(body[:low], True)
    if len(data) > DATA_LIMIT or "<<<" in data or ">>>" in data:
        raise ValueError("the data block does not fit")
    return data, sender


def _precedence(value):
    value = (value or "").strip().lower()
    return value if value in ("bulk", "list", "junk") else ("other" if value else "none")


def _sender_kind(sender_domain, h):
    if h.get("list_unsubscribe_present") or h.get("list_id") or (h.get("precedence") or "").lower() in {"bulk", "list", "junk"}:
        return "bulk_list"
    if sender_domain in FREEMAIL:
        return "freemail"
    return "organisation" if sender_domain else "unknown"


def view_from_data(data: str, sender_domain: str) -> str:
    value = json.loads(data)
    h = value["headers"]
    subject, body = h.get("subject") or "", value.get("body") or ""
    sender, reply_to = h.get("from") or "", h.get("reply_to") or ""
    recipients = len([x for x in (h.get("to") or "").split(",") if x.strip()]) + \
        len([x for x in (h.get("cc") or "").split(",") if x.strip()])
    facts = (f"sender: {_sender_kind(sender_domain or '', h)}; list-unsubscribe: {'yes' if h.get('list_unsubscribe_present') else 'no'}; "
             f"list-id: {'yes' if h.get('list_id') else 'no'}; precedence: {_precedence(h.get('precedence'))}; "
             f"reply-chain: {'yes' if h.get('in_reply_to_present') else 'no'}; attachments: {int(h.get('attachments') or 0)}; "
             f"recipients: {recipients}; truncated: {'yes' if value.get('body_truncated') else 'no'}")
    return facts + "\nfrom: " + sender + "\nreply-to: " + reply_to + "\nsubject: " + subject + "\n\n" + body


def view(raw: bytes) -> str:
    """Raw mail bytes -> the Bengu-1 input text."""
    data, sender = data_block(raw)
    return view_from_data(data, sender.split("@")[1] if "@" in sender else "")