coruna-lab-web/frontend/tools/_scheme_patch.py

"""Keep Deployment/Reporting URL scheme as ``https://%@``.

Core/type0x01 build URLs with the cstring/CFString format ``https://%@``.
``daily.html`` plugin URLs use ``https://[HOST_PLACEHOLDER]/...``.

Older lab builds patched scheme to cleartext for proxy testing:

- ``https://%@`` → ``http://%@\\0`` (9 visible chars + NUL pad), CFString length 10 → 9
- mistaken same-length form ``http://%@/`` (caused ``host//path``)

This module leaves pristine ``https://%@`` alone and restores any of those
legacy http forms back to ``https://%@`` (CFString length 10).
"""

from __future__ import annotations

import struct

from _path_patch import _arm64e_target, _fat_slices, _parse_macho

HTTPS = b"https://%@"
# Previous mistaken same-length patch (caused host//path).
LEGACY_SLASH = b"http://%@/"
# 10-byte slot from http cleartext patch: 9-char format + NUL pad.
HTTP_SLOT = b"http://%@\x00"
HTTP_VISIBLE = b"http://%@"

# Back-compat aliases for tests / importers that still use the old names.
OLD = HTTPS
NEW_SLOT = HTTP_SLOT
NEW_VISIBLE = HTTP_VISIBLE


def _patch_thin_scheme(data: bytes, *, expect_hits: int, label: str) -> bytes:
    macho = _parse_macho(data, label=label)
    base_subtype = macho.cpu_subtype & 0x00FFFFFF
    architecture = (
        "arm64e" if base_subtype == 2 else "arm64" if base_subtype == 0 else ""
    )
    if architecture not in {"arm64", "arm64e"}:
        raise SystemExit(f"{label}: unsupported CPU subtype {base_subtype}")

    cstring = macho.section("__TEXT", "__cstring")
    region = bytearray(data[cstring.offset : cstring.offset + cstring.size])

    https_count = region.count(HTTPS)
    slash_count = region.count(LEGACY_SLASH)
    http_slot_count = region.count(HTTP_SLOT)

    if https_count == expect_hits and slash_count == 0 and http_slot_count == 0:
        # Already https://%@ — nothing to do.
        return data

    if slash_count == expect_hits and https_count == 0 and http_slot_count == 0:
        source = LEGACY_SLASH
    elif http_slot_count == expect_hits and https_count == 0 and slash_count == 0:
        source = HTTP_SLOT
    else:
        raise SystemExit(
            f"{label}: expected {expect_hits} {HTTPS!r} (or legacy http forms) in "
            f"__cstring; found https={https_count} slash={slash_count} "
            f"http_slot={http_slot_count}"
        )

    hits: list[int] = []
    start = 0
    while True:
        hit = region.find(source, start)
        if hit < 0:
            break
        hits.append(hit)
        start = hit + 1
    if len(hits) != expect_hits:
        raise SystemExit(
            f"{label}: __cstring scheme hits={len(hits)}, expected {expect_hits}"
        )

    for hit in hits:
        region[hit : hit + len(HTTPS)] = HTTPS

    buf = bytearray(data)
    buf[cstring.offset : cstring.offset + cstring.size] = region

    cfstring = macho.section("__DATA_CONST", "__cfstring")
    if cfstring.size % 32:
        raise SystemExit(f"{label}: __cfstring size is not record-aligned")

    patched_lengths = 0
    for hit in hits:
        path_vmaddr = cstring.addr + hit
        for record_offset in range(
            cfstring.offset, cfstring.offset + cfstring.size, 32
        ):
            raw = struct.unpack_from("<Q", buf, record_offset + 16)[0]
            length = struct.unpack_from("<Q", buf, record_offset + 24)[0]
            target = raw if architecture == "arm64" else _arm64e_target(raw)
            if target != path_vmaddr:
                continue
            if length not in (len(HTTPS), len(HTTP_VISIBLE)):
                raise SystemExit(
                    f"{label}: scheme CFString length={length}, expected "
                    f"{len(HTTPS)} or {len(HTTP_VISIBLE)}"
                )
            struct.pack_into("<Q", buf, record_offset + 24, len(HTTPS))
            patched_lengths += 1
            break
        else:
            raise SystemExit(
                f"{label}: no CFString pointing at scheme cstring vmaddr {path_vmaddr:#x}"
            )

    if patched_lengths != expect_hits:
        raise SystemExit(
            f"{label}: patched {patched_lengths} CFString lengths, expected {expect_hits}"
        )
    return bytes(buf)


def ensure_deployment_scheme_https(
    data: bytes,
    *,
    expect_hits: int,
    label: str,
) -> bytes:
    """Ensure ``https://%@`` (CFString length 10) in thin or fat dylibs.

    ``expect_hits`` is the total number of format strings across the whole file
    (2 for fat core, 1 for thin type0x01). Restores legacy lab http patches.
    """
    slices = _fat_slices(data, label=label)
    if slices is None:
        return _patch_thin_scheme(data, expect_hits=expect_hits, label=label)

    if expect_hits % len(slices) != 0:
        raise SystemExit(
            f"{label}: expect_hits={expect_hits} not divisible by fat slices={len(slices)}"
        )
    per_slice = expect_hits // len(slices)
    buf = bytearray(data)
    for index, slice_info in enumerate(slices):
        thin = bytes(buf[slice_info.offset : slice_info.offset + slice_info.size])
        patched = _patch_thin_scheme(
            thin, expect_hits=per_slice, label=f"{label}[slice{index}]"
        )
        if len(patched) != slice_info.size:
            raise SystemExit(f"{label}: fat slice size changed after scheme patch")
        buf[slice_info.offset : slice_info.offset + slice_info.size] = patched
    return bytes(buf)