Insert spam-to-junk include after the require block and self-test it

.pre-commit-config.yaml

@@ -42,3 +42,9 @@ repos:
         entry: ci/validate-no-secrets.sh
         language: system
         pass_filenames: false
+      - id: sieve_reconcile_selftest
+        name: sieve-reconcile selftest
+        entry: python3 apps/base/mail/resources/sieve-reconcile.py --selftest
+        language: system
+        files: sieve-reconcile\.py$
+        pass_filenames: false

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

@@ -6,18 +6,20 @@ import sys
 import urllib.parse
 import urllib.request
 
-URL = os.environ["STALWART_URL"].rstrip("/")
-AUTH = "Basic " + base64.b64encode(
-    f"{os.environ['STALWART_USER']}:{os.environ['STALWART_PASSWORD']}".encode()
-).decode()
 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"')
-REQUIRE = re.compile(r"\s*require\b")
+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"):
-    req = urllib.request.Request(URL + path, data=data, headers={"Authorization": AUTH})
+    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:
@@ -33,23 +35,52 @@ def jmap(*calls):
     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):
-    lines = script.splitlines(keepends=True)
-    i, in_require = 0, False
-    while i < len(lines):
-        line = lines[i].strip()
-        if in_require or REQUIRE.match(lines[i]):
-            in_require = not line.endswith(";")
-        elif line and not line.startswith("#"):
+    at = 0
+    for name, _, end in commands(script):
+        if name != "require":
             break
-        i += 1
-    return HEADER + "".join(lines[:i]) + "".join(lines[i:])
+        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}))
+        (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):
@@ -69,9 +100,9 @@ def reconcile(account):
     content = ""
     if target:
         content = http(f"/jmap/download/{acct}/{target['blobId']}/script").decode()
-        if target["isActive"] and HAS_INCLUDE.search(content):
+        if target["isActive"] and has_include(content):
             return
-    if not HAS_INCLUDE.search(content):
+    if not has_include(content):
         content = with_include(content)
     blob = json.loads(http(f"/jmap/upload/{acct}/", content.encode(), "application/sieve"))["blobId"]
     if target:
@@ -89,11 +120,39 @@ def reconcile(account):
     print(f"{account['name']}: {action}", flush=True)
 
 
-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)
+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)