#!/usr/bin/env python3
"""Restore user apps, app data and permissions onto a freshly flashed, rooted phone.

Sources (from the backup folder):
  apks/<pkg>/*.apk                     APK + splits
  rooted-fresh/dumpsys-package.txt     installer, granted runtime permissions
  rooted-fresh/packages-3.txt          which packages are user apps
  rooted-fresh/packages-disabled.txt   frozen/disabled state
  rooted/apps-manifest.json            special app-ops (overlay, all-files, ...)
  <work>/data/<pkg>, <work>/user_de/0/<pkg>   unpacked data-fresh.tar

The device comes from ANDROID_SERIAL, as for adb. Usage: restore.py BACKUP_DIR WORK_DIR [--data-only] [pkg ...]   (no pkgs = all user apps)
  --data-only  only put the data back (apps already installed and updated): keeps the current APK,
               installer, permissions and frozen state. Never re-run a full restore over updated apps.
Every action and its result is appended to WORK_DIR/restore.log.
"""

import json
import re
import shlex
import subprocess
import sys
import time
from pathlib import Path

SKIP = {"com.topjohnwu.magisk"}  # installed fresh, its data must not be reused
# Dropped on restore: regenerated by the app, and the bulk of most backups.
VOLATILE = ["cache", "code_cache", "app_webview/Default/Cache", "app_webview/Default/Code Cache"]

backup, work = Path(sys.argv[1]), Path(sys.argv[2])
data_only = "--data-only" in sys.argv[3:]
only = set(sys.argv[3:]) - {"--data-only"}
log_file = (work / "restore.log").open("a")


def log(msg: str) -> None:
    line = f"{time.strftime('%H:%M:%S')} {msg}"
    print(line, flush=True)
    log_file.write(line + "\n")
    log_file.flush()


def adb(*args: str, timeout: int = 300, stdin=None) -> subprocess.CompletedProcess:
    return subprocess.run(["adb", *args], capture_output=True, text=stdin is None,
                          timeout=timeout, stdin=stdin)


def su(cmd: str, timeout: int = 300) -> str:
    r = adb("shell", f"su -c {shlex.quote(cmd)}", timeout=timeout)
    return (r.stdout + r.stderr).strip()


def parse_dumpsys(text: str) -> dict:
    """installer + granted runtime permissions per package."""
    out = {}
    for block in re.split(r"\n  Package \[", text)[1:]:
        name = block.split("]", 1)[0]
        installer = re.search(r"\n    installerPackageName=(\S+)", block)
        runtime = block.split("runtime permissions:", 1)
        perms = re.findall(r"\n        ([\w.]+): granted=true", runtime[1]) if len(runtime) > 1 else []
        out[name] = {"installer": installer.group(1) if installer else None, "perms": sorted(set(perms))}
    return out


fresh = backup / "rooted-fresh"
meta = parse_dumpsys((fresh / "dumpsys-package.txt").read_text(errors="replace"))
manifest = json.loads((backup / "rooted/apps-manifest.json").read_text())["apps"]
user_pkgs = sorted(set(re.findall(r"=([\w.]+)\s+installer=", (fresh / "packages-3.txt").read_text())) - SKIP)
disabled = set(re.findall(r"package:([\w.]+)", (fresh / "packages-disabled.txt").read_text()))
targets = [p for p in user_pkgs if not only or p in only]
log(f"=== restore start: {len(targets)} packages")

# Pass 1: install every APK set, so installers like Obtainium exist before pass 2 names them.
for pkg in [] if data_only else targets:
    apks = sorted(str(p) for p in (backup / "apks" / pkg).glob("*.apk"))
    if not apks:
        log(f"{pkg}: NO APK, skipped")
        continue
    # --bypass-low-target-sdk-block: Android 14+ refuses targetSdk < 24 (Automagic workarounds).
    r = subprocess.run(["adb", "install-multiple", "-r", "-d", "-g",
                        "--bypass-low-target-sdk-block", *apks],
                       capture_output=True, text=True, timeout=900)
    log(f"{pkg}: install {'OK' if 'Success' in r.stdout else 'FAILED ' + (r.stdout + r.stderr).strip()[-200:]}")

installed = set(re.findall(r"package:([\w.]+)", adb("shell", "pm list packages").stdout))

# Pass 2: data, ownership, SELinux labels, installer, permissions, disabled state.
for pkg in targets:
    if pkg not in installed:
        continue
    info = meta.get(pkg, {})
    su(f"am force-stop {pkg}")
    uid = su(f"stat -c %u /data/data/{pkg}")
    if not uid.isdigit():
        log(f"{pkg}: no data dir after install ({uid}), data skipped")
        continue
    cache_gid = int(uid) - 10000 + 20000  # u0_aNNN_cache

    members = [m for m in (f"data/{pkg}", f"user_de/0/{pkg}") if (work / m).is_dir()]
    if members:
        for m in members:
            su(f"find /data/{m} -mindepth 1 -maxdepth 1 ! -name lib -exec rm -rf {{}} +")
        excludes = [f"--exclude={m}/{v}" for m in members for v in VOLATILE]
        tar = subprocess.Popen(["tar", "-cf", "-", "-C", str(work), *excludes, *members],
                               stdout=subprocess.PIPE)
        r = adb("exec-in", "su -c 'tar -xf - -C /data'", stdin=tar.stdout, timeout=1800)
        tar.wait()
        # exec-in returns once our side is sent; the phone's tar may still be writing.
        su("while pgrep -x tar >/dev/null; do sleep 0.3; done", timeout=600)
        for m in members:
            su(f"chown -R {uid}:{uid} /data/{m}; "
               f"for d in cache code_cache; do mkdir -p /data/{m}/$d; chown {uid}:{cache_gid} /data/{m}/$d; "
               # restorecon gives only app_data_file:s0; the per-app MLS categories that
               # installd put on the top directory must be copied onto its contents.
               f"chmod 2771 /data/{m}/$d; done; chcon -hR \"$(stat -c %C /data/{m})\" /data/{m}",
               timeout=900)
        wrong = su(f"find {' '.join('/data/' + m for m in members)} ! -user {uid} | wc -l")
        log(f"{pkg}: data {'+'.join(members)} rc={r.returncode} uid={uid} wrong_owner={wrong}")
    else:
        log(f"{pkg}: no data in backup")
    if data_only:
        continue

    inst = info.get("installer")
    if inst and inst != "null" and inst in installed | {"com.android.vending"}:
        log(f"{pkg}: installer {inst} -> {su(f'pm set-installer {pkg} {inst}') or 'OK'}")

    # -g granted everything at install; revoke what was not granted before.
    wanted = set(info.get("perms", []))
    now = re.findall(r"\n\s+([\w.]+): granted=true", su(f"dumpsys package {pkg} | sed -n '/runtime permissions:/,$p'"))
    revoked = [p for p in set(now) - wanted]
    for p in revoked:
        su(f"pm revoke {pkg} {p}")
    for p in wanted - set(now):
        su(f"pm grant {pkg} {p}")
    ops = manifest.get(pkg, {}).get("specialAppOps", {})
    for op, mode in ops.items():
        su(f"appops set {pkg} {op} {mode}")
    log(f"{pkg}: perms granted={len(wanted)} revoked={len(revoked)} appops={','.join(ops) or '-'}")

    if pkg in disabled:
        log(f"{pkg}: disable-user -> {su(f'pm disable-user --user 0 {pkg}')}")

log("=== restore done")