#!/usr/bin/env python3
"""Minimal Qualcomm DIAG EFS2 client over a serial diag port.

Packet layouts and opcodes replicate JohnBel/EfsTools exactly. Used to
back up and replace policyman RF-band config on a rooted Sony Xperia 1 VI.
"""
import sys, os, time, struct, glob, serial

DIAG_SUBSYS = 0x4B
SUBSYS_EFS = 19
# EFS opcodes (EfsTools QcdmEfsCommand)
HELLO, QUERY, OPEN, CLOSE, READ, WRITE = 0, 1, 2, 3, 4, 5
UNLINK, MKDIR, OPENDIR, READDIR, CLOSEDIR, STAT = 8, 9, 11, 12, 13, 15
# open flags
O_WRONLY, O_RDONLY, O_CREAT = 0o1, 0o0, 0o100

# CRC-16 CCITT reflected (HDLC FCS-16)
_TAB = []
for _i in range(256):
    _c = _i
    for _ in range(8):
        _c = (_c >> 1) ^ 0x8408 if (_c & 1) else (_c >> 1)
    _TAB.append(_c)

def fcs16(data):
    fcs = 0xFFFF
    for b in data:
        fcs = (fcs >> 8) ^ _TAB[(fcs ^ b) & 0xFF]
    return (~fcs) & 0xFFFF

def hdlc_encode(payload):
    crc = fcs16(payload)
    frame = bytes(payload) + struct.pack("<H", crc)
    out = bytearray()
    for b in frame:
        if b == 0x7E: out += b"\x7d\x5e"
        elif b == 0x7D: out += b"\x7d\x5d"
        else: out.append(b)
    out.append(0x7E)
    return bytes(out)

def hdlc_decode(frame):
    frame = frame.rstrip(b"\x7e")
    out = bytearray(); esc = False
    for b in frame:
        if esc: out.append(b ^ 0x20); esc = False
        elif b == 0x7D: esc = True
        else: out.append(b)
    return bytes(out)

def _responds_to_diag(port, timeout=1.0):
    """Return True if port answers a DIAG version request (cmd 0x00)."""
    try:
        s = serial.Serial(port, 115200, timeout=0.3)
    except Exception:
        return False
    try:
        s.reset_input_buffer()
        s.write(hdlc_encode([0x00]))
        buf, deadline = bytearray(), time.time() + timeout
        while time.time() < deadline:
            b = s.read(1)
            if not b:
                continue
            if b[0] == 0x7E:
                if buf:
                    dec = hdlc_decode(bytes(buf))
                    buf = bytearray()
                    if len(dec) >= 3 and dec[0] == 0x00 and \
                       fcs16(dec[:-2]) == struct.unpack("<H", dec[-2:])[0]:
                        return True
            else:
                buf.append(b[0])
        return False
    finally:
        s.close()

def find_diag_port():
    """Find the DIAG serial port. DIAGPORT wins; otherwise probe candidates,
    interface 0 (-if00-) first, and return the one that answers DIAG."""
    env = os.environ.get("DIAGPORT")
    if env:
        return env
    candidates = (sorted(glob.glob("/dev/serial/by-id/*-if00-port0")) +
                  sorted(glob.glob("/dev/serial/by-id/*")) +
                  sorted(glob.glob("/dev/ttyUSB*")))
    seen = set()
    for p in candidates:
        real = os.path.realpath(p)
        if real in seen:
            continue
        seen.add(real)
        if _responds_to_diag(p):
            return p
    raise IOError("no responding DIAG port found; is the phone in diag mode "
                  "(setprop sys.usb.config diag,serial_cdev,rmnet,adb)?")

