#!/usr/bin/env python3
"""Rebuild the Anniversary Collection M1R analysis ELF from extracted ISO files.

Python 3.9+; standard library only. Required inputs: DATA12.AIF, SLUS_208.33.

Usage:
    python build_m1r_elf.py --source /path/to/disc-files
    python build_m1r_elf.py --source . --output analysis/M1R.relocated.elf
    python build_m1r_elf.py --source . --verify-against existing/M1R.relocated.elf

The default module analysis base is 0x00400000. The result combines the MM1-6
module with the collection shell, resolves SNR1 relocations and adds symbols.
It is for static analysis, not a bootable replacement. Ghidra annotations,
local-variable names, inferred types, and external-call references are separate
from the ELF and are not reproduced here. Two unresolved shell SYS_FATAL calls
remain documented in the report, as in the original analysis ELF.

The source directory must contain files extracted from the ISO (or be the root
of a mounted ISO). This script does not parse ISO images themselves.
"""
import argparse
import collections
import csv
import hashlib
import io
import json
import os
from pathlib import Path
import re
import struct
import sys
import tempfile

DEFAULT_BASE = 0x00400000
REFERENCE_SHA256 = 'ba8edc31dbc71a76abf78e26a03275ed5ea32c9c3822b4685e90b7688d6e6a52'
REFERENCE_SOURCES = {
    'M1R.rel': '744e493250a4612a8e0c375bf9e3ac40bb43e8f1f27aa1e1a98250631b1cc4f1',
    'M1R.elf': '39d5af73929e3d2e137a306ce634ab2cd934a4809d8c38426a78ac3c62905d87',
    'SLUS_208.33': '396c625e926c8726b1c0439b75e0192eef84483f657b7a5e40292fb81cce86e5',
}


def require(condition, message):
    if not condition:
        raise ValueError(message)


def check_range(data, offset, size, description):
    require(0 <= offset <= len(data) and 0 <= size <= len(data) - offset,
            f'{description} extends outside its file')


def cstring(data, offset):
    require(0 <= offset < len(data), 'String offset outside file')
    return data[offset:data.index(0, offset)].decode('ascii')


def elf_sections(data):
    check_range(data, 0, 52, 'ELF header')
    require(data[:7] == b'\x7fELF\x01\x01\x01', 'Expected little-endian ELF32')
    require(struct.unpack_from('<H', data, 18)[0] == 8, 'Expected a MIPS ELF')
    start = struct.unpack_from('<I', data, 32)[0]
    stride, count, string_index = struct.unpack_from('<HHH', data, 46)
    require(stride == 40, 'Unexpected ELF section header size')
    require(0 < string_index < count, 'Invalid section-name table index')
    check_range(data, start, count * stride, 'ELF section table')
    sections = [list(struct.unpack_from('<10I', data, start + i * stride)) for i in range(count)]
    for section in sections:
        if section[1] != 8:
            check_range(data, section[4], section[5], 'ELF section')
    strings = sections[string_index]
    names = data[strings[4]:strings[4] + strings[5]]
    return sections, [cstring(names, s[0]) for s in sections], string_index


def section_at(sections, address):
    return next((i for i, s in enumerate(sections)
                 if s[2] & 2 and s[3] <= address < s[3] + s[5]), None)


def collect_shell_exports(data):
    sections, names, _ = elf_sections(data)
    table = sections[names.index('.sndata')]
    result = {}
    for offset in range(table[4], table[4] + table[5] - 11, 4):
        pointer, address, metadata = struct.unpack_from('<III', data, offset)
        if metadata >> 16 != 2:
            continue
        i = section_at(sections, pointer)
        target = section_at(sections, address)
        if i is None or target is None or sections[i][1] != 1:
            continue
        s = sections[i]
        name_offset = s[4] + pointer - s[3]
        try:
            name = cstring(data, name_offset)
        except (UnicodeDecodeError, ValueError):
            continue
        if not re.fullmatch(r'[A-Za-z_][A-Za-z_0-9]*', name):
            continue
        key = (name, metadata & 0xFFFF)
        entry = (address, names[target])
        require(key not in result or result[key] == entry, 'Conflicting shell exports')
        result[key] = entry
    return result


