#!/usr/bin/env python3
"""Emit SQL that copies apps' keystore2 keys from an old ROM's DB into the live one, under the new UIDs.

Only works when both ROMs use the same TEE (compare keyentry.km_uuid). Keys super-encrypted with the
old lock-screen credential (blobmetadata tag 0, EncryptedBy) cannot be moved and are skipped.

Usage: keystore-transplant.py OLD_DB [--replace] OLD_UID:NEW_UID [...] > transplant.sql
  --replace  first delete the app's keys that already exist on the new ROM (needed when it regenerated them)
Apply with the apps force-stopped: `su -c 'sqlite3 /data/misc/keystore/persistent.sqlite < transplant.sql'`, reboot.
"""

import secrets
import sqlite3
import sys

DOMAIN_APP = 0
TAG_ENCRYPTED_BY = 0


def lit(v: object) -> str:
    if v is None:
        return "NULL"
    if isinstance(v, bytes):
        return f"X'{v.hex()}'"
    if isinstance(v, str):
        return "'" + v.replace("'", "''") + "'"
    return str(v)


def delete_keys(where: str) -> list[str]:
    ids = f"(SELECT id FROM keyentry WHERE {where})"
    return [f"DELETE FROM blobmetadata WHERE blobentryid IN (SELECT id FROM blobentry WHERE keyentryid IN {ids});",
            f"DELETE FROM blobentry WHERE keyentryid IN {ids};",
            f"DELETE FROM keyparameter WHERE keyentryid IN {ids};",
            f"DELETE FROM keymetadata WHERE keyentryid IN {ids};",
            f"DELETE FROM grant WHERE keyentryid IN {ids};",
            f"DELETE FROM keyentry WHERE {where};"]


args = sys.argv[1:]
old = sqlite3.connect(f"file:{args.pop(0)}?mode=ro", uri=True)
replace = "--replace" in args
pairs = [tuple(map(int, a.split(":"))) for a in args if a != "--replace"]
# Older keystore2 schemas have no blobentry.state; every blob there is current (0).
blob_state = "state" if any(c[1] == "state" for c in old.execute("PRAGMA table_info(blobentry)")) else "0"

out = ["BEGIN;"]
for old_ns, new_ns in pairs:
    if replace:
        out += delete_keys(f"domain={DOMAIN_APP} AND namespace={new_ns}")
    keys = old.execute("SELECT id, key_type, alias, state, km_uuid FROM keyentry WHERE domain=? AND namespace=?",
                       (DOMAIN_APP, old_ns)).fetchall()
    for kid, key_type, alias, state, km_uuid in keys:
        blobs = old.execute(f"SELECT id, subcomponent_type, blob, {blob_state} FROM blobentry WHERE keyentryid=?", (kid,)).fetchall()
        if any(old.execute("SELECT 1 FROM blobmetadata WHERE blobentryid=? AND tag=?", (b[0], TAG_ENCRYPTED_BY)).fetchone()
               for b in blobs):
            print(f"-- skip {old_ns} {alias!r}: super-encrypted with the old lock-screen credential", file=sys.stderr)
            continue
        new_id = secrets.randbits(63) * (1 if secrets.randbits(1) else -1)
        # The old key wins over one the app generated under the same alias on the new ROM.
        out += delete_keys(f"domain={DOMAIN_APP} AND namespace={new_ns} AND alias={lit(alias)}")
        out.append(f"INSERT INTO keyentry (id,key_type,domain,namespace,alias,state,km_uuid) VALUES "
                   f"({new_id},{key_type},{DOMAIN_APP},{new_ns},{lit(alias)},{state},{lit(km_uuid)});")
        for bid, sub, blob, bstate in blobs:
            out.append(f"INSERT INTO blobentry (subcomponent_type,keyentryid,blob,state) VALUES ({sub},{new_id},{lit(blob)},{bstate});")
            for tag, data in old.execute("SELECT tag, data FROM blobmetadata WHERE blobentryid=?", (bid,)):
                out.append(f"INSERT INTO blobmetadata (blobentryid,tag,data) VALUES ((SELECT max(id) FROM blobentry WHERE keyentryid={new_id}),{tag},{lit(data)});")
        for tag, data, level in old.execute("SELECT tag, data, security_level FROM keyparameter WHERE keyentryid=?", (kid,)):
            out.append(f"INSERT INTO keyparameter (keyentryid,tag,data,security_level) VALUES ({new_id},{tag},{lit(data)},{level});")
        for tag, data in old.execute("SELECT tag, data FROM keymetadata WHERE keyentryid=?", (kid,)):
            out.append(f"INSERT INTO keymetadata (keyentryid,tag,data) VALUES ({new_id},{tag},{lit(data)});")
        print(f"-- {old_ns}->{new_ns} {alias!r}: {len(blobs)} blobs", file=sys.stderr)
out.append("COMMIT;")
print("\n".join(out))