class Diag:
    def __init__(self, port=None):
        self.port = port or find_diag_port()
        self.s = serial.Serial(self.port, 115200, timeout=0.4)
        self.s.reset_input_buffer()

    def _read_frame(self, timeout=4.0):
        buf = bytearray(); deadline = time.time() + timeout
        while time.time() < deadline:
            b = self.s.read(1)
            if not b:
                continue
            if b[0] == 0x7E:
                if not buf:
                    continue  # leading terminator
                dec = hdlc_decode(bytes(buf))
                buf = bytearray()
                if len(dec) < 3:
                    continue
                payload, crc = dec[:-2], struct.unpack("<H", dec[-2:])[0]
                if fcs16(payload) != crc:
                    continue  # corrupt / partial, keep reading
                return payload
            else:
                buf.append(b[0])
        raise TimeoutError("no diag frame")

    def subsys(self, efs_cmd, body=b"", timeout=4.0):
        """Send an EFS2 subsystem command, return matching response payload."""
        req = struct.pack("<BBH", DIAG_SUBSYS, SUBSYS_EFS, efs_cmd) + body
        self.s.write(hdlc_encode(req))
        deadline = time.time() + timeout
        while time.time() < deadline:
            p = self._read_frame(timeout=timeout)
            if len(p) >= 4 and p[0] == DIAG_SUBSYS and p[1] == SUBSYS_EFS \
               and struct.unpack("<H", p[2:4])[0] == efs_cmd:
                return p
            # otherwise async log/event frame -> ignore
        raise TimeoutError("no matching response for efs cmd %d" % efs_cmd)

    # ---- EFS operations ----
    def hello(self):
        win, wbyte, ver = 0x100000, 0x100000, 1
        body = struct.pack("<IIIIIIiii", win, wbyte, win, wbyte, win, wbyte, ver, ver, ver)
        body += b"\xff\xff\xff\xff"
        p = self.subsys(HELLO, body)
        return struct.unpack("<i", p[28:32])[0]  # version

    def stat(self, name):
        p = self.subsys(STAT, name.encode() + b"\x00")
        err = struct.unpack("<i", p[4:8])[0]
        mode = struct.unpack("<i", p[8:12])[0]
        size = struct.unpack("<i", p[12:16])[0]
        return err, mode, size

    def opendir(self, name):
        p = self.subsys(OPENDIR, name.encode() + b"\x00")
        d = struct.unpack("<i", p[4:8])[0]
        err = struct.unpack("<i", p[8:12])[0]
        return d, err

    def readdir(self, d, seq):
        p = self.subsys(READDIR, struct.pack("<ii", d, seq))
        err = struct.unpack("<i", p[12:16])[0]
        etype = struct.unpack("<i", p[16:20])[0]
        size = struct.unpack("<i", p[24:28])[0]
        name = p[40:].split(b"\x00", 1)[0].decode("ascii", "replace")
        return err, etype, size, name

    def closedir(self, d):
        self.subsys(CLOSEDIR, struct.pack("<i", d))

    def listdir(self, path):
        d, err = self.opendir(path)
        if err != 0 or d < 0:
            raise IOError("opendir %s err=%d" % (path, err))
        entries = []; seq = 1
        while True:
            err, etype, size, name = self.readdir(d, seq)
            if err != 0 or not name:
                break
            entries.append((name, etype, size)); seq += 1
        self.closedir(d)
        return entries

    def open(self, name, flags, perm):
        body = struct.pack("<ii", flags, perm) + name.encode() + b"\x00"
        p = self.subsys(OPEN, body)
        fd = struct.unpack("<i", p[4:8])[0]
        err = struct.unpack("<i", p[8:12])[0]
        return fd, err

    def close(self, fd):
        self.subsys(CLOSE, struct.pack("<i", fd))

    def read(self, fd, size, offset):
        p = self.subsys(READ, struct.pack("<iII", fd, size, offset))
        nread = struct.unpack("<i", p[12:16])[0]
        err = struct.unpack("<i", p[16:20])[0]
        return err, p[20:20 + nread]

    def write(self, fd, offset, chunk):
        body = struct.pack("<iI", fd, offset) + chunk
        p = self.subsys(WRITE, body)
        nwr = struct.unpack("<i", p[12:16])[0]
        err = struct.unpack("<i", p[16:20])[0]
        return err, nwr

    def unlink(self, name):
        p = self.subsys(UNLINK, name.encode() + b"\x00")
        return struct.unpack("<i", p[4:8])[0]

    def read_file(self, name):
        err, mode, size = self.stat(name)
        if err != 0:
            raise IOError("stat %s err=%d" % (name, err))
        fd, err = self.open(name, O_RDONLY, 0)
        if err != 0 or fd < 0:
            raise IOError("open(r) %s err=%d" % (name, err))
        data = bytearray(); off = 0
        while off < size:
            err, chunk = self.read(fd, min(1024, size - off), off)
            if err != 0:
                self.close(fd); raise IOError("read %s err=%d" % (name, err))
            if not chunk:
                break
            data += chunk; off += len(chunk)
        self.close(fd)
        return bytes(data)

    def write_file(self, name, data):
        # replicate EfsTools: delete existing, then create+write
        err, mode, size = self.stat(name)
        if err == 0:
            self.unlink(name)
        fd, err = self.open(name, O_WRONLY | O_CREAT, 0o777)
        if err != 0 or fd < 0:
            raise IOError("open(w) %s err=%d" % (name, err))
        off = 0
        while off < len(data):
            chunk = data[off:off + 1024]
            err, nwr = self.write(fd, off, chunk)
            if err != 0 or nwr <= 0:
                self.close(fd); raise IOError("write %s err=%d nwr=%d" % (name, err, nwr))
            off += nwr
        self.close(fd)
        return off