def patch_word(word, kind, value):
    value &= 0xFFFFFFFF
    require(kind in (2, 4, 5, 6), f'Unknown relocation type {kind}')
    if kind != 2:
        require(word & 0xFC00003F != 0x0D, 'Loader rejects relocation over BREAK')
    if kind == 2:
        return value
    if kind == 4:
        require(word >> 26 in (2, 3), 'R_MIPS_26 site is not J or JAL')
        require(value % 4 == 0, 'Unaligned branch target')
        return (word & 0xFC000000) | ((value >> 2) & 0x03FFFFFF)
    immediate = (((value + 0x8000) & 0xFFFFFFFF) >> 16) if kind == 5 else value & 0xFFFF
    return (word & 0xFFFF0000) | immediate


def make_elf(template, blocks, symbols, entry):
    """Write conventional ELF32 sections/segments and a standard symbol table."""
    header = bytearray(template[:52])
    struct.pack_into('<III', header, 24, entry, 52, 0)
    struct.pack_into('<HHHHHH', header, 40, 52, 32, len(blocks), 40, len(blocks) + 4, len(blocks) + 1)
    output = header + bytearray(32 * len(blocks))
    sections = [[0] * 10]
    shstrings = bytearray(b'\0')
    phdrs = []

    def append(payload, align=16):
        output.extend(b'\0' * (-len(output) % align))
        offset = len(output)
        output.extend(payload)
        return offset

    def name_offset(name):
        offset = len(shstrings)
        shstrings.extend(name.encode('ascii') + b'\0')
        return offset

    for block in blocks:
        no = name_offset(block['name'])
        offset = append(block['data']) if block['type'] != 8 else 0
        block['file_offset'] = offset
        sections.append([no, block['type'], block['flags'], block['address'],
                         offset, block['size'], 0, 0, 16, 0])
        pflags = 4 | (2 if block['flags'] & 1 else 0) | (1 if block['flags'] & 4 else 0)
        phdrs.append([1, offset, block['address'], block['address'],
                      len(block['data']), block['size'], pflags, 16])
    shname, symname, strname = (name_offset(n) for n in ('.shstrtab', '.symtab', '.strtab'))
    strings = bytearray(b'\0')
    entries = bytearray(16)
    for name, address, section, metadata in symbols:
        no = len(strings)
        strings.extend(name.encode('ascii') + b'\0')
        entries.extend(struct.pack('<IIIBBH', no, address, 0, 0x10, 0, section))
    sections.append([shname, 3, 0, 0, append(shstrings), len(shstrings), 0, 0, 1, 0])
    sections.append([symname, 2, 0, 0, append(entries), len(entries), len(blocks) + 3, 1, 4, 16])
    sections.append([strname, 3, 0, 0, append(strings), len(strings), 0, 0, 1, 0])
    shoff = append(b''.join(struct.pack('<10I', *s) for s in sections))
    struct.pack_into('<I', output, 32, shoff)
    for i, p in enumerate(phdrs):
        struct.pack_into('<8I', output, 52 + i * 32, *p)
    return bytes(output)


