#!/usr/bin/env python3
"""Format a real exFAT image from scratch — the danos exFAT test volume.

Pure Python 3 standard library (no mkfs.exfat / mtools). It writes a valid exFAT
filesystem — a Main Boot Sector + its boot-region checksum + a backup region, the
32-bit FAT, an allocation bitmap, an up-case table (with its checksum), and a root
directory whose entry sets a real exFAT reader (and the danos exfat engine) mount
and walk. Mirrors tools/make-fat-image.py in spirit.

    make-exfat-image.py [--serial <hex>] [--label <name>] <out.img> <size-MiB>
    make-exfat-image.py --verify <out.img>

The image seeds one file, HELLO.TXT, so a mount can be proven by reading it.
"""

import struct
import sys

SECTOR = 512
UPCASE_UNITS = 256  # a-z -> A-Z, the rest identity; covers ASCII names


def align_up(value, to):
    return (value + to - 1) // to * to


def rotr16(v):
    return ((v >> 1) | (v << 15)) & 0xFFFF


def rotr32(v):
    return ((v >> 1) | (v << 31)) & 0xFFFFFFFF


def boot_checksum(region):
    """32-bit rotate-right sum over the boot region, skipping VolumeFlags
    (106,107) and PercentInUse (112) of the first sector."""
    checksum = 0
    for i, byte in enumerate(region):
        if i in (106, 107, 112):
            continue
        checksum = (rotr32(checksum) + byte) & 0xFFFFFFFF
    return checksum


def upcase_checksum(table_bytes):
    checksum = 0
    for byte in table_bytes:
        checksum = (rotr32(checksum) + byte) & 0xFFFFFFFF
    return checksum


def set_checksum(entries):
    """16-bit rotate-right sum over a directory-entry set, skipping its own two
    checksum bytes (offset 2..3 of the first entry)."""
    checksum = 0
    for i, byte in enumerate(entries):
        if i in (2, 3):
            continue
        checksum = (rotr16(checksum) + byte) & 0xFFFF
    return checksum


def name_hash(upcased_units):
    h = 0
    for unit in upcased_units:
        h = (rotr16(h) + (unit & 0xFF)) & 0xFFFF
        h = (rotr16(h) + (unit >> 8)) & 0xFFFF
    return h


def ascii_upper(unit):
    return unit - ord("a") + ord("A") if ord("a") <= unit <= ord("z") else unit


def solve_geometry(total_sectors, spc):
    """Solve for cluster count / FAT length / heap offset that fit. The FAT sits
    after the main + backup boot regions (24 sectors)."""
    fat_offset = 24
    fat_length = 1
    while True:
        heap_offset = align_up(fat_offset + fat_length, spc)
        cluster_count = (total_sectors - heap_offset) // spc
        needed = ((cluster_count + 2) * 4 + SECTOR - 1) // SECTOR
        if needed <= fat_length:
            return cluster_count, fat_offset, fat_length, heap_offset
        fat_length = needed