TARGET_DIR = "/policyman"
TARGETS = ["band_set_01.xml", "plmn_mcc_supported_01.xml", "policies.xml"]
SRC = "/tmp/claude-1000/-home-grmrgecko-Downloads-XperiFirm-5-8-1--by-Igor-Eisberg-/aa2c5f73-89ed-4884-8a6f-7f146ae4f002/scratchpad/xda_xml"
BACKUP = "/tmp/claude-1000/-home-grmrgecko-Downloads-XperiFirm-5-8-1--by-Igor-Eisberg-/aa2c5f73-89ed-4884-8a6f-7f146ae4f002/scratchpad/policyman_backup"

def main():
    cmd = sys.argv[1] if len(sys.argv) > 1 else "probe"
    d = Diag()
    print("using diag port %s" % d.port)
    ver = d.hello()
    print("EFS hello ok, version=%d" % ver)

    if cmd == "probe":
        entries = d.listdir(TARGET_DIR)
        print("%s: %d entries" % (TARGET_DIR, len(entries)))
        for name, etype, size in entries:
            mark = " <== TARGET" if name in TARGETS else ""
            print("  %-40s type=%d size=%d%s" % (name, etype, size, mark))
        for t in TARGETS:
            err, mode, size = d.stat("%s/%s" % (TARGET_DIR, t))
            print("stat %s -> err=%d mode=%o size=%d" % (t, err, mode & 0xffff, size))

    elif cmd == "backup":
        os.makedirs(BACKUP, exist_ok=True)
        for t in TARGETS:
            data = d.read_file("%s/%s" % (TARGET_DIR, t))
            open("%s/%s" % (BACKUP, t), "wb").write(data)
            print("backed up %s (%d bytes)" % (t, len(data)))

    elif cmd == "write":
        for t in TARGETS:
            data = open("%s/%s" % (SRC, t), "rb").read()
            n = d.write_file("%s/%s" % (TARGET_DIR, t), data)
            print("wrote %s (%d bytes)" % (t, n))

    elif cmd == "verify":
        ok = True
        for t in TARGETS:
            want = open("%s/%s" % (SRC, t), "rb").read()
            got = d.read_file("%s/%s" % (TARGET_DIR, t))
            same = want == got
            ok = ok and same
            print("verify %s: %s (device %d bytes, source %d bytes)" %
                  (t, "MATCH" if same else "DIFFER", len(got), len(want)))
        print("ALL VERIFIED" if ok else "MISMATCH!")

    d.s.close()

if __name__ == "__main__":
    main()