# SPDX-License-Identifier: GPL-2.0-or-later
"""The CDJ-900NXS update file (C900NXS.UPD, firmware 1.31): read, rebuild with
a new MAIN application, and check that nothing else changed.

The file: three decimal section lengths, one per line, then the sections
DRIV, MAIN and PANL. Each section is a 32-byte header ("CDJ-900NXS  MAINVer1.31"
...; its last byte selects the updater's mode, "0" = application only), then
Motorola S-records, then a CRC-16/XMODEM of all that, little-endian.

MAIN's S-records fill the player's NOR flash from 0: below 0x40000 the boot
loader and the recovery loader, from 0x40000 the application as one packed
region ([u32 length][LZSS stream][u16 additive checksum of length and stream]).
The boot loader checks that sum and unpacks the region to 0x08000000. The
updater in application-only mode writes 0x40000-0x7E0000 only.

This module never changes the boot loader or recovery records, DRIV or PANL:
build_main() copies them byte for byte and verify() checks it again on the
finished file.
"""
import binascii
import hashlib
import re
import struct

STOCK_UPD_SHA256 = 'e96bca7d1ac983985992a8e5fa7813e1ebf04bef514e0c6cfd0a8375f667f7f9'
STOCK_APP_SHA256 = '1ac3e30cde292865413950fe159663b3958e85939ed0a5d0122f796809a5f8ef'
STOCK_APP_SIZE = 16519056
APP_START = 0x40000             # the application's packed region in flash
APP_END = 0x7E0000              # what the updater may write, end exclusive
RECORD_DATA = 32                # data bytes per new S2 record, as the stock file
ADDRESS_BYTES = {0: 2, 1: 2, 2: 3, 3: 4, 5: 2, 6: 3, 7: 4, 8: 3, 9: 2}
HEADER = re.compile(rb'CDJ-900NXS  (DRIV|MAIN|PANL)Ver(\d\.\d\d)')


class UpdError(Exception):
    pass


def sha256(data):
    return hashlib.sha256(data).hexdigest()


def crc(data):
    return binascii.crc_hqx(data, 0)


class Record:
    __slots__ = ('kind', 'address', 'data')

    def __init__(self, kind, address, data):
        self.kind, self.address, self.data = kind, address, bytes(data)

    def __eq__(self, other):
        return (self.kind, self.address, self.data) == (other.kind, other.address, other.data)

    def encode(self):
        width = ADDRESS_BYTES[self.kind]
        body = bytes([width + len(self.data) + 1]) + self.address.to_bytes(width, 'big') + self.data
        body += bytes([~sum(body) & 0xFF])
        return b'S%d' % self.kind + body.hex().upper().encode() + b'\r\n'


class Section:
    def __init__(self, blob):
        self.blob = blob
        if len(blob) < 36:
            raise UpdError('a section is truncated')
        if crc(blob[:-2]) != struct.unpack_from('<H', blob, len(blob) - 2)[0]:
            raise UpdError('a section fails its CRC')
        self.header = blob[:32]
        m = HEADER.match(self.header)
        if not m:
            raise UpdError('unexpected section header %r' % self.header)
        self.name, self.version = m.group(1).decode(), m.group(2).decode()
        text = blob[32:-2]
        if not text.endswith(b'\r\n'):
            raise UpdError('%s: the S-records do not end in a line break' % self.name)
        self.records = []
        for n, line in enumerate(text[:-2].split(b'\r\n'), 1):
            if not re.fullmatch(rb'S[0-9][0-9A-F]+', line):
                raise UpdError('%s: S-record %d is malformed' % (self.name, n))
            kind = int(line[1:2])
            if kind not in ADDRESS_BYTES:
                raise UpdError('%s: S-record %d has type %d' % (self.name, n, kind))
            raw = bytes.fromhex(line[2:].decode())
            width = ADDRESS_BYTES[kind]
            if len(raw) != raw[0] + 1 or len(raw) < width + 2 or sum(raw) & 0xFF != 0xFF:
                raise UpdError('%s: S-record %d fails its length or checksum' % (self.name, n))
            self.records.append(Record(kind, int.from_bytes(raw[1:1 + width], 'big'), raw[1 + width:-1]))
        if self.records[0].kind != 0 or self.records[-1].kind not in (7, 8, 9):
            raise UpdError('%s: no S-record header or end' % self.name)

    def data_records(self):
        return [r for r in self.records if r.kind in (1, 2, 3)]

    def flash(self, size=APP_END):
        """The bytes the records put into flash below `size` (0xFF elsewhere)."""
        image = bytearray(b'\xff' * size)
        for r in self.data_records():
            if r.address + len(r.data) > size:
                raise UpdError('%s: a record reaches past 0x%X' % (self.name, size))
            image[r.address:r.address + len(r.data)] = r.data
        return bytes(image)


