unlockedRestrictions/tools/verify_bundle.py

#!/usr/bin/env python3
"""Offline verification harness for a built .plugin bundle. Runs without a phone.

Usage: python3 tools/verify_bundle.py <bundle.plugin>
"""
from __future__ import annotations

import ast
import json
import sys
import traceback
import types
from dataclasses import dataclass
from enum import Enum, auto
from pathlib import Path

REPO_ROOT = Path(__file__).resolve().parent.parent
PLUGIN_JSON = REPO_ROOT / "plugin.json"

_FAILURES = []


def check(name: str, condition: bool, detail: str = "") -> None:
    if condition:
        print(f"PASS {name}")
    else:
        message = f"FAIL {name}"
        if detail:
            message += f": {detail}"
        print(message, file=sys.stderr)
        _FAILURES.append(name)


class HookStrategy(Enum):
    DEFAULT = auto()
    MODIFY = auto()
    CANCEL = auto()
    MODIFY_FINAL = auto()


@dataclass
class HookResult:
    strategy: HookStrategy = HookStrategy.DEFAULT
    params: object = None


LOG_LINES: list[str] = []
HOOKED: list[tuple[str, str, object]] = []

# Classes the plugin is expected to hook, mapped to the hooked method names.
EXPECTED_HOOKS = {
    ("org.telegram.messenger.FlagSecureReason", "isSecuredNow"),
    ("android.view.Window", "setFlags"),
    ("org.telegram.messenger.MessagesController", "isChatNoForwards"),
    ("org.telegram.messenger.MessagesController", "isUserNoForwards"),
    ("org.telegram.messenger.MessagesController", "isPeerNoForwards"),
    ("org.telegram.ui.ChatActivity", "hasSelectedNoforwardsMessage"),
    ("org.telegram.messenger.MessageObject", "canForwardMessage"),
    ("org.telegram.messenger.SendMessagesHelper", "sendMessage"),
    ("org.telegram.ui.Cells.DialogCell", "buildLayout"),
    ("org.telegram.ui.Components.ChatAvatarContainer", "setTitle"),
    ("org.telegram.messenger.AndroidUtilities", "escape"),
    ("org.telegram.ui.ProfileActivity", "updateProfileData"),
}


class MethodHook:
    def before_hooked_method(self, param):
        pass

    def after_hooked_method(self, param):
        pass


class MethodReplacement:
    pass


class FakeJavaClass:
    def __init__(self, name):
        self.name = name

    def getName(self):
        return self.name


class BasePlugin:
    def __init__(self):
        self._settings = {}

    def get_setting(self, key, default=None):
        return self._settings.get(key, default)

    def set_setting(self, key, value, reload_settings=False):
        self._settings[key] = value

    def add_on_send_message_hook(self, priority=0):
        pass

    def hook_all_methods(self, hook_class, method_name, xposed_hook=None, priority=None, **kwargs):
        HOOKED.append((hook_class.getName(), method_name, xposed_hook))
        return [object()]

    def log(self, message):
        LOG_LINES.append(str(message))

    def on_plugin_load(self):
        pass

    def on_plugin_unload(self):
        pass

    def create_settings(self):
        return []


def fake_log(*args, **kwargs):
    LOG_LINES.append(" ".join(str(a) for a in args))


def fake_find_class(name):
    return FakeJavaClass(name)


class _SettingsRow:
    def __init__(self, **kwargs):
        self.__dict__.update(kwargs)


class FakeArrayList(list):
    def isEmpty(self):
        return not self

    def add(self, item):
        self.append(item)

    def size(self):
        return len(self)

    def get(self, i):
        return self[i]


class FakeMessageObject:
    def __init__(self, protected=False, noforwards=False):
        self.protected = protected
        self.messageOwner = types.SimpleNamespace(noforwards=noforwards)


class FakeParam:
    def __init__(self, args=(), this=None, result=None):
        self.args = None if args is None else list(args)
        self.thisObject = this
        self._result = result
        self.result_set = False

    def getResult(self):
        return self._result

    def setResult(self, value):
        self._result = value
        self.result_set = True


def _module(name, **attrs):
    mod = types.ModuleType(name)
    mod.__dict__.update(attrs)
    sys.modules[name] = mod
    return mod