def relocate(rel, elf, shell, base):
    check_range(rel, 0, 20, 'SNR1 header')
    require(rel[:4] == b'SNR1', 'Not an SNR1 module')
    require(base % 0x10000 == 0 and 0 < base < 0x10000000, 'Use an aligned analysis base')
    ms, mn, _ = elf_sections(elf)
    ss, sn, _ = elf_sections(shell)
    exports = collect_shell_exports(shell)
    ro, rc, so, sc = struct.unpack_from('<4I', rel, 4)
    require(ro >= 20 and sc > 0, 'Invalid SNR1 table header')
    check_range(rel, ro, rc * 12, 'SNR1 relocation table')
    check_range(rel, so, sc * 12, 'SNR1 symbol table')
    require(ro + rc * 12 == so and so + sc * 12 == len(rel), 'Unexpected REL layout')
    symbols = []
    resolved = [base]
    for i in range(sc):
        no, addr, meta = struct.unpack_from('<3I', rel, so + 12 * i)
        name = cstring(rel, no)
        if not i:
            require(not name and not addr and not meta, 'Invalid null symbol')
            symbols.append((name, addr, meta))
            continue
        if meta >> 16 == 1:
            require(addr == 0, f'Import is already relocated: {name}')
            require((name, meta & 0xFFFF) in exports,
                    f'Unresolved module import: {name} (ID {meta & 0xFFFF})')
            addr = exports[name, meta & 0xFFFF][0]
        else:
            require(meta >> 16 == 2 and section_at(ms, addr) is not None, 'Bad defined symbol')
            addr += base
        resolved.append(addr)
        symbols.append((name, addr, meta))
    patched = bytearray(rel)
    touched = set()
    patches = []
    type_counts = collections.Counter()
    for i in range(rc):
        site, info, addend = struct.unpack_from('<3I', rel, ro + i * 12)
        kind, index = info & 0xFF, info >> 8
        si = section_at(ms, site)
        require(si is not None and ms[si][1] == 1, 'Relocation outside initialized ELF sections')
        require(site + 4 <= ms[si][3] + ms[si][5] and site + 4 <= ro,
                'Relocation crosses a section or payload boundary')
        require(site % 4 == 0 and site not in touched, 'Unaligned or repeated relocation')
        require(index < sc, 'Invalid symbol index')
        touched.add(site)
        target = (resolved[index] + addend) & 0xFFFFFFFF
        old = struct.unpack_from('<I', rel, site)[0]
        new = patch_word(old, kind, target)
        if kind == 4:
            require((base + site + 4) >> 28 == target >> 28, 'Branch exceeds MIPS jump region')
        struct.pack_into('<I', patched, site, new)
        patches.append((site, kind, target, old, new, index, addend))
        type_counts[kind] += 1

    blocks = []
    module_indices = {}
    shell_indices = {}
    for source, sections, names, shift, prefix, mapping in [
        (shell, ss, sn, 0, 'shell', shell_indices),
        (elf, ms, mn, base, 'm1', module_indices),
    ]:
        for si, (s, name) in enumerate(zip(sections, names)):
            if not s[2] & 2 or not s[5] or s[1] not in (1, 8):
                continue
            payload = source[s[4]:s[4] + s[5]] if s[1] == 1 else b''
            if prefix == 'm1' and s[1] == 1:
                require(payload == rel[s[3]:s[3] + s[5]], f'ELF/REL mismatch: {name}')
                payload = bytes(patched[s[3]:s[3] + s[5]])
            mapping[si] = len(blocks) + 1
            blocks.append(dict(name=f'.{prefix}{name}', type=s[1], flags=s[2] & 7,
                               address=s[3] + shift, size=s[5], data=payload))
    by_address = sorted(blocks, key=lambda b: b['address'])
    require(all(a['address'] + a['size'] <= b['address'] for a, b in zip(by_address, by_address[1:])),
            'Module and shell memory overlap')

    output_symbols = []
    used_names = set()
    for name, addr, meta in symbols[1:]:
        si = section_at(ms, addr - base) if meta >> 16 == 2 else section_at(ss, addr)
        block = module_indices[si] if meta >> 16 == 2 else shell_indices[si]
        output_symbols.append((name, addr, block, meta))
        used_names.add(name)
    for (name, identifier), (addr, section_name) in exports.items():
        if name in used_names:
            continue
        si = section_at(ss, addr)
        if si in shell_indices:
            output_symbols.append((name, addr, shell_indices[si], 0x20000 | identifier))
            used_names.add(name)

    # The shell has a small, already-addressed SNR1 root. Resolve its _prolog
    # call into this active module; preserve unsupported imports as undefined.
    shell_root = ss[sn.index('.sndata')]
    require(shell[shell_root[4]:shell_root[4] + 4] == b'SNR1', 'Invalid shell SNR1 root')
    sr, src, sy, syc = struct.unpack_from('<4I', shell, shell_root[4] + 4)

    def shell_offset(va, size=1):
        si = section_at(ss, va)
        require(si is not None and ss[si][1] == 1, 'Bad shell file address')
        require(size >= 0 and va + size <= ss[si][3] + ss[si][5],
                'Shell record crosses a section boundary')
        return ss[si][4] + va - ss[si][3]

    module_definitions = {name: addr for name, addr, meta in symbols if meta >> 16 == 2}
    shell_patches = []
    unresolved_shell = []
    shell_offset(sr, src * 12)
    shell_offset(sy, syc * 12)
    for i in range(src):
        site, info, addend = struct.unpack_from('<3I', shell, shell_offset(sr) + 12 * i)
        require(info >> 8 < syc, 'Invalid shell symbol index')
        no, addr, meta = struct.unpack_from('<3I', shell, shell_offset(sy) + (info >> 8) * 12)
        name = cstring(shell, shell_offset(no))
        if name not in module_definitions:
            unresolved_shell.append(dict(site=site, name=name, kind=info & 0xFF))
            if name not in used_names:
                output_symbols.append((name, 0, 0, meta))
                used_names.add(name)
            continue
        target = module_definitions[name] + addend
        shell_offset(site, 4)
        require(site % 4 == 0, 'Unaligned shell relocation')
        if info & 255 == 4:
            require((site + 4) >> 28 == target >> 28, 'Shell branch exceeds MIPS jump region')
        si = section_at(ss, site)
        block = blocks[shell_indices[si] - 1]
        data = bytearray(block['data'])
        offset = site - block['address']
        old = struct.unpack_from('<I', data, offset)[0]
        new = patch_word(old, info & 255, target)
        struct.pack_into('<I', data, offset, new)
        block['data'] = bytes(data)
        shell_patches.append(dict(site=site, target=target, name=name))

    composite = make_elf(elf, blocks, output_symbols, struct.unpack_from('<I', shell, 24)[0])
    actual = composite
    out_sections, out_names, _ = elf_sections(actual)

    def out_word(address):
        si = section_at(out_sections, address)
        require(si is not None and out_sections[si][1] == 1, 'Missing output section')
        s = out_sections[si]
        return struct.unpack_from('<I', actual, s[4] + address - s[3])[0]

    # Independently decode the resulting words back to their expected values.
    for site, kind, target, old, new, index, addend in patches:
        word = out_word(base + site)
        if kind == 2:
            require(word == target, 'R32 verification failed')
        elif kind == 4:
            decoded = ((base + site + 4) & 0xF0000000) | ((word & 0x03FFFFFF) << 2)
            require(decoded == target and word >> 26 == old >> 26, 'J/JAL verification failed')
        elif kind == 5:
            expected = ((target >> 16) + (1 if target & 0x8000 else 0)) & 0xFFFF
            require(word & 0xFFFF == expected and word >> 16 == old >> 16, 'HI16 verification failed')
        else:
            require(word & 0xFFFF == target & 0xFFFF and word >> 16 == old >> 16, 'LO16 verification failed')
    text = out_sections[out_names.index('.m1.text')]
    zero_before = sum(struct.unpack_from('<I', rel, a)[0] == 0x0C000000
                      for a in range(text[3] - base, text[3] - base + text[5], 4))
    zero_after = sum(struct.unpack_from('<I', actual, a)[0] == 0x0C000000
                     for a in range(text[4], text[4] + text[5], 4))
    require(zero_after == 0, 'Zero-target calls remain in module')
    for patch in shell_patches:
        word = out_word(patch['site'])
        require(((patch['site'] + 4) & 0xF0000000) | ((word & 0x03FFFFFF) << 2)
                == patch['target'], 'Shell call verification failed')
    symtab = out_sections[out_names.index('.symtab')]
    strtab = out_sections[symtab[6]]
    require(symtab[5] == (len(output_symbols) + 1) * 16, 'Wrong output symbol count')
    for index, (name, addr, block, _) in enumerate(output_symbols, 1):
        no, value, size, info, other, si = struct.unpack_from('<IIIBBH', actual, symtab[4] + index * 16)
        require((cstring(actual, strtab[4] + no), value, size, info, other, si)
                == (name, addr, 0, 0x10, 0, block), 'Output symbol verification failed')
    symbol_csv = io.StringIO(newline='')
    writer = csv.writer(symbol_csv)
    writer.writerow(['name', 'analysis_address', 'section'])
    for name, addr, si, meta in output_symbols:
        writer.writerow([name, f'0x{addr:08X}', out_names[si] if si else 'UND'])
    report = dict(module_analysis_base=f'0x{base:08X}', module_relocations_applied=rc,
        relocation_types=dict(type_counts), module_jal_zero_before=zero_before,
        module_jal_zero_after=zero_after, shell_relocations_applied=shell_patches,
        unresolved_shell_calls=unresolved_shell,
        symbols=len(output_symbols),
        sources={name: hashlib.sha256(b).hexdigest() for name, b in
                 [('M1R.rel', rel), ('M1R.elf', elf), ('SLUS_208.33', shell)]},
        output_sha256=hashlib.sha256(actual).hexdigest(),
        validation='Every module relocation decoded and verified; all output symbols checked; zero JAL 0 in module.',
        purpose='Static analysis composite. Analysis base is chosen, not captured from a running console.')
    return composite, symbol_csv.getvalue().encode('utf-8'), report


