#!/usr/bin/env python3
"""Post-install patches for offlineimap on Alpine.

Run as root in the Docker build (RUN python3 /tmp/patch.py).
Fixes two upstream issues that persist in offlineimap 8.0.0:

1. UnicodeEncodeError when writing messages with non-ASCII bytes
   (Maildir.save thread + IMAP header generation).
2. SSL SECLEVEL=2 in OpenSSL 3.x prevents connecting to servers with
   weak DH keys (DH_KEY_TOO_SMALL). Force SECLEVEL=1 in imaplib2.

Paths are resolved dynamically so the patch survives Python version
bumps across Alpine releases (was hardcoded to python3.12).
"""

import glob
import sys

# ---------------------------------------------------------------
# 1. Maildir.py — wrap fd.write(msg.as_bytes(...)) in a try/except
#    that falls back to UTF-8 encoded as_string on UnicodeEncodeError
# ---------------------------------------------------------------

maildir_path = glob.glob(
    "/usr/lib/python%d.%d/site-packages/offlineimap/folder/Maildir.py"
    % (sys.version_info.major, sys.version_info.minor)
)[0]

with open(maildir_path, "r") as f:
    maildir = f.read()

maildir = maildir.replace(
    "fd.write(msg.as_bytes(policy=output_policy))",
    "try:\n            fd.write(msg.as_bytes(policy=output_policy))\n        except UnicodeEncodeError:\n            fd.write(msg.as_string(policy=output_policy).encode('utf-8'))",
)

with open(maildir_path, "w") as f:
    f.write(maildir)

# ---------------------------------------------------------------
# 2. IMAP.py — add a helper and replace msg.as_bytes(policy=...) calls
#    with the helper that falls back on UnicodeEncodeError
# ---------------------------------------------------------------

imap_path = glob.glob(
    "/usr/lib/python%d.%d/site-packages/offlineimap/folder/IMAP.py"
    % (sys.version_info.major, sys.version_info.minor)
)[0]

with open(imap_path, "r") as f:
    imap_py = f.read()

imap_patch = """
def get_msg_bytes(msg, policy):
    try:
        return msg.as_bytes(policy=policy)
    except UnicodeEncodeError:
        return msg.as_string(policy=policy).encode('utf-8')
"""
imap_py = imap_patch + imap_py.replace(
    "msg.as_bytes(policy=output_policy)", "get_msg_bytes(msg, output_policy)"
)

with open(imap_path, "w") as f:
    f.write(imap_py)

# ---------------------------------------------------------------
# 3. imaplib2.py — force SECLEVEL=1 to allow weak DH keys
#    (moved here from a sed one-liner in the Dockerfile)
# ---------------------------------------------------------------

imaplib2_paths = glob.glob(
    "/usr/lib/python%d.%d/site-packages/imaplib2/imaplib2.py"
    % (sys.version_info.major, sys.version_info.minor)
)

if imaplib2_paths:
    imaplib2_path = imaplib2_paths[0]
    with open(imaplib2_path, "r") as f:
        imaplib2_py = f.read()

    old = "ctx = ssl.SSLContext(ssl_version)"
    new = (
        'ctx = ssl.SSLContext(ssl_version)\n'
        '                ctx.set_ciphers("DEFAULT:@SECLEVEL=1")'
    )

    if old in imaplib2_py:
        imaplib2_py = imaplib2_py.replace(old, new)
        with open(imaplib2_path, "w") as f:
            f.write(imaplib2_py)
        print("Patched imaplib2 SECLEVEL=1")
    else:
        print("imaplib2 already patched or pattern not found — skipping")
else:
    print("imaplib2 not found — skipping SECLEVEL patch")