#!/usr/bin/env python3
# injectDriver.py - write a 4096-byte $E000 opponent-module image into a copy of the Modem Wars disk.
#
#   python3 injectDriver.py <module.bin> <in.d64> <out.d64>
#
# The module occupies track 18 sector 7 ($E000-$E0FF, stored in the clear because the loader never
# encrypts track 18) followed by track 34 sectors 1-15 ($E100-$EFFF, stored with the EA sector cipher).
import sys, os

SECTORS_PER_TRACK = [21]*17 + [19]*7 + [18]*6 + [17]*5      # tracks 1..35


def trackOffset(track):
    return sum(SECTORS_PER_TRACK[:track-1]) * 256


def sectorOffset(track, sector):
    return trackOffset(track) + sector * 256


def encryptSector(track, sector, plain):
    """The cipher the drive code applies at $0531; track 18 is stored in the clear."""
    if track == 18:
        return bytes(plain)
    a = (track | 0xC0) << 1
    carry = (a >> 8) & 1
    a &= 0xFF
    a = a + sector + carry
    carry = (a >> 8) & 1
    a &= 0xFF
    out = bytearray(256)
    key = 0
    for y in range(256):
        if y != 0:
            a = y ^ key
        newCarry = (a >> 7) & 1
        a = ((a << 1) | carry) & 0xFF
        carry = newCarry
        key = a
        out[y] = plain[y] ^ key
    return bytes(out)


# where each page of the module lives on the disk
def placement():
    places = [(18, 7)]                                      # $E000-$E0FF
    places += [(34, s) for s in range(1, 16)]               # $E100-$EFFF
    return places


def main():
    modulePath, inPath, outPath = sys.argv[1], sys.argv[2], sys.argv[3]
    module = open(modulePath, "rb").read()
    if len(module) != 4096:
        raise SystemExit(f"module must be exactly 4096 bytes, got {len(module)}")
    image = bytearray(open(inPath, "rb").read())
    if len(image) < 174848:
        raise SystemExit("not a 35-track .d64")
    for i, (track, sector) in enumerate(placement()):
        plain = module[i*256:(i+1)*256]
        off = sectorOffset(track, sector)
        image[off:off+256] = encryptSector(track, sector, plain)
    open(outPath, "wb").write(image)
    print(f"wrote {outPath}: module installed in track 18 sector 7 and track 34 sectors 1-15")


if __name__ == "__main__":
    main()