argocd-apps/apps/base/mail/resources/sieve-reconcile.py

import base64
import json
import os
import re
import sys
import urllib.parse
import urllib.request

USING = ["urn:ietf:params:jmap:core", "urn:ietf:params:jmap:sieve", "urn:stalwart:jmap"]
HEADER = 'require ["include"];\ninclude :global "spam-to-junk";\n'
HAS_INCLUDE = re.compile(r'\binclude\b[^;]*"spam-to-junk"')
TOKENS = re.compile(r'"(?:\\.|[^"\\])*"|#[^\n]*|/\*.*?\*/', re.S)


def auth():
    user = f"{os.environ['STALWART_USER']}:{os.environ['STALWART_PASSWORD']}"
    return "Basic " + base64.b64encode(user.encode()).decode()


def http(path, data=None, ctype="application/json"):
    url = os.environ["STALWART_URL"].rstrip("/") + path
    req = urllib.request.Request(url, data=data, headers={"Authorization": auth()})
    if data is not None:
        req.add_header("Content-Type", ctype)
    with urllib.request.urlopen(req, timeout=30) as resp:
        return resp.read()


def jmap(*calls):
    body = json.dumps({"using": USING, "methodCalls": [[m, a, str(i)] for i, (m, a) in enumerate(calls)]})
    responses = json.loads(http("/jmap", body.encode()))["methodResponses"]
    for name, args, _ in responses:
        if name == "error":
            raise RuntimeError(args)
    return [args for _, args, _ in responses]


def blank(script):
    # Same-length copy with comments as spaces and string contents as "x", so offsets match.
    def mask(m):
        t = m.group(0)
        if t.startswith('"'):
            return '"' + re.sub(r"[^\n]", "x", t[1:-1]) + '"'
        return re.sub(r"[^\n]", " ", t)
    return TOKENS.sub(mask, script)


def commands(script):
    s, start, depth = blank(script), None, 0
    for p, c in enumerate(s):
        if start is None and not c.isspace():
            start = p
        depth += {"{": 1, "}": -1}.get(c, 0)
        if start is not None and depth == 0 and c in ";}":
            yield re.match(r"\w*", s[start:]).group(0).lower(), start, p + 1
            start = None


def with_include(script):
    at = 0
    for name, _, end in commands(script):
        if name != "require":
            break
        at = end
    if at:
        eol = script.find("\n", at)
        eol = len(script) if eol < 0 else eol + 1
        if not blank(script)[at:eol].strip():
            at = eol
        if not script[:at].endswith("\n"):
            return script[:at] + "\n" + HEADER + script[at:]
    return script[:at] + HEADER + script[at:]


def has_include(script):
    uncommented = TOKENS.sub(lambda m: m.group(0) if m.group(0).startswith('"') else " ", script)
    return bool(HAS_INCLUDE.search(uncommented))


def users():
    ids, position = [], 0
    while True:
        (page,) = jmap(("x:Account/query", {"position": position, "limit": 500, "calculateTotal": True}))
        ids += page["ids"]
        position += len(page["ids"])
        if not page["ids"] or position >= page.get("total", position):
            break
    for start in range(0, len(ids), 500):
        (got,) = jmap(("x:Account/get", {"ids": ids[start:start + 500], "properties": ["@type", "name"]}))
        yield from (a for a in got["list"] if a["@type"] == "User")


def reconcile(account):
    acct = account["id"]
    (got,) = jmap(("SieveScript/get", {"accountId": acct}))
    scripts = got["list"]
    target = next((s for s in scripts if s["isActive"]), None) or next(
        (s for s in scripts if s["name"] == "default"), None
    )
    content = ""
    if target:
        content = http(f"/jmap/download/{acct}/{target['blobId']}/script").decode()
        if target["isActive"] and has_include(content):
            return
    if not has_include(content):
        content = with_include(content)
    blob = json.loads(http(f"/jmap/upload/{acct}/", content.encode(), "application/sieve"))["blobId"]
    if target:
        args = {"accountId": acct, "update": {target["id"]: {"blobId": blob}},
                "onSuccessActivateScript": target["id"]}
        action = f"added spam-to-junk include to {target['name']!r}"
    else:
        args = {"accountId": acct, "create": {"d": {"name": "default", "blobId": blob}},
                "onSuccessActivateScript": "#d"}
        action = "created and activated 'default'"
    (res,) = jmap(("SieveScript/set", args))
    failed = res.get("notCreated") or res.get("notUpdated")
    if failed:
        raise RuntimeError(failed)
    print(f"{account['name']}: {action}", flush=True)


SELFTEST_CASES = {
    "empty": "",
    "no require": "keep;\n",
    "no trailing newline": 'require "fileinto";',
    "plain": 'require ["fileinto"];\nif true { fileinto "A"; }\n',
    "multiline after comment": '# mine\nrequire ["fileinto",\n  "mailbox"];\nkeep;\n',
    "trailing comment": 'require "fileinto"; # needed; really\nfileinto "A";',
    "bracketed comment": '/* require "x"; */ require /* ; */ ["fileinto"];\nrequire "mailbox";\nstop;\n',
}


def selftest():
    for label, script in SELFTEST_CASES.items():
        out = with_include(script)
        names = [n for n, _, _ in commands(out)]
        first = next((i for i, n in enumerate(names) if n != "require"), len(names))
        assert "require" not in names[first:], (label, out)
        assert names[first:first + 1] == ["include"], (label, out)
        assert has_include(out) and not has_include(script), label
        assert out.replace(HEADER, "", 1).rstrip("\n") == script.rstrip("\n"), label
    assert not has_include('# include :global "spam-to-junk";\n')
    print("selftest ok")


if __name__ == "__main__":
    if sys.argv[1:] == ["--selftest"]:
        selftest()
        sys.exit(0)
    errors = 0
    for account in users():
        try:
            reconcile(account)
        except Exception as err:
            errors += 1
            print(f"{account['name']}: {err}", file=sys.stderr, flush=True)
    sys.exit(1 if errors else 0)