def encode_section(header, records):
    body = header + b''.join(r.encode() for r in records)
    return body + struct.pack('<H', crc(body))


def parse(raw):
    """The three sections of an update file, every length and checksum checked."""
    pos, sizes = 0, []
    for _ in range(3):
        end = raw.index(b'\r\n', pos)
        if not raw[pos:end].isdigit():
            raise UpdError('the file does not start with three section lengths')
        sizes.append(int(raw[pos:end]))
        pos = end + 2
    if pos + sum(sizes) != len(raw):
        raise UpdError('the section lengths do not add up to the file')
    sections = []
    for size in sizes:
        sections.append(Section(raw[pos:pos + size]))
        pos += size
    if [s.name for s in sections] != ['DRIV', 'MAIN', 'PANL']:
        raise UpdError('the sections are not DRIV, MAIN, PANL')
    return sections


def join(sections_blobs):
    return b''.join(b'%d\r\n' % len(b) for b in sections_blobs) + b''.join(sections_blobs)


# --- the packed application -------------------------------------------------

def unpack_lzss(stream):
    """The boot loader's LZSS (0xA0000900): a 4096-byte ring of 4078 spaces and
    18 zeros, written from 4078; flag bits LSB first, 1 = literal; a match is
    2 bytes, 12-bit ring position and length 3-18."""
    ring = bytearray(b' ' * 4078 + b'\0' * 18)
    r, out, pos, n = 4078, bytearray(), 0, len(stream)
    while pos < n:
        flags = stream[pos]
        pos += 1
        for bit in range(8):
            if pos >= n:
                break
            if flags >> bit & 1:
                b = stream[pos]
                pos += 1
                out.append(b)
                ring[r] = b
                r = (r + 1) & 0xFFF
            else:
                if pos + 2 > n:
                    raise UpdError('the packed application ends inside a match')
                lo, hi = stream[pos], stream[pos + 1]
                pos += 2
                at, count = lo | (hi & 0xF0) << 4, (hi & 0x0F) + 3
                for k in range(count):
                    b = ring[(at + k) & 0xFFF]
                    out.append(b)
                    ring[r] = b
                    r = (r + 1) & 0xFFF
    return bytes(out)


def region(flash):
    """(application, packed region bytes) from the flash image."""
    size = struct.unpack_from('<I', flash, APP_START)[0]
    end = APP_START + 4 + size
    if end + 2 > APP_END:
        raise UpdError('the packed application runs past 0x%X' % APP_END)
    stored = struct.unpack_from('<H', flash, end)[0]
    if stored != sum(flash[APP_START:end]) & 0xFFFF:
        raise UpdError('the packed application fails its additive checksum')
    return unpack_lzss(flash[APP_START + 4:end]), flash[APP_START:end + 2]


def pack_region(stream):
    """[length][stream][additive checksum], padded with 0xFF to 32 bytes."""
    out = struct.pack('<I', len(stream)) + stream
    out += struct.pack('<H', sum(out) & 0xFFFF)
    return out + b'\xff' * (-len(out) % RECORD_DATA)


# --- the stock file -----------------------------------------------------------

