214 lines
8.7 KiB
Python
214 lines
8.7 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
import struct
|
|
from typing import BinaryIO
|
|
|
|
|
|
class Fat16Error(ValueError):
|
|
pass
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class FatEntry:
|
|
directory_offset: int
|
|
first_cluster: int
|
|
size: int
|
|
|
|
|
|
class Fat16Image:
|
|
def __init__(self, path: Path, *, writable: bool = False):
|
|
self.path = Path(path)
|
|
self.handle: BinaryIO = self.path.open("r+b" if writable else "rb")
|
|
self.writable = writable
|
|
self._load_geometry()
|
|
|
|
def __enter__(self) -> "Fat16Image":
|
|
return self
|
|
|
|
def __exit__(self, *_args: object) -> None:
|
|
self.handle.close()
|
|
|
|
def _load_geometry(self) -> None:
|
|
self.handle.seek(0)
|
|
mbr = self.handle.read(512)
|
|
if len(mbr) != 512 or mbr[510:512] != b"\x55\xaa":
|
|
raise Fat16Error("image does not contain a valid MBR")
|
|
entry = mbr[446:462]
|
|
self.partition_lba = struct.unpack_from("<I", entry, 8)[0]
|
|
self.partition_sectors = struct.unpack_from("<I", entry, 12)[0]
|
|
self.partition_offset = self.partition_lba * 512
|
|
if entry[4] not in {0x04, 0x06, 0x0E, 0x0C} or not self.partition_lba:
|
|
raise Fat16Error("first image partition is not a supported FAT partition")
|
|
self.handle.seek(self.partition_offset)
|
|
bpb = self.handle.read(512)
|
|
self.bytes_per_sector = struct.unpack_from("<H", bpb, 11)[0]
|
|
self.sectors_per_cluster = bpb[13]
|
|
self.reserved_sectors = struct.unpack_from("<H", bpb, 14)[0]
|
|
self.fat_count = bpb[16]
|
|
self.root_entries = struct.unpack_from("<H", bpb, 17)[0]
|
|
self.fat_sectors = struct.unpack_from("<H", bpb, 22)[0]
|
|
if self.bytes_per_sector != 512 or not self.fat_sectors or not self.root_entries:
|
|
raise Fat16Error("first image partition is not FAT16")
|
|
self.cluster_bytes = self.bytes_per_sector * self.sectors_per_cluster
|
|
self.fat_offset = self.partition_offset + self.reserved_sectors * self.bytes_per_sector
|
|
self.root_offset = self.fat_offset + self.fat_count * self.fat_sectors * self.bytes_per_sector
|
|
self.root_bytes = self.root_entries * 32
|
|
self.data_offset = self.root_offset + self.root_bytes
|
|
data_sectors = self.partition_sectors - (
|
|
self.reserved_sectors + self.fat_count * self.fat_sectors + self.root_bytes // self.bytes_per_sector
|
|
)
|
|
self.cluster_count = data_sectors // self.sectors_per_cluster
|
|
if not 4085 <= self.cluster_count < 65525:
|
|
raise Fat16Error("FAT partition cluster count is not FAT16")
|
|
|
|
@staticmethod
|
|
def _name83(name: str) -> bytes:
|
|
value = name.upper()
|
|
if value.count(".") > 1:
|
|
raise Fat16Error("FAT filename must use 8.3 form")
|
|
base, dot, extension = value.partition(".")
|
|
allowed = "ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_$~!#%&-{}()@'`"
|
|
if not 1 <= len(base) <= 8 or len(extension) > 3 or any(c not in allowed for c in base + extension):
|
|
raise Fat16Error(f"unsupported FAT 8.3 filename: {name}")
|
|
return base.ljust(8).encode("ascii") + extension.ljust(3).encode("ascii")
|
|
|
|
def _read_fat(self) -> bytearray:
|
|
self.handle.seek(self.fat_offset)
|
|
return bytearray(self.handle.read(self.fat_sectors * self.bytes_per_sector))
|
|
|
|
def _fat_value(self, fat: bytes, cluster: int) -> int:
|
|
return struct.unpack_from("<H", fat, cluster * 2)[0]
|
|
|
|
def _set_fat_value(self, fat: bytearray, cluster: int, value: int) -> None:
|
|
struct.pack_into("<H", fat, cluster * 2, value)
|
|
|
|
def _write_fats(self, fat: bytes) -> None:
|
|
if not self.writable:
|
|
raise Fat16Error("image was opened read-only")
|
|
for index in range(self.fat_count):
|
|
self.handle.seek(self.fat_offset + index * self.fat_sectors * self.bytes_per_sector)
|
|
self.handle.write(fat)
|
|
|
|
def find(self, name: str) -> FatEntry | None:
|
|
expected = self._name83(name)
|
|
self.handle.seek(self.root_offset)
|
|
root = self.handle.read(self.root_bytes)
|
|
for index in range(self.root_entries):
|
|
offset = index * 32
|
|
raw = root[offset : offset + 32]
|
|
if raw[0] == 0x00:
|
|
break
|
|
if raw[0] == 0xE5 or raw[11] == 0x0F:
|
|
continue
|
|
if raw[:11] == expected:
|
|
return FatEntry(
|
|
self.root_offset + offset,
|
|
struct.unpack_from("<H", raw, 26)[0],
|
|
struct.unpack_from("<I", raw, 28)[0],
|
|
)
|
|
return None
|
|
|
|
def _chain(self, first_cluster: int, fat: bytes) -> list[int]:
|
|
if first_cluster < 2:
|
|
return []
|
|
result: list[int] = []
|
|
seen: set[int] = set()
|
|
cluster = first_cluster
|
|
while 2 <= cluster < 0xFFF8:
|
|
if cluster in seen or cluster > self.cluster_count + 1:
|
|
raise Fat16Error("FAT cluster chain is invalid")
|
|
seen.add(cluster)
|
|
result.append(cluster)
|
|
cluster = self._fat_value(fat, cluster)
|
|
return result
|
|
|
|
def _cluster_offset(self, cluster: int) -> int:
|
|
return self.data_offset + (cluster - 2) * self.cluster_bytes
|
|
|
|
def read_file(self, name: str) -> bytes:
|
|
entry = self.find(name)
|
|
if entry is None:
|
|
raise FileNotFoundError(name)
|
|
fat = self._read_fat()
|
|
remaining = entry.size
|
|
result = bytearray()
|
|
for cluster in self._chain(entry.first_cluster, fat):
|
|
self.handle.seek(self._cluster_offset(cluster))
|
|
chunk = self.handle.read(min(self.cluster_bytes, remaining))
|
|
result.extend(chunk)
|
|
remaining -= len(chunk)
|
|
if remaining == 0:
|
|
break
|
|
if remaining:
|
|
raise Fat16Error(f"FAT file is truncated: {name}")
|
|
return bytes(result)
|
|
|
|
def free_bytes(self) -> int:
|
|
fat = self._read_fat()
|
|
free_clusters = sum(
|
|
1
|
|
for cluster in range(2, self.cluster_count + 2)
|
|
if self._fat_value(fat, cluster) == 0
|
|
)
|
|
return free_clusters * self.cluster_bytes
|
|
|
|
def _free_chain(self, first_cluster: int, fat: bytearray) -> None:
|
|
for cluster in self._chain(first_cluster, fat):
|
|
self._set_fat_value(fat, cluster, 0)
|
|
|
|
def _directory_slot(self, name: str) -> int:
|
|
existing = self.find(name)
|
|
if existing is not None:
|
|
return existing.directory_offset
|
|
self.handle.seek(self.root_offset)
|
|
root = self.handle.read(self.root_bytes)
|
|
for index in range(self.root_entries):
|
|
if root[index * 32] in {0x00, 0xE5}:
|
|
return self.root_offset + index * 32
|
|
raise Fat16Error("FAT16 root directory is full")
|
|
|
|
def write_file(self, name: str, data: bytes, *, contiguous: bool = False) -> None:
|
|
if not self.writable:
|
|
raise Fat16Error("image was opened read-only")
|
|
content = bytes(data)
|
|
fat = self._read_fat()
|
|
existing = self.find(name)
|
|
if existing is not None:
|
|
self._free_chain(existing.first_cluster, fat)
|
|
required = max(1, (len(content) + self.cluster_bytes - 1) // self.cluster_bytes)
|
|
free = [cluster for cluster in range(2, self.cluster_count + 2) if self._fat_value(fat, cluster) == 0]
|
|
if contiguous:
|
|
clusters: list[int] = []
|
|
run: list[int] = []
|
|
previous = -2
|
|
for cluster in free:
|
|
run = run + [cluster] if cluster == previous + 1 else [cluster]
|
|
previous = cluster
|
|
if len(run) == required:
|
|
clusters = run
|
|
break
|
|
else:
|
|
clusters = free[:required]
|
|
if len(clusters) != required:
|
|
raise Fat16Error("FAT16 partition does not have enough free clusters")
|
|
for index, cluster in enumerate(clusters):
|
|
next_cluster = clusters[index + 1] if index + 1 < len(clusters) else 0xFFFF
|
|
self._set_fat_value(fat, cluster, next_cluster)
|
|
chunk = content[index * self.cluster_bytes : (index + 1) * self.cluster_bytes]
|
|
self.handle.seek(self._cluster_offset(cluster))
|
|
self.handle.write(chunk)
|
|
if len(chunk) < self.cluster_bytes:
|
|
self.handle.write(bytes(self.cluster_bytes - len(chunk)))
|
|
self._write_fats(fat)
|
|
directory_offset = self._directory_slot(name)
|
|
record = bytearray(32)
|
|
record[:11] = self._name83(name)
|
|
record[11] = 0x20
|
|
struct.pack_into("<H", record, 26, clusters[0])
|
|
struct.pack_into("<I", record, 28, len(content))
|
|
self.handle.seek(directory_offset)
|
|
self.handle.write(record)
|
|
self.handle.flush()
|