class ExfatImage:
    def __init__(self, size_mib, volume_id=0x1234ABCD, label="DANOS"):
        self.total_sectors = size_mib * 1024 * 1024 // SECTOR
        self.spc = 8  # 4 KiB clusters
        self.volume_id = volume_id & 0xFFFFFFFF
        self.label = label
        self.cluster_count, self.fat_offset, self.fat_length, self.heap_offset = solve_geometry(self.total_sectors, self.spc)
        if self.cluster_count < 16:
            sys.exit(f"error: image too small for exFAT ({self.cluster_count} clusters)")
        self.cluster_bytes = self.spc * SECTOR
        # Layout: the allocation bitmap (as many clusters as it needs — one per
        # 8*cluster_bytes clusters of the volume), then the up-case table, the root
        # directory, and the seeded file. A single-cluster bitmap (small volumes,
        # e.g. the 48 MiB fixture) puts root at cluster 4, as before.
        self.bitmap_bytes = (self.cluster_count + 7) // 8
        self.bitmap_clusters = (self.bitmap_bytes + self.cluster_bytes - 1) // self.cluster_bytes
        self.bitmap_cluster = 2
        self.upcase_cluster = self.bitmap_cluster + self.bitmap_clusters
        self.root_cluster = self.upcase_cluster + 1
        self.hello_cluster = self.root_cluster + 1
        self.image = bytearray(self.total_sectors * SECTOR)

    def cluster_offset(self, cluster):
        return (self.heap_offset + (cluster - 2) * self.spc) * SECTOR

    def set_fat(self, cluster, value):
        struct.pack_into("<I", self.image, self.fat_offset * SECTOR + cluster * 4, value)

    def mark_allocated(self, cluster):
        bit = cluster - 2
        pos = self.cluster_offset(2) + bit // 8
        self.image[pos] |= 1 << (bit % 8)

    def main_boot_sector(self):
        sector = bytearray(SECTOR)
        sector[0:3] = b"\xEB\x76\x90"          # jump boot
        sector[3:11] = b"EXFAT   "             # filesystem name
        # 11..64 MustBeZero (already zero)
        struct.pack_into("<Q", sector, 72, self.total_sectors)   # volume length
        struct.pack_into("<I", sector, 80, self.fat_offset)      # fat offset
        struct.pack_into("<I", sector, 84, self.fat_length)      # fat length
        struct.pack_into("<I", sector, 88, self.heap_offset)     # cluster heap offset
        struct.pack_into("<I", sector, 92, self.cluster_count)   # cluster count
        struct.pack_into("<I", sector, 96, self.root_cluster)    # first cluster of root
        struct.pack_into("<I", sector, 100, self.volume_id)      # volume serial number
        struct.pack_into("<H", sector, 104, 0x0100)              # filesystem revision 1.0
        sector[108] = 9                                          # bytes per sector shift (512)
        sector[109] = self.spc.bit_length() - 1                  # sectors per cluster shift
        sector[110] = 1                                          # number of FATs
        sector[111] = 0x80                                       # drive select
        sector[112] = 0xFF                                       # percent in use (unknown)
        sector[510] = 0x55
        sector[511] = 0xAA
        return sector

    def build(self):
        # Main boot region (sectors 0..11): VBR, eight extended boot sectors, OEM
        # parameters, reserved, then the checksum sector.
        vbr = self.main_boot_sector()
        self.image[0:SECTOR] = vbr
        for s in range(1, 9):  # extended boot sectors carry the 0xAA550000 signature
            struct.pack_into("<I", self.image, s * SECTOR + 508, 0xAA550000)
        # sectors 9 (OEM) and 10 (reserved) stay zero
        region = bytes(self.image[0 : 11 * SECTOR])
        checksum = boot_checksum(region)
        for i in range(SECTOR // 4):
            struct.pack_into("<I", self.image, 11 * SECTOR + i * 4, checksum)
        # Backup boot region (sectors 12..23) is a copy of 0..11.
        self.image[12 * SECTOR : 24 * SECTOR] = self.image[0 : 12 * SECTOR]

        # FAT: reserved entries, then a single-cluster chain per metadata object,
        # except the bitmap which spans self.bitmap_clusters (a real FAT chain).
        self.set_fat(0, 0xFFFFFFF8)
        self.set_fat(1, 0xFFFFFFFF)
        used = []
        for i in range(self.bitmap_clusters):
            cluster = self.bitmap_cluster + i
            self.set_fat(cluster, 0xFFFFFFFF if i == self.bitmap_clusters - 1 else cluster + 1)
            used.append(cluster)
        for cluster in (self.upcase_cluster, self.root_cluster, self.hello_cluster):
            self.set_fat(cluster, 0xFFFFFFFF)
            used.append(cluster)

        # Allocation bitmap: every metadata/file cluster in use.
        for cluster in used:
            self.mark_allocated(cluster)

        # Up-case table: 256 explicit units, a-z -> A-Z.
        upcase = bytearray(UPCASE_UNITS * 2)
        for i in range(UPCASE_UNITS):
            struct.pack_into("<H", upcase, i * 2, ascii_upper(i))
        off = self.cluster_offset(self.upcase_cluster)
        self.image[off : off + len(upcase)] = upcase
        table_checksum = upcase_checksum(upcase)

        # Seed file HELLO.TXT (contiguous, one cluster).
        content = b"exfat hello danos\n"
        off = self.cluster_offset(self.hello_cluster)
        self.image[off : off + len(content)] = content

        # Root directory: bitmap, up-case, volume label, HELLO set.
        root = self.cluster_offset(self.root_cluster)
        # 0x81 Allocation Bitmap
        struct.pack_into("<BBB", self.image, root, 0x81, 0, 0)
        struct.pack_into("<I", self.image, root + 20, self.bitmap_cluster)
        struct.pack_into("<Q", self.image, root + 24, self.bitmap_bytes)
        # 0x82 Up-case Table
        struct.pack_into("<B", self.image, root + 32, 0x82)
        struct.pack_into("<I", self.image, root + 32 + 4, table_checksum)
        struct.pack_into("<I", self.image, root + 32 + 20, self.upcase_cluster)
        struct.pack_into("<Q", self.image, root + 32 + 24, UPCASE_UNITS * 2)
        # 0x83 Volume Label
        label_units = [ord(c) for c in self.label[:11]]
        struct.pack_into("<BB", self.image, root + 64, 0x83, len(label_units))
        for i, u in enumerate(label_units):
            struct.pack_into("<H", self.image, root + 64 + 2 + i * 2, u)
        # HELLO.TXT set: File (0x85) + Stream (0xC0) + Name (0xC1)
        name = "HELLO.TXT"
        self.write_file_set(root + 96, name, first_cluster=self.hello_cluster, length=len(content))

    def write_file_set(self, offset, name, first_cluster, length):
        entries = bytearray(32 * 3)
        # File entry
        entries[0] = 0x85
        entries[1] = 2  # stream + one name entry
        struct.pack_into("<H", entries, 4, 0x20)  # attributes: archive
        # Stream entry
        entries[32 + 0] = 0xC0
        entries[32 + 1] = 0x01 | 0x02  # allocation possible + no FAT chain (contiguous)
        entries[32 + 3] = len(name)
        upname = [ascii_upper(ord(c)) for c in name]
        struct.pack_into("<H", entries, 32 + 4, name_hash(upname))
        struct.pack_into("<Q", entries, 32 + 8, length)   # valid data length
        struct.pack_into("<I", entries, 32 + 20, first_cluster)
        struct.pack_into("<Q", entries, 32 + 24, length)  # data length
        # File Name entry
        entries[64 + 0] = 0xC1
        for i, c in enumerate(name):
            struct.pack_into("<H", entries, 64 + 2 + i * 2, ord(c))
        struct.pack_into("<H", entries, 2, set_checksum(entries))
        self.image[offset : offset + len(entries)] = entries

    def serialize(self):
        self.build()
        return bytes(self.image)


def verify(path):
    with open(path, "rb") as handle:
        data = handle.read()
    if len(data) < 512 or data[510] != 0x55 or data[511] != 0xAA:
        sys.exit("verify: missing 0x55AA boot signature")
    if data[3:11] != b"EXFAT   ":
        sys.exit("verify: not an exFAT boot sector")
    if any(data[11:64]):
        sys.exit("verify: MustBeZero region is not zero")
    fat_offset = struct.unpack_from("<I", data, 80)[0]
    heap_offset = struct.unpack_from("<I", data, 88)[0]
    cluster_count = struct.unpack_from("<I", data, 92)[0]
    root_cluster = struct.unpack_from("<I", data, 96)[0]
    spc = 1 << data[109]
    # Boot checksum sector 11 must match a fresh checksum over sectors 0..10.
    expected = boot_checksum(data[0 : 11 * SECTOR])
    got = struct.unpack_from("<I", data, 11 * SECTOR)[0]
    if expected != got:
        sys.exit(f"verify: boot checksum mismatch (0x{got:08X} != 0x{expected:08X})")
    # Walk the root directory for the HELLO.TXT set and check its checksum.
    root = (heap_offset + (root_cluster - 2) * spc) * SECTOR
    found = False
    for i in range(spc * SECTOR // 32):
        entry = root + i * 32
        if data[entry] == 0x00:
            break
        if data[entry] == 0x85:
            secondary = data[entry + 1]
            total = (secondary + 1) * 32
            stored = struct.unpack_from("<H", data, entry + 2)[0]
            if set_checksum(data[entry : entry + total]) != stored:
                sys.exit("verify: a file set checksum is wrong")
            found = True
    if not found:
        sys.exit("verify: no file set in the root directory")
    print(f"make-exfat-image: {path} OK "
          f"({cluster_count} clusters of {spc * SECTOR} bytes, fat@{fat_offset}, heap@{heap_offset})")


def main(argv):
    if len(argv) == 3 and argv[1] == "--verify":
        verify(argv[2])
        return 0
    argv = list(argv)
    volume_id = 0x1234ABCD
    label = "DANOS"
    i = 1
    while i < len(argv):
        if argv[i] == "--serial" and i + 1 < len(argv):
            volume_id = int(argv[i + 1], 16)
            del argv[i : i + 2]
        elif argv[i] == "--label" and i + 1 < len(argv):
            label = argv[i + 1]
            del argv[i : i + 2]
        else:
            i += 1
    if len(argv) != 3:
        sys.exit("usage: make-exfat-image.py [--serial <hex>] [--label <name>] <out.img> <size-MiB>\n"
                 "       make-exfat-image.py --verify <out.img>")
    out_path = argv[1]
    size_mib = int(argv[2])
    image = ExfatImage(size_mib, volume_id, label)
    with open(out_path, "wb") as handle:
        handle.write(image.serialize())
    print(f"make-exfat-image: wrote {out_path} "
          f"({size_mib} MiB exFAT, {image.cluster_count} clusters, serial 0x{image.volume_id:08X})")
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv))