#!/usr/bin/env python3
# decryptDisk.py - undo the Electronic Arts sector cipher used by the Modem Wars 1541 loader
# (drive code at $0531): every sector except track 18 is XORed with a rolling key.
import os, sys, glob

SECTORS_PER_TRACK = [21]*17 + [19]*7 + [18]*6 + [17]*5

def decryptSector(track, sector, data):
    if track == 18:
        return bytes(data)
    # lda track ; ora #$C0 ; asl a  (carry = 1 since bit 7 is set) ; adc sector
    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          # tya ; eor key
        # rol a (carry from the previous rol / adc)
        newCarry = (a >> 7) & 1
        a = ((a << 1) | carry) & 0xFF
        carry = newCarry
        key = a                  # sta key
        out[y] = data[y] ^ key
    return bytes(out)

def main():
    src = sys.argv[1]
    dst = sys.argv[2]
    os.makedirs(dst, exist_ok=True)
    img = bytearray()
    for track in range(1, 36):
        n = SECTORS_PER_TRACK[track-1]
        trackBytes = bytearray()
        for s in range(n):
            raw = open(f"{src}/t{track:02d}s{s:02d}.bin", "rb").read()
            dec = decryptSector(track, s, raw)
            open(f"{dst}/t{track:02d}s{s:02d}.bin", "wb").write(dec)
            trackBytes += dec
        open(f"{dst}/track{track:02d}.bin", "wb").write(trackBytes)
        img += trackBytes
    open(f"{dst}/decrypted.d64", "wb").write(img)
    print(f"wrote {len(img)} bytes")

main()