def install_fakes() -> None:
    _module(
        "base_plugin",
        BasePlugin=BasePlugin,
        HookResult=HookResult,
        HookStrategy=HookStrategy,
        MethodHook=MethodHook,
        MethodReplacement=MethodReplacement,
    )
    _module("android_utils", log=fake_log, run_on_ui_thread=lambda fn, delay=0: fn())
    _module(
        "hook_utils",
        find_class=fake_find_class,
        set_private_field=lambda obj, name, value: setattr(obj, name, value) or True,
        get_private_field=lambda obj, name: getattr(obj, name, None),
    )
    _module("client_utils", get_last_fragment=lambda: None)
    ui_mod = _module("ui", __path__=[])
    ui_mod.settings = _module(
        "ui.settings",
        **{n: type(n, (_SettingsRow,), {}) for n in ("Header", "Input", "Text", "Switch", "Divider")},
    )
    ui_mod.bulletin = _module("ui.bulletin", BulletinHelper=types.SimpleNamespace(
        show_info=fake_log, show_error=fake_log, show_success=fake_log,
    ))
    java_mod = _module("java", __path__=[])
    java_mod.util = _module("java.util", ArrayList=FakeArrayList)
    org = _module("org", __path__=[])
    org.telegram = _module("org.telegram", __path__=[])
    org.telegram.messenger = _module("org.telegram.messenger", MessageObject=FakeMessageObject)


KEY_TO_DUNDER = {
    "id": "__id__",
    "name": "__name__",
    "description": "__description__",
    "author": "__author__",
    "version": "__version__",
    "icon": "__icon__",
    "app_version": "__app_version__",
    "sdk_version": "__sdk_version__",
    "requirements": "__requirements__",
}


def expected_dunders(meta: dict) -> dict:
    result = {}
    for key, dunder in KEY_TO_DUNDER.items():
        if key not in meta:
            continue
        value = meta[key]
        if key in ("description", "icon") and value == "":
            continue
        if key == "requirements" and value == []:
            continue
        result[dunder] = value
    return result