class Stock:
    """The official update, checked to be exactly the 1.31 file."""

    def __init__(self, raw):
        if sha256(raw) != STOCK_UPD_SHA256:
            raise UpdError('this is not the original CDJ-900NXS 1.31 update (SHA-256 %s)' % sha256(raw))
        self.raw = raw
        self.sections = parse(raw)
        main = self.sections[1]
        if main.version != '1.31' or main.header[31:32] != b'0':
            raise UpdError('MAIN is not 1.31 in application-only mode')
        kinds = {r.kind for r in main.records}
        if not kinds <= {0, 2, 7}:
            raise UpdError('MAIN has unexpected record types %s' % sorted(kinds))
        self.flash = main.flash()
        self.application, _ = region(self.flash)
        if len(self.application) != STOCK_APP_SIZE or sha256(self.application) != STOCK_APP_SHA256:
            raise UpdError('the 1.31 application inside does not match')

    def boot_records(self):
        """Boot loader and recovery: every data record below the application."""
        recs = self.sections[1].data_records()
        low = [r for r in recs if r.address < APP_START]
        if any(r.address + len(r.data) > APP_START for r in low):
            raise UpdError('a stock record crosses into the application')
        return low


def relabel(header, version):
    old = HEADER.match(header).group(0)
    new = old[:-4] + version.encode()
    if not re.fullmatch(r'\d\.\d\d', version):
        raise UpdError('the version label must look like 1.32')
    out = header.replace(old, new, 1)
    if out[:len(old) - 4] != header[:len(old) - 4] or out[len(old):] != header[len(old):]:
        raise UpdError('relabelling would change more than the version')
    return out


def build(stock, application, packed_stream, label='1.32'):
    """The new update: stock DRIV and PANL, MAIN with the stock header (version
    label only changed), the stock S0, the stock boot loader and recovery
    records, the new application region, and the stock S7 end record."""
    main = stock.sections[1]
    reg = pack_region(packed_stream)
    if APP_START + len(reg) > APP_END:
        raise UpdError('the packed application needs %d bytes, the flash has %d'
                       % (len(reg), APP_END - APP_START))
    records = [main.records[0]] + stock.boot_records()
    records += [Record(2, APP_START + off, reg[off:off + RECORD_DATA]) for off in range(0, len(reg), RECORD_DATA)]
    records.append(main.records[-1])
    blob = encode_section(relabel(main.header, label) if label else main.header, records)
    return join([stock.sections[0].blob, blob, stock.sections[2].blob])


def verify(raw, stock, application, label='1.32'):
    """Every check on a finished file; returns the list of what was checked."""
    done = []
    sections = parse(raw)
    done.append('section lengths, CRC-16 of every section, every S-record checksum')
    if sections[0].blob != stock.sections[0].blob or sections[2].blob != stock.sections[2].blob:
        raise UpdError('DRIV or PANL differs from stock')
    done.append('DRIV and PANL byte for byte as stock')
    main, smain = sections[1], stock.sections[1]
    if main.header != (relabel(smain.header, label) if label else smain.header):
        raise UpdError('the MAIN header differs beyond its version label')
    if main.header[31:32] != b'0':
        raise UpdError('the MAIN header no longer selects application-only mode')
    done.append('MAIN header as stock but for the label %s; application-only mode kept' % (label or '1.31'))
    if main.records[0] != smain.records[0] or main.records[-1] != smain.records[-1]:
        raise UpdError('the S0 header or S7 entry record changed')
    done.append('S0 header record and S7 entry record as stock')
    low = [r for r in main.data_records() if r.address < APP_START]
    if low != stock.boot_records():
        raise UpdError('the boot loader or recovery records changed')
    flash = main.flash()
    if flash[:APP_START] != stock.flash[:APP_START]:
        raise UpdError('the boot loader or recovery bytes changed')
    done.append('boot loader and recovery (flash 0x0-0x3FFFF): every record and byte as stock')
    if any(r.address < APP_START < r.address + len(r.data) for r in main.data_records()):
        raise UpdError('a record crosses into the boot area')
    done.append('no record reaches past 0x%X or into the boot area' % APP_END)
    app, reg = region(flash)
    done.append('packed application: additive checksum')
    if app != application:
        raise UpdError('the update does not unpack to the patched application')
    done.append('unpacked with the boot loader\'s LZSS: exactly the patched application (SHA-256 %s)'
                % sha256(app)[:16])
    done.append('flash used by the application: %d of %d bytes' % (len(reg), APP_END - APP_START))
    return done