#!/usr/bin/env python3
# Aurora: Sahara command-mode probe (serial/HWID/PK_HASH), no storage access.
# Uploads a RAM-only firehose loader ONLY if --upload is given AND the PK_HASH
# is the Qualcomm test root ("old" = unfused). Otherwise returns device to
# image-transfer-pending state and exits.
import sys, json, logging, time
from edlclient.Library.Connection.usblib import usb_class
from edlclient.Library.sahara import sahara
from edlclient.Library.sahara_defs import cmd_t, sahara_mode_t
from edlclient.Config.qualcomm_config import root_cert_hash

UNFUSED = root_cert_hash["old"]
out = sys.argv[1]
loader = sys.argv[2] if len(sys.argv) > 2 else None
upload = "--upload" in sys.argv

cdc = usb_class(portconfig=[[0x05c6, 0x9008, -1]], loglevel=logging.INFO)
if not cdc.connect():
    sys.exit("no 05c6:9008")
s = sahara(cdc, loglevel=logging.INFO)
s.programmer = "none"  # disable loaderdb auto-selection
r = s.connect()
print("connect:", {k: v for k, v in r.items() if k != "data"})
if r.get("mode") != "sahara" or r.get("cmd") != cmd_t.SAHARA_HELLO_REQ:
    sys.exit("device not at Sahara HELLO; not touching it further")
h = r["data"]
ver = h.version
s.cmd_info(version=ver)
info = dict(sahara_version=ver, version_supported=h.version_supported,
            hello_mode=h.mode, cmd_packet_length=h.cmd_packet_length,
            serial=s.serials, hwid=s.hwidstr, msm_id=s.msm_str, oem_id=s.oem_str,
            model_id=s.model_id, pkhash=s.pkhash,
            pkhash_is_qualcomm_test_root=(s.pkhash == UNFUSED))
try:
    sv = s.cmdexec_get_sbl_version()
    info["sbl_sw_version"] = f"{sv:08x}"
except Exception as e:
    info["sbl_sw_version"] = f"err: {e}"
json.dump(info, open(out, "w"), indent=2)
print(json.dumps(info, indent=2))

# back to image transfer pending -> target re-sends HELLO
s.cmd_modeswitch(sahara_mode_t.SAHARA_MODE_IMAGE_TX_PENDING)
r = s.connect()
print("after modeswitch:", {k: v for k, v in r.items() if k != "data"})
if info["serial"] != "<EMMC_SERIAL>":
    sys.exit("serial mismatch (expected Aurora <EMMC_SERIAL>) -> STOP")
if upload:
    if not info["pkhash_is_qualcomm_test_root"]:
        sys.exit("PK_HASH is not Qualcomm test root -> NOT uploading. STOP.")
    s.programmer = loader
    m = s.upload_loader(version=ver)
    print("upload_loader ->", repr(m))