#!/usr/bin/env python3
"""exfat.py -- read and (minimally) write an exFAT file system in an image
file, without mounting it (mounting needs root or FUSE).

  exfat.py IMAGE[@OFFSET] info          volume label, cluster size, number of FATs
  exfat.py IMAGE[@OFFSET] ls [DIR]      list a directory ("/" by default)
  exfat.py IMAGE[@OFFSET] cat PATH      print a file's bytes to stdout
  exfat.py IMAGE[@OFFSET] put DIR       copy the files directly in DIR into
                                        the root directory (fresh volume only)

OFFSET is where the file system starts inside IMAGE, in bytes (for a
partition of a whole-disk image), e.g. foundry.img@95420416.

Foundry OS uses it for two things:
  * `put`: mkfs.exfat (exfatprogs) can format a volume but not fill it, so
    the share partition's README is written with this (`make image`).
  * `ls`/`cat`: after `make test` the host reads back what Rustbelt wrote
    to the share, independently of the kernel's driver.
fsck.exfat (exfatprogs) checks the result both times.

`put` is deliberately tiny: every file gets one contiguous run of free
clusters with the NoFatChain flag (so the FAT itself is not touched), the
same fixed timestamp, and an entry in the root directory. That is all valid
exFAT, and enough for a few files on a fresh volume. (The kernel's
build/exfat_populate.rs takes the same approach; this is a separate,
independent implementation.)
"""
import os
import struct
import sys

ATTR_DIRECTORY = 0x10
ATTR_ARCHIVE = 0x20
NO_FAT_CHAIN = 0x02
ALLOC_POSSIBLE = 0x01
# 2026-01-01 00:00:00, exFAT/DOS format: year-1980 <<25 | month <<21 | day <<16.
TIMESTAMP = (46 << 25) | (1 << 21) | (1 << 16)


def rotr16_add(acc, byte):
    return ((((acc & 1) << 15) | (acc >> 1)) + byte) & 0xFFFF