def verify_hooks(plugin, hooks) -> None:
    flag_secure = 0x2000
    strip = hooks[("android.view.Window", "setFlags")]
    for enabled in (True, False):
        plugin.set_setting("allow_screenshots", enabled)
        param = FakeParam(args=[flag_secure | 0x80, flag_secure | 0x80])
        strip.before_hooked_method(param)
        want = 0x80 if enabled else flag_secure | 0x80
        check(
            f"Window.setFlags strips FLAG_SECURE only when enabled (enabled={enabled})",
            param.args[0] == want,
            f"flags={param.args[0]:#x}",
        )

    plugin.set_setting("allow_screenshots", True)
    param = FakeParam(result=True)
    hooks[("org.telegram.messenger.FlagSecureReason", "isSecuredNow")].after_hooked_method(param)
    check("FlagSecureReason.isSecuredNow forced false", param.getResult() is False)

    no_forwards = hooks[("org.telegram.messenger.MessagesController", "isChatNoForwards")]
    for enabled in (True, False):
        plugin.set_setting("allow_forwarding", enabled)
        param = FakeParam(result=True)
        no_forwards.after_hooked_method(param)
        check(
            f"isChatNoForwards forced false only when enabled (enabled={enabled})",
            param.getResult() is (not enabled),
            f"got {param.getResult()!r}",
        )

    plugin.set_setting("allow_forwarding", True)
    can_forward = hooks[("org.telegram.messenger.MessageObject", "canForwardMessage")]
    message = FakeMessageObject(noforwards=True)
    seen = []
    # The re-invoked original answers "forwardable" only if noforwards is unset.
    method = types.SimpleNamespace(
        invoke=lambda obj: seen.append(obj.messageOwner.noforwards) or not obj.messageOwner.noforwards
    )
    param = FakeParam(this=message, result=False)
    param.method = method
    can_forward.after_hooked_method(param)
    check(
        "canForwardMessage re-evaluates without noforwards and restores the flag",
        param.getResult() is True and seen == [False] and message.messageOwner.noforwards is True,
        f"result={param.getResult()!r} seen={seen} flag={message.messageOwner.noforwards!r}",
    )
    param = FakeParam(this=FakeMessageObject(noforwards=False), result=False)
    param.method = method
    can_forward.after_hooked_method(param)
    check("canForwardMessage keeps a refusal unrelated to noforwards", param.getResult() is False)

    escape = hooks[("org.telegram.messenger.AndroidUtilities", "escape")]
    badge = type(escape).after_hooked_method.__globals__
    badge["_building_title"] = "Chat"
    results = []
    for arg in ("Chat", "Other"):
        param = FakeParam(args=[arg], result=arg)
        escape.after_hooked_method(param)
        results.append(param.getResult())
    badge["_building_title"] = None
    param = FakeParam(args=["Chat"], result="Chat")
    escape.after_hooked_method(param)
    results.append(param.getResult())
    check(
        "escape() prefixes only the title being laid out",
        results == ["\u2003 Chat", "Other", "Chat"],
        f"{results}",
    )

    intercept = hooks[("org.telegram.messenger.SendMessagesHelper", "sendMessage")]
    copier = type(intercept).before_hooked_method.__globals__["copier"]

    cases = [
        ([], []),
        ([0, 0, 0], []),
        ([7], []),
        ([7, 7, 7], [(0, 3)]),
        ([0, 7, 7, 0, 9, 9, 9], [(1, 3), (4, 7)]),
        ([7, 7, 9, 9], [(0, 2), (2, 4)]),
        ([7, 0, 7], []),
        ([5] * 12, [(0, 10), (10, 12)]),
        ([5] * 11, [(0, 10)]),
    ]
    bad = [(k, copier.album_runs(k), want) for k, want in cases if copier.album_runs(k) != want]
    check("album_runs groups consecutive album items, max 10 per group", not bad, f"{bad}")

    # Stub out the Java-facing parts of the bundled copier module.
    copied = []
    copier.is_protected = lambda m: m.protected
    copier.copy_messages = lambda **kwargs: copied.append(kwargs)

    free, locked = FakeMessageObject(), FakeMessageObject(protected=True)
    param = FakeParam(args=[FakeArrayList([free, locked]), -100, False, True, True, 0, 0])
    intercept.before_hooked_method(param)
    check(
        "mixed forward: protected message copied, the rest forwarded",
        len(copied) == 1
        and copied[0]["messages"] == [locked]
        and copied[0]["peer"] == -100
        and copied[0]["hide_caption"] is True
        and list(param.args[0]) == [free]
        and not param.result_set,
        f"copied={copied} args0={list(param.args[0])} result_set={param.result_set}",
    )

    copied.clear()
    param = FakeParam(args=[FakeArrayList([locked]), 5, False, False, True, 0, 0])
    intercept.before_hooked_method(param)
    check(
        "all-protected forward: nothing left for the original sendMessage to forward",
        len(copied) == 1 and list(param.args[0]) == [],
        f"copied={copied} args0={list(param.args[0])}",
    )

    copied.clear()

    def broken_check(m):
        raise RuntimeError("boom")

    copier.is_protected = broken_check
    param = FakeParam(args=[FakeArrayList([free]), 5, False, False, True, 0, 0])
    intercept.before_hooked_method(param)
    check(
        "failed protection check copies instead of forwarding",
        len(copied) == 1 and list(param.args[0]) == [],
        f"copied={copied} args0={list(param.args[0])}",
    )
    copier.is_protected = lambda m: m.protected

    copied.clear()
    param = FakeParam(args=[FakeArrayList([free]), 5, False, False, True, 0, 0])
    intercept.before_hooked_method(param)
    check(
        "unprotected forward passes through untouched",
        not copied and not param.result_set and list(param.args[0]) == [free],
    )

    plugin.set_setting("allow_forwarding", False)
    param = FakeParam(args=[FakeArrayList([locked]), 5, False, False, True, 0, 0])
    intercept.before_hooked_method(param)
    check("forward interception off when setting disabled", not copied and not param.result_set)

    plugin.set_setting("allow_forwarding", True)
    param = FakeParam(args=[object()])
    intercept.before_hooked_method(param)
    check("sendMessage(SendMessageParams) overload ignored", not copied and not param.result_set)


