#!/usr/bin/env python3
"""
embed_payload.py <input.elf> <output_blob.cpp>

Reads a RISC-V ELF payload, applies the opcode shuffle (driven by
JANUS_OPCODE_SEED from config.hpp), XOR-encrypts, and emits a C++ source file
with the encrypted blob and relocation table for the loader.

Requires: pyelftools  (pip install pyelftools)
"""

import sys
import struct
import re
from elftools.elf.elffile import ELFFile
from elftools.elf.relocation import RelocationSection

# Must stay in sync with include/janus/config.hpp 
OPCODE_SEED = 0xDEADBEEF
XOR_KEY     = 0x4A
XOR_MIX     = 0x67

def gen_perm(seed: int) -> list[int]:
    """Fisher-Yates with LCG — must match gen_perm() in src/vm/core.cpp."""
    p = list(range(32))
    for i in range(31, 0, -1):
        seed = (seed * 1664525 + 1013904223) & 0xFFFFFFFF
        j = (seed >> 17) % (i + 1)
        p[i], p[j] = p[j], p[i]
    return p

def xor_crypt(data: bytearray) -> bytearray:
    return bytearray(b ^ (XOR_KEY ^ ((i * XOR_MIX) & 0xFF)) for i, b in enumerate(data))

def shuffle_opcodes(data: bytearray, text_start: int, text_size: int, perm: list[int]) -> None:
    """Apply opcode permutation to every 32-bit RISC-V instruction in .text."""
    for off in range(text_start, text_start + text_size, 4):
        if off + 4 > len(data):
            break
        word = struct.unpack_from('<I', data, off)[0]
        if (word & 0x3) != 0x3:
            continue  # compressed or data — skip
        real_op  = (word >> 2) & 0x1F
        shuf_op  = perm[real_op]
        word = (word & ~(0x1F << 2)) | (shuf_op << 2)
        struct.pack_into('<I', data, off, word)

def main():
    if len(sys.argv) != 3:
        print(f"usage: {sys.argv[0]} <input.elf> <output_blob.cpp>", file=sys.stderr)
        sys.exit(1)

    elf_path, out_path = sys.argv[1], sys.argv[2]
    perm = gen_perm(OPCODE_SEED)

    with open(elf_path, 'rb') as f:
        elf = ELFFile(f)

        # Collect load segments into a flat image starting at virtual address 0
        load_segs = sorted(
            [s for s in elf.iter_segments() if s.header.p_type == 'PT_LOAD'],
            key=lambda s: s.header.p_vaddr
        )
        if not load_segs:
            print("No PT_LOAD segments", file=sys.stderr)
            sys.exit(1)

        img_end = max(s.header.p_vaddr + s.header.p_memsz for s in load_segs)
        image   = bytearray(img_end)

        for seg in load_segs:
            vaddr = seg.header.p_vaddr
            data  = seg.data()
            image[vaddr:vaddr + len(data)] = data

        # Find .text section bounds for opcode shuffle
        text_start = text_size = 0
        sec = elf.get_section_by_name('.text')
        if sec:
            text_start = sec.header.sh_addr
            text_size  = sec.header.sh_size

        # Collect R_RISCV_64 / R_RISCV_RELATIVE relocations
        reloc_offsets: list[int] = []
        for section in elf.iter_sections():
            if not isinstance(section, RelocationSection):
                continue
            sym_table = elf.get_section(section.header.sh_link)
            for reloc in section.iter_relocations():
                rtype = reloc.entry.r_info_type
                # R_RISCV_64 = 2, R_RISCV_RELATIVE = 3
                if rtype in (2, 3):
                    reloc_offsets.append(reloc.entry.r_offset)

        # Entry point
        entry_offset = elf.header.e_entry

    # Apply shuffle then encrypt
    shuffle_opcodes(image, text_start, text_size, perm)
    encrypted = xor_crypt(image)

    # Emit C++ blob
    hex_bytes  = ', '.join(f'0x{b:02X}' for b in encrypted)
    reloc_vals = ', '.join(str(o) for o in reloc_offsets) if reloc_offsets else '0'
    reloc_cnt  = len(reloc_offsets)

    cpp = f"""\
// AUTO-GENERATED by tools/embed_payload.py — do not edit by hand.
#include <cstdint>
#include <cstddef>

extern const uint8_t  kPayload[]      = {{ {hex_bytes} }};
extern const size_t   kPayloadSize    = {len(encrypted)};
extern const uint64_t kEntryOffset    = {entry_offset}ULL;
extern const uint64_t kRelocOffsets[] = {{ {reloc_vals} }};
extern const size_t   kRelocCount     = {reloc_cnt};
"""

    with open(out_path, 'w') as f:
        f.write(cpp)

    print(f"[janus] payload {len(encrypted)} bytes, {reloc_cnt} relocs → {out_path}")

if __name__ == '__main__':
    main()