#!/usr/bin/env python3

import argparse
import os
import platform
import shlex
import subprocess
import time

from run_redroid import DEVICE_PROFILES


def run(cmd, check=True, capture_output=False):
    return subprocess.run(
        cmd,
        check=check,
        text=True,
        capture_output=capture_output,
    )


def adb(serial, args, check=True, capture_output=False):
    return run(
        ["adb", "-s", serial] + args,
        check=check,
        capture_output=capture_output,
    )


class Backend:
    def __init__(self, serial="", container_name="", container_runtime="docker"):
        self.serial = serial
        self.container_name = container_name
        self.container_runtime = container_runtime
        self._root_checked = False
        self._root_available = False

    @property
    def mode(self):
        return "container" if self.container_name else "adb"

    def connect(self):
        if self.mode == "adb" and ":" in self.serial:
            run(["adb", "connect", self.serial], check=False)

    def has_root(self):
        if self.mode == "container":
            return True
        if not self._root_checked:
            result = adb(
                self.serial,
                ["shell", "su", "-mm", "-c", "id"],
                check=False,
                capture_output=True,
            )
            self._root_available = result.returncode == 0 and "uid=0(" in result.stdout
            self._root_checked = True
        return self._root_available

    def exec_shell(self, command, check=True, capture_output=False, root=False):
        if self.mode == "container":
            return run(
                [
                    self.container_runtime,
                    "exec",
                    self.container_name,
                    "/system/bin/sh",
                    "-c",
                    command,
                ],
                check=check,
                capture_output=capture_output,
            )
        if root and self.has_root():
            return adb(
                self.serial,
                ["shell", "su", "-mm", "-c", command],
                check=check,
                capture_output=capture_output,
            )
        return adb(
            self.serial,
            ["shell", command],
            check=check,
            capture_output=capture_output,
        )

    def wait_for_boot(self, timeout):
        deadline = time.time() + timeout
        if self.mode == "adb":
            run(["adb", "-s", self.serial, "wait-for-device"])
        while time.time() < deadline:
            boot = self.exec_shell(
                "getprop sys.boot_completed",
                check=False,
                capture_output=True,
            ).stdout.strip()
            if boot == "1":
                return
            time.sleep(2)
        raise TimeoutError("Timed out waiting for Android boot completion")


def try_shell(backend, command, root=False):
    backend.exec_shell(command, check=False, root=root)


def put_setting(backend, namespace, key, value):
    try_shell(backend, "settings put {} {} {}".format(namespace, key, value))


def delete_setting(backend, namespace, key):
    try_shell(backend, "settings delete {} {}".format(namespace, key))


def detect_gateway_ip():
    system = platform.system().lower()
    if system == "linux":
        result = run(
            ["sh", "-lc", "ip route show default 2>/dev/null | awk 'NR==1 {print $3}'"],
            capture_output=True,
        )
        gateway = result.stdout.strip()
        if gateway:
            return gateway
    elif system == "windows":
        result = run(
            [
                "powershell",
                "-NoProfile",
                "-Command",
                "(Get-NetRoute -DestinationPrefix '0.0.0.0/0' | "
                "Sort-Object RouteMetric | Select-Object -First 1).NextHop",
            ],
            capture_output=True,
        )
        gateway = result.stdout.strip()
        if gateway:
            return gateway
    return ""


def configure_proxy(backend, proxy_host, proxy_port):
    if proxy_host == "none":
        delete_setting(backend, "global", "http_proxy")
        return "disabled"
    host = proxy_host
    if proxy_host == "auto":
        host = detect_gateway_ip()
    if not host:
        return "unresolved"
    put_setting(backend, "global", "http_proxy", "{}:{}".format(host, proxy_port))
    return "{}:{}".format(host, proxy_port)


def shell_quote(value):
    return shlex.quote(value)


def apply_device_profile(backend, profile):
    if profile == "none":
        return "skipped"
    props = DEVICE_PROFILES[profile]
    resetprop = backend.exec_shell(
        "if [ -x /debug_ramdisk/resetprop ]; then echo /debug_ramdisk/resetprop; "
        "elif command -v resetprop >/dev/null 2>&1; then command -v resetprop; fi",
        check=False,
        capture_output=True,
    ).stdout.strip()
    tool = resetprop or "setprop"
    for item in props:
        key, value = item.split("=", 1)
        try_shell(
            backend,
            "{} {} {}".format(
                tool,
                shell_quote(key),
                shell_quote(value),
            ),
            root=True,
        )
    return "{} via {}".format(profile, tool)


