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

Command-line tool to list, pull, push and back up files in a rooted phone's
modem EFS. Packet layouts and opcodes replicate JohnBel/EfsTools exactly.
Originally written to back up and replace the policyman RF-band config on a
Sony Xperia 1 VI. Run with -h for usage.
"""
import sys, os, time, struct, glob, argparse, 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


TYPE_NAMES = {0: "file", 1: "dir", 2: "link", 3: "immovable", 15: "item"}


def cmd_ls(d, args):
    entries = d.listdir(args.path)
    for name, etype, size in entries:
        print("%-42s %-9s %8d" % (name, TYPE_NAMES.get(etype, str(etype)), size))
    print("%d entries in %s" % (len(entries), args.path))


def cmd_stat(d, args):
    err, mode, size = d.stat(args.path)
    if err:
        raise SystemExit("stat %s: efs error %d" % (args.path, err))
    print("%s  mode=%o  size=%d" % (args.path, mode & 0xffff, size))


def cmd_pull(d, args):
    local = args.local or os.path.basename(args.path.rstrip("/"))
    data = d.read_file(args.path)
    with open(local, "wb") as f:
        f.write(data)
    print("pulled %s -> %s (%d bytes)" % (args.path, local, len(data)))


def cmd_push(d, args):
    with open(args.local, "rb") as f:
        data = f.read()
    n = d.write_file(args.path, data)
    verified = d.read_file(args.path) == data
    print("pushed %s -> %s (%d bytes), read-back %s"
          % (args.local, args.path, n, "matches" if verified else "MISMATCH"))
    if not verified:
        raise SystemExit(2)


def cmd_backup(d, args):
    written = [0]

    def walk(efs_dir, local_dir):
        os.makedirs(local_dir, exist_ok=True)
        for name, etype, size in d.listdir(efs_dir):
            efs_path = efs_dir.rstrip("/") + "/" + name
            local_path = os.path.join(local_dir, name)
            if etype == 1:  # directory
                walk(efs_path, local_path)
            elif etype in (0, 15):  # regular file or item file
                try:
                    data = d.read_file(efs_path)
                except Exception as e:
                    sys.stderr.write("  skip %s: %s\n" % (efs_path, e))
                    continue
                with open(local_path, "wb") as f:
                    f.write(data)
                written[0] += 1
                print("  %s (%d bytes)" % (efs_path, len(data)))

    walk(args.path, args.dest)
    print("backed up %d files from %s to %s" % (written[0], args.path, args.dest))


def build_parser():
    p = argparse.ArgumentParser(
        prog="efs2.py",
        description="Minimal Qualcomm DIAG EFS2 client: list, pull, push and "
                    "back up files in a rooted phone's modem EFS.",
        epilog="Put the phone in diag mode first (adb shell su -c 'setprop "
               "sys.usb.config diag,serial_cdev,rmnet,adb') and bind the option "
               "driver. The diag port is auto-detected; override with --port or "
               "the DIAGPORT env var.")
    p.add_argument("-p", "--port", help="diag serial port (default: auto-detect)")
    sub = p.add_subparsers(dest="command", required=True)

    s = sub.add_parser("ls", help="list an EFS directory")
    s.add_argument("path")
    s.set_defaults(func=cmd_ls)

    s = sub.add_parser("stat", help="stat an EFS file")
    s.add_argument("path")
    s.set_defaults(func=cmd_stat)

    s = sub.add_parser("pull", help="download one EFS file")
    s.add_argument("path", help="EFS path, e.g. /policyman/band_set_01.xml")
    s.add_argument("local", nargs="?", help="local path (default: the basename)")
    s.set_defaults(func=cmd_pull)

    s = sub.add_parser("push", help="upload/overwrite one EFS file, then verify")
    s.add_argument("local", help="local file to upload")
    s.add_argument("path", help="destination EFS path")
    s.set_defaults(func=cmd_push)

    s = sub.add_parser("backup", help="recursively download an EFS directory")
    s.add_argument("path", help="EFS directory, e.g. /policyman")
    s.add_argument("dest", help="local directory to write into")
    s.set_defaults(func=cmd_backup)
    return p


def main():
    args = build_parser().parse_args()
    d = Diag(args.port)
    sys.stderr.write("using diag port %s (efs v%d)\n" % (d.port, d.hello()))
    try:
        args.func(d, args)
    finally:
        d.s.close()


if __name__ == "__main__":
    main()