def name_crc(name):
    """The disc's FIO_MakeCRC: uppercase alphanumerics; other bytes become zero."""
    value = 0xFFFFFFFF
    for char in name.upper().encode('ascii'):
        char = char if 48 <= char <= 57 or 65 <= char <= 90 else 0
        value ^= char << 24
        for _ in range(8):
            value = ((value << 1) ^ (0x04C11DB7 if value & 0x80000000 else 0)) & 0xFFFFFFFF
    return value


def read_module_archive(path):
    """Extract only the two root entries needed from the hashed AIF directory."""
    length = path.stat().st_size
    wanted = {name_crc(name): name for name in ('M1R.elf', 'M1R.rel')}
    found = {}
    records = {}
    with path.open('rb') as stream:
        first = stream.read(8)
        require(len(first) == 8, 'Truncated AIF header')
        info_size, root = struct.unpack('<II', first)
        require(16 <= info_size <= length, 'Invalid AIF directory size')
        stream.seek(0)
        header = stream.read(info_size)
        pending = [(root, ())]
        visited = set()
        while pending:
            pos, parents = pending.pop()
            require(pos % 16 == 0 and pos not in visited, 'Invalid or cyclic AIF directory')
            check_range(header, pos, 16, 'AIF directory record')
            visited.add(pos)
            crc, dirs_at, files_at, dirs, files = struct.unpack_from('<IIIHH', header, pos)
            parents = parents + ((crc,) if crc else ())
            for index in range(dirs):
                pending.append((dirs_at + index * 16, parents))
            if files:
                require(files_at % 16 == 0, 'Unaligned AIF file table')
                check_range(header, files_at, files * 16, 'AIF file table')
            for index in range(files):
                record = files_at + index * 16
                crc, offset, size, flags = struct.unpack_from('<IIII', header, record)
                require(info_size <= offset <= length and size <= length - offset,
                        'AIF entry extends outside the archive')
                if parents or crc not in wanted:
                    continue
                name = wanted[crc]
                require(name not in found, f'Duplicate AIF entry for {name}')
                require(flags == 0, f'Unsupported AIF flags for {name}: {flags:#x}')
                stream.seek(offset)
                data = stream.read(size)
                require(len(data) == size, f'Truncated AIF payload: {name}')
                found[name] = data
                records[name] = dict(record_offset=record, offset=offset, size=size,
                                     name_crc=f'0x{crc:08X}')
    require(len(found) == 2, 'DATA12.AIF is missing root entries M1R.elf and/or M1R.rel')
    return found, records