def asset_path(*parts):
    return os.path.join(os.path.dirname(os.path.abspath(__file__)), *parts)


def install_extreme_script(backend):
    if backend.mode != "container" and not backend.has_root():
        return "root unavailable"

    local_script = asset_path("assets", "redroid_extreme_post_fs_data.sh")
    if not os.path.isfile(local_script):
        return "missing asset"

    tmp_path = "/data/local/tmp/20-redroid-extreme.sh"
    final_path = "/data/adb/post-fs-data.d/20-redroid-extreme.sh"

    if backend.mode == "container":
        run(
            [
                backend.container_runtime,
                "cp",
                local_script,
                "{}:{}".format(backend.container_name, tmp_path),
            ]
        )
    else:
        run(["adb", "-s", backend.serial, "push", local_script, tmp_path])

    try_shell(
        backend,
        "mkdir -p /data/adb/post-fs-data.d && "
        "cp {} {} && "
        "chown 0:0 {} && "
        "chmod 0755 {} && "
        "rm -f {}".format(tmp_path, final_path, final_path, final_path, tmp_path),
        root=True,
    )
    return final_path


def main():
    parser = argparse.ArgumentParser(
        description="Apply post-boot hardening to a redroid instance."
    )
    parser.add_argument(
        "--serial",
        default="127.0.0.1:5555",
        help="ADB serial, for example 127.0.0.1:5555",
    )
    parser.add_argument(
        "--container-name",
        help="Operate on a running redroid container directly instead of using ADB",
    )
    parser.add_argument(
        "--container-runtime",
        default="docker",
        choices=["docker", "podman"],
        help="Container runtime when --container-name is used",
    )
    parser.add_argument(
        "--timeout",
        type=int,
        default=300,
        help="Seconds to wait for Android to finish booting",
    )
    parser.add_argument(
        "--proxy-host",
        default="auto",
        help="Proxy host, auto, or none",
    )
    parser.add_argument(
        "--proxy-port",
        type=int,
        default=20171,
        help="Proxy port when proxy is enabled",
    )
    parser.add_argument(
        "--device-profile",
        choices=["none"] + sorted(DEVICE_PROFILES.keys()),
        default="none",
        help="Apply a runtime device profile through resetprop when available",
    )
    parser.add_argument(
        "--install-extreme-script",
        action="store_true",
        help="Install the validated 5570 post-fs-data script into /data/adb/post-fs-data.d",
    )
    args = parser.parse_args()

    backend = Backend(
        serial=args.serial,
        container_name=args.container_name or "",
        container_runtime=args.container_runtime,
    )
    backend.connect()
    backend.wait_for_boot(args.timeout)

    put_setting(backend, "global", "device_provisioned", "1")
    put_setting(backend, "secure", "user_setup_complete", "1")
    put_setting(backend, "secure", "location_mode", "3")
    put_setting(backend, "global", "package_verifier_enable", "0")
    put_setting(backend, "global", "verifier_verify_adb_installs", "0")

    try_shell(backend, "pm disable-user --user 0 com.android.setupwizard")
    try_shell(backend, "pm disable-user --user 0 com.google.android.setupwizard")
    try_shell(backend, "appops set com.android.shell android:mock_location allow")

    proxy_state = configure_proxy(backend, args.proxy_host, args.proxy_port)
    profile_state = apply_device_profile(backend, args.device_profile)
    extreme_script_state = "skipped"
    if args.install_extreme_script:
        extreme_script_state = install_extreme_script(backend)

    if backend.mode == "container":
        print("Container: {}".format(backend.container_name))
    else:
        print("ADB serial: {}".format(backend.serial))
    print("Boot status: completed")
    print("SetupWizard: disabled")
    print("Provisioned: yes")
    print("Mock location for shell: allowed")
    print("HTTP proxy: {}".format(proxy_state))
    print("Device profile: {}".format(profile_state))
    print("Extreme script: {}".format(extreme_script_state))


if __name__ == "__main__":
    main()