class Volume:
    def __init__(self, path, offset):
        self.path, self.offset = path, offset
        with open(path, "rb") as f:
            f.seek(offset)
            boot = f.read(512)
        if boot[3:11] != b"EXFAT   ":
            sys.exit(f"exfat.py: no exFAT file system at offset {offset} of {path}")
        self.sector = 1 << boot[108]
        self.cluster = self.sector << boot[109]
        self.fat_at = struct.unpack_from("<I", boot, 80)[0] * self.sector
        self.heap_at = struct.unpack_from("<I", boot, 88)[0] * self.sector
        self.count = struct.unpack_from("<I", boot, 92)[0]
        self.root = struct.unpack_from("<I", boot, 96)[0]
        self.num_fats = boot[110]
        self.volume_flags = struct.unpack_from("<H", boot, 106)[0]
        self.f = open(path, "r+b" if os.access(path, os.W_OK) else "rb")
        self.label, self.bitmap_cluster, self.upcase = "", None, []
        for e, _ in self.entries(self.read_chain(self.root, None, True)):
            if e[0] == 0x83:
                self.label = e[2 : 2 + 2 * e[1]].decode("utf-16-le")
            elif e[0] == 0x81:
                self.bitmap_cluster = struct.unpack_from("<I", e, 20)[0]
                self.bitmap_len = struct.unpack_from("<Q", e, 24)[0]
            elif e[0] == 0x82:
                start, length = struct.unpack_from("<IQ", e, 20)
                raw = self.read_chain(start, length, False)
                units = struct.unpack(f"<{len(raw) // 2}H", raw)
                i = 0
                while i < len(units):  # 0xFFFF n = "n identity entries"
                    if units[i] == 0xFFFF and i + 1 < len(units):
                        self.upcase.extend(range(len(self.upcase), len(self.upcase) + units[i + 1]))
                        i += 2
                    else:
                        self.upcase.append(units[i])
                        i += 1

    # --- low level -------------------------------------------------------
    def pos(self, cluster):
        return self.offset + self.heap_at + (cluster - 2) * self.cluster

    def read_at(self, at, n):
        self.f.seek(at)
        return self.f.read(n)

    def write_at(self, at, data):
        self.f.seek(at)
        self.f.write(data)

    def fat(self, cluster):
        return struct.unpack("<I", self.read_at(self.offset + self.fat_at + 4 * cluster, 4))[0]

    def clusters(self, first, length, chained):
        """The clusters of a file: consecutive (NoFatChain) or via the FAT."""
        if first == 0:
            return []
        if not chained:
            n = max(1, -(-length // self.cluster))
            return list(range(first, first + n))
        out, c = [], first
        while 2 <= c < self.count + 2 and len(out) <= self.count:
            out.append(c)
            c = self.fat(c)
        return out

    def read_chain(self, first, length, chained):
        data = b"".join(self.read_at(self.pos(c), self.cluster) for c in self.clusters(first, length or 0, chained))
        return data if length is None else data[:length]

    @staticmethod
    def entries(data):
        for i in range(0, len(data), 32):
            yield data[i : i + 32], i

    # --- directories -----------------------------------------------------
    def list_dir(self, first, length, chained):
        """[(name, attrs, first_cluster, size, chained)] of one directory."""
        data = self.read_chain(first, length, chained)
        out, i = [], 0
        while i < len(data):
            t = data[i]
            if t == 0x00:
                break
            if t == 0x85:
                secondary = data[i + 1]
                attrs = struct.unpack_from("<H", data, i + 4)[0]
                stream = data[i + 32 : i + 64]
                flags, name_len = stream[1], stream[3]
                first_c, size = struct.unpack_from("<IQ", stream, 20)
                name = b"".join(data[i + 32 * k + 2 : i + 32 * k + 32] for k in range(2, secondary + 1))
                name = name.decode("utf-16-le")[:name_len]
                out.append((name, attrs, first_c, size, not flags & NO_FAT_CHAIN))
                i += 32 * (secondary + 1)
            else:
                i += 32
        return out

    def lookup(self, path):
        cur = (self.root, None, True, ATTR_DIRECTORY)
        for part in [p for p in path.split("/") if p]:
            if not cur[3] & ATTR_DIRECTORY:
                sys.exit(f"exfat.py: {path}: not a directory")
            for name, attrs, first, size, chained in self.list_dir(cur[0], cur[1], cur[2]):
                if self.up(name) == self.up(part):
                    cur = (first, size, chained, attrs)
                    break
            else:
                sys.exit(f"exfat.py: {path}: no such file or directory")
        return cur

    def up(self, name):
        return [self.upcase[u] if u < len(self.upcase) else u for u in name.encode("utf-16-le")[::2]]

    # --- writing (fresh volumes) -----------------------------------------
    def bitmap(self):
        return bytearray(self.read_chain(self.bitmap_cluster, self.bitmap_len, False))

    def entry_set(self, name, attrs, first, size):
        units = list(struct.unpack(f"<{len(name.encode('utf-16-le')) // 2}H", name.encode("utf-16-le")))
        names = -(-len(units) // 15)
        s = bytearray(32 * (2 + names))
        s[0], s[1] = 0x85, 1 + names
        struct.pack_into("<HHIII", s, 4, attrs, 0, TIMESTAMP, TIMESTAMP, TIMESTAMP)
        h = 0
        for u in units:
            up = self.upcase[u] if u < len(self.upcase) else u
            h = rotr16_add(rotr16_add(h, up & 0xFF), up >> 8)
        s[32] = 0xC0
        s[33] = ALLOC_POSSIBLE | (NO_FAT_CHAIN if first else 0)
        s[35] = len(units)
        struct.pack_into("<HHQII Q".replace(" ", ""), s, 36, h, 0, size, 0, first, size)
        for k in range(names):
            chunk = units[15 * k : 15 * k + 15]
            s[64 + 32 * k] = 0xC1
            struct.pack_into(f"<{len(chunk)}H", s, 64 + 32 * k + 2, *chunk)
        c = 0
        for i, b in enumerate(s):
            if i not in (2, 3):
                c = rotr16_add(c, b)
        struct.pack_into("<H", s, 2, c)
        return bytes(s)

    def put(self, src_dir):
        bm = self.bitmap()
        used = lambda c: bm[(c - 2) // 8] >> ((c - 2) % 8) & 1
        nxt = max(c for c in range(2, self.count + 2) if used(c)) + 1
        root_clusters = self.clusters(self.root, None, True)
        root = bytearray(self.read_chain(self.root, None, True))
        slot = next(i for e, i in self.entries(root) if e[0] == 0x00)
        for name in sorted(os.listdir(src_dir)):
            path = os.path.join(src_dir, name)
            if not os.path.isfile(path):
                continue
            data = open(path, "rb").read()
            n = -(-len(data) // self.cluster)
            if nxt + n > self.count + 2:
                sys.exit("exfat.py: volume full")
            first = nxt if n else 0
            for c in range(nxt, nxt + n):
                bm[(c - 2) // 8] |= 1 << ((c - 2) % 8)
                self.write_at(self.pos(c), data[(c - nxt) * self.cluster : (c - nxt + 1) * self.cluster].ljust(self.cluster, b"\0"))
            nxt += n
            es = self.entry_set(name, ATTR_ARCHIVE, first, len(data))
            if slot + len(es) > len(root):
                sys.exit("exfat.py: root directory full (put only fills its first cluster(s))")
            root[slot : slot + len(es)] = es
            slot += len(es)
        for k, c in enumerate(root_clusters):
            self.write_at(self.pos(c), root[k * self.cluster : (k + 1) * self.cluster])
        for k, c in enumerate(self.clusters(self.bitmap_cluster, self.bitmap_len, False)):
            self.write_at(self.pos(c), bytes(bm[k * self.cluster : (k + 1) * self.cluster]).ljust(self.cluster, b"\0"))
        in_use = sum(used(c) for c in range(2, self.count + 2))
        # PercentInUse (boot sector byte 112) is not covered by the boot checksum.
        self.write_at(self.offset + 112, bytes([in_use * 100 // self.count]))
        self.f.flush()


def main():
    if len(sys.argv) < 3:
        sys.exit(__doc__)
    img, _, off = sys.argv[1].partition("@")
    vol = Volume(img, int(off or 0))
    cmd, args = sys.argv[2], sys.argv[3:]
    if cmd == "info":
        print(f"label {vol.label!r}, {vol.cluster}-byte clusters, {vol.count} clusters, "
              f"{vol.num_fats} FAT(s){' (TexFAT)' if vol.num_fats == 2 else ''}, volume flags {vol.volume_flags:#x}")
    elif cmd == "ls":
        first, size, chained, attrs = vol.lookup(args[0] if args else "/")
        for name, a, _, sz, _ in vol.list_dir(first, size if first != vol.root else None, chained):
            print(f"{'<DIR>' if a & ATTR_DIRECTORY else sz:>10}  {name}")
    elif cmd == "cat":
        first, size, chained, attrs = vol.lookup(args[0])
        sys.stdout.buffer.write(vol.read_chain(first, size, chained))
    elif cmd == "put":
        vol.put(args[0])
    else:
        sys.exit(__doc__)


if __name__ == "__main__":
    main()