#!/usr/bin/env python3
"""adfTool.py - read files out of an Amiga OFS/FFS disk image (.adf).

    adfTool.py list IMAGE                 every file, one path per line
    adfTool.py extract IMAGE PATH OUT     one file (PATH like Devs/serial.device)
    adfTool.py extractall IMAGE OUTDIR    every file, keeping the drawer tree

Read-only. Used by the installer's `import` (a Workbench disk image is how
most people own AmigaOS files such as Devs/serial.device) and by the Amiga
gates. Handles DD (880 KB) and HD (1.76 MB) images; names match case-
insensitively, as AmigaDOS does.
"""

import os
import struct
import sys

BLOCK_SIZE = 512
HASH_SLOTS = 72            # hash-table entries in a 512-byte header block
TYPE_HEADER = 2            # primary type of a root / drawer / file header
SEC_TYPE_ROOT = 1
SEC_TYPE_USERDIR = 2
SEC_TYPE_FILE = -3
OFS_DATA_HEADER = 24       # OFS data blocks start with a 24-byte header


class AdfImage:
    def __init__(self, path):
        with open(path, "rb") as f:
            self.image = f.read()
        if len(self.image) % BLOCK_SIZE != 0 or self.image[0:3] != b"DOS":
            raise ValueError("%s is not an AmigaDOS disk image" % path)
        self.isFfs = (self.image[3] & 1) != 0
        self.rootBlock = (len(self.image) // BLOCK_SIZE) // 2

    def block(self, number):
        return self.image[number * BLOCK_SIZE:(number + 1) * BLOCK_SIZE]

    @staticmethod
    def long(data, offset):
        # Negative offsets count from the end of the block, as the
        # AmigaDOS structure definitions do.
        if offset < 0:
            offset += BLOCK_SIZE
        return struct.unpack(">l", data[offset:offset + 4])[0]

    @staticmethod
    def name(data):
        length = data[BLOCK_SIZE - 80]
        return data[BLOCK_SIZE - 79:BLOCK_SIZE - 79 + length].decode("latin-1")

    def entries(self, dirBlock):
        # (name, header block, secondary type) for everything in a drawer.
        data = self.block(dirBlock)
        found = []
        for slot in range(HASH_SLOTS):
            number = self.long(data, 24 + 4 * slot)
            while number:
                header = self.block(number)
                if self.long(header, 0) == TYPE_HEADER:
                    found.append((self.name(header), number, self.long(header, -4)))
                number = self.long(header, -16)
        return sorted(found, key=lambda e: e[0].upper())

    def walk(self, dirBlock=None, prefix=""):
        # Yield (path, header block) for every file under a drawer.
        if dirBlock is None:
            dirBlock = self.rootBlock
        for entryName, number, secType in self.entries(dirBlock):
            path = prefix + entryName
            if secType == SEC_TYPE_USERDIR:
                yield from self.walk(number, path + "/")
            elif secType == SEC_TYPE_FILE:
                yield path, number

    def find(self, path):
        current = self.rootBlock
        for part in path.split("/"):
            match = [n for entryName, n, _ in self.entries(current) if entryName.upper() == part.upper()]
            if not match:
                raise KeyError("not found in image: %s" % path)
            current = match[0]
        return current

    def read(self, headerBlock):
        header = self.block(headerBlock)
        size = self.long(header, -188)
        data = bytearray()
        chunk = header
        while True:
            count = self.long(chunk, 8)
            # Data-block pointers fill the table from the END (slot 71 first).
            pointers = [self.long(chunk, 24 + 4 * i) for i in range(HASH_SLOTS)][::-1][:count]
            for pointer in pointers:
                blockData = self.block(pointer)
                if self.isFfs:
                    data += blockData
                else:
                    data += blockData[OFS_DATA_HEADER:OFS_DATA_HEADER + self.long(blockData, 12)]
            extension = self.long(chunk, -8)
            if not extension:
                break
            chunk = self.block(extension)
        return bytes(data[:size])


def main(argv):
    if len(argv) < 3 or argv[1] not in ("list", "extract", "extractall"):
        sys.stderr.write(__doc__)
        return 2
    try:
        image = AdfImage(argv[2])
        if argv[1] == "list":
            for path, _ in image.walk():
                print(path)
        elif argv[1] == "extract":
            if len(argv) != 5:
                sys.stderr.write(__doc__)
                return 2
            with open(argv[4], "wb") as f:
                f.write(image.read(image.find(argv[3])))
        else:
            if len(argv) != 4:
                sys.stderr.write(__doc__)
                return 2
            for path, number in image.walk():
                out = os.path.join(argv[3], *path.split("/"))
                os.makedirs(os.path.dirname(out), exist_ok=True)
                with open(out, "wb") as f:
                    f.write(image.read(number))
    except (OSError, ValueError, KeyError) as error:
        sys.stderr.write("adfTool.py: %s\n" % error)
        return 1
    return 0


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