def disc_file(source, name):
    # Accommodate lowercase dumps and ISO tools retaining the ;1 suffix.
    matches = [p for p in source.iterdir()
               if p.is_file() and p.name.split(';', 1)[0].upper() == name.upper()]
    require(len(matches) == 1,
            f'Expected exactly one {name} in {source}; found {len(matches)}')
    return matches[0].resolve()


def write_outputs(outputs, sources, force):
    """Preflight all destinations; then atomically write and read back each file."""
    for path, data in outputs.items():
        for source in sources:
            require(path != source and not (path.exists() and os.path.samefile(path, source)),
                    f'Output would overwrite an input: {path}')
        require(not path.exists() or path.is_file(), f'Output is not a regular file: {path}')
        require(force or not path.exists() or path.read_bytes() == data,
                f'Different output already exists: {path}; choose another path or use --force')
    for path, data in outputs.items():
        if not path.exists() or path.read_bytes() != data:
            path.parent.mkdir(parents=True, exist_ok=True)
            temporary = None
            try:
                with tempfile.NamedTemporaryFile(dir=path.parent, prefix=path.name + '.',
                                                 suffix='.tmp', delete=False) as stream:
                    temporary = Path(stream.name)
                    stream.write(data)
                os.replace(temporary, path)
            finally:
                if temporary is not None and temporary.exists():
                    temporary.unlink()
        require(path.read_bytes() == data, f'Output read-back failed: {path}')


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument('--source', type=Path, default=Path.cwd(),
                        help='Folder containing DATA12.AIF and SLUS_208.33 (default: current folder)')
    parser.add_argument('-o', '--output', type=Path, default=Path('M1R.relocated.elf'),
                        help='Output ELF file (default: ./M1R.relocated.elf)')
    parser.add_argument('--base', type=lambda s: int(s, 0), default=DEFAULT_BASE,
                        help='Module analysis base (default: 0x00400000)')
    parser.add_argument('--verify-against', type=Path,
                        help='Require byte-for-byte equality with an existing ELF before writing')
    parser.add_argument('--force', action='store_true',
                        help='Replace differing generated outputs; never replaces input files')
    args = parser.parse_args(argv)
    try:
        source = args.source.resolve()
        require(source.is_dir(), '--source must be a folder of extracted or mounted ISO files')
        archive = disc_file(source, 'DATA12.AIF')
        shell_path = disc_file(source, 'SLUS_208.33')
        module, records = read_module_archive(archive)
        composite, symbols, report = relocate(module['M1R.rel'], module['M1R.elf'],
                                               shell_path.read_bytes(), args.base)
        known_input = report['sources'] == REFERENCE_SOURCES
        if known_input and args.base == DEFAULT_BASE:
            require(report['output_sha256'] == REFERENCE_SHA256,
                    'Known-input build does not match the reference ELF hash')
        if args.verify_against:
            require(composite == args.verify_against.read_bytes(),
                    f'Rebuilt ELF differs from {args.verify_against}')
        report['archive_entries'] = records
        report['reference_match'] = report['output_sha256'] == REFERENCE_SHA256
        report['known_reference_inputs'] = known_input
        report['file_validation'] = 'Generated outputs are read back and compared byte for byte.'
        output = args.output.resolve()
        outputs = {
            output: composite,
            output.with_suffix('.symbols.csv'): symbols,
            output.with_suffix('.report.json'): (json.dumps(report, indent=2) + '\n').encode('utf-8'),
        }
        require(len(outputs) == 3, 'Output filename collides with a sidecar filename')
        protected = {archive, shell_path, Path(__file__).resolve()}
        if args.verify_against:
            protected.add(args.verify_against.resolve())
        # An identical ELF used as its own reference is safe and needs no write.
        if output in protected and args.verify_against and output == args.verify_against.resolve():
            require(output not in {archive, shell_path, Path(__file__).resolve()},
                    'Output would overwrite an input')
            protected.remove(output)
        write_outputs(outputs, protected, args.force)
    except (OSError, ValueError, KeyError, IndexError, struct.error) as error:
        print(f'error: {error}', file=sys.stderr)
        return 1
    print(f'ELF: {output}')
    print(f'Module base: {args.base:#010x}; relocations: {report["module_relocations_applied"]:,}; '
          f'symbols: {report["symbols"]:,}; remaining module JAL 0: {report["module_jal_zero_after"]}')
    print(f'SHA-256: {report["output_sha256"]}')
    if report['reference_match']:
        print('Matches the original analysis ELF.')
    if args.verify_against:
        print(f'Byte-for-byte comparison passed: {args.verify_against}')
    print(f'Symbols: {output.with_suffix(".symbols.csv")}')
    print(f'Report: {output.with_suffix(".report.json")}')
    return 0


if __name__ == '__main__':
    sys.exit(main())