def main(argv=None) -> int:
    argv = sys.argv[1:] if argv is None else argv
    if len(argv) != 1:
        print("usage: tools/verify_bundle.py <bundle.plugin>", file=sys.stderr)
        return 1
    bundle_path = Path(argv[0])
    if not bundle_path.is_file():
        print(f"error: bundle not found at {bundle_path}", file=sys.stderr)
        return 1

    install_fakes()
    source = bundle_path.read_text(encoding="utf-8")

    meta = json.loads(PLUGIN_JSON.read_text(encoding="utf-8"))
    modules_before = set(sys.modules)
    ns = {"__name__": meta["id"]}
    exec(compile(source, str(bundle_path), "exec"), ns)
    new_modules = set(sys.modules) - modules_before

    expected = expected_dunders(meta)
    mismatches = [
        f"{dunder}: expected {value!r}, got {ns.get(dunder)!r}"
        for dunder, value in expected.items()
        if ns.get(dunder) != value
    ]
    check("metadata constants match plugin.json", not mismatches, "; ".join(mismatches))

    tree = ast.parse(source, filename=str(bundle_path))
    top_level_assigned = {
        target.id
        for node in tree.body
        if isinstance(node, ast.Assign)
        for target in node.targets
        if isinstance(target, ast.Name)
    }
    check(
        "metadata exposed as top-level literal assignments (AST)",
        set(expected) <= top_level_assigned,
        f"missing: {set(expected) - top_level_assigned}",
    )

    plugin_classes = [
        value
        for value in ns.values()
        if isinstance(value, type) and issubclass(value, BasePlugin) and value is not BasePlugin
    ]
    check(
        "exactly one BasePlugin subclass defined",
        len(plugin_classes) == 1,
        f"found {len(plugin_classes)}",
    )

    plugin_cls = plugin_classes[0] if len(plugin_classes) == 1 else None
    if plugin_cls is not None:
        check(
            "plugin class __module__ matches executed namespace",
            plugin_cls.__module__ == ns["__name__"],
            f"__module__={plugin_cls.__module__!r} != {ns['__name__']!r}",
        )

        plugin = plugin_cls()
        LOG_LINES.clear()
        plugin.on_plugin_load()
        check(
            "on_plugin_load emits load log line",
            any("unlocked_restrictions loaded" in line for line in LOG_LINES),
            f"log lines: {LOG_LINES}",
        )
        hooked = {(cls, method) for cls, method, _ in HOOKED}
        check(
            "on_plugin_load hooks every target method",
            EXPECTED_HOOKS <= hooked,
            f"missing: {EXPECTED_HOOKS - hooked}",
        )

        switches = [row for row in plugin.create_settings() if type(row).__name__ == "Switch"]
        check(
            "settings expose all toggles, enabled by default",
            {(row.key, row.default) for row in switches}
            == {
                ("allow_screenshots", True),
                ("allow_forwarding", True),
                ("protected_lock_badge", True),
            },
            f"got {[(row.key, row.default) for row in switches]}",
        )

        verify_hooks(plugin, {(cls, method): hook for cls, method, hook in HOOKED})

    # Raise inside a bundled module to prove linecache registration: the
    # traceback must name the original source file and show the failing line.
    screenshots = ns.get("screenshots")
    traceback_text = ""
    if screenshots is not None:
        try:
            screenshots.StripFlagSecureHook(lambda: True).before_hooked_method(FakeParam(args=None))
        except TypeError:
            traceback_text = traceback.format_exc()
    check(
        "hook traceback names src/screenshots.py",
        "src/screenshots.py" in traceback_text,
        f"traceback:\n{traceback_text}",
    )
    check(
        "traceback includes the raising source line text",
        "if len(args) != 2" in traceback_text,
        f"traceback:\n{traceback_text}",
    )

    bundled = {"main", "hooking", "screenshots", "forwarding", "copier", "badge", "protection"}
    check(
        "sys.modules gained no entry for bundled module names",
        not (bundled & new_modules),
        f"leaked: {bundled & new_modules}",
    )

    if _FAILURES:
        print(f"\n{len(_FAILURES)} check(s) failed: {', '.join(_FAILURES)}", file=sys.stderr)
        return 1
    print("\nall checks passed")
    return 0


if __name__ == "__main__":
    sys.exit(main())