#!/usr/bin/env python3
"""
Generate encrypted shellcode blobs for the redstone-jit task.

blob1: executes (ret), flag1 is plaintext in the blob body —
       participant must examine decrypted memory to find it.
       flag1 is ALSO the XOR key for blob2, so it serves double duty.
blob2: prints "Enter code: " prompt, reads 8 bytes, computes a
       polynomial rolling hash (Horner, mod 2^16), compares to a
       pre-computed target.  Correct input is "enderman" (0x8592).
       Encrypted with flag1 as cycling key.

Outputs blobs.h with the encrypted arrays.
"""

# ---- Embedded data ----
flag1    = b"vsosh{3ncrypt3d_by_st3v3}\x00"    # 26 bytes (25 chars + null)
flag2    = b"vsosh{d3crypt3d_by_y0u}\n\x00"    # 25 bytes (23 chars + \n + null)
prompt   = b"Enter code: "                       # 12 bytes (no null needed)
flag1_key = flag1[:-1]                           # 25 bytes cycling XOR key

assert len(flag1)    == 26, f"flag1 len={len(flag1)}"
assert len(flag2)    == 25, f"flag2 len={len(flag2)}"
assert len(prompt)   == 12, f"prompt len={len(prompt)}"
assert len(flag1_key)== 25, f"flag1_key len={len(flag1_key)}"

# ---- Assemble blob1 (27 bytes) ----
#
# Memory layout:
#   [0]     C3              ret  (blob executes and immediately returns)
#   [1-26]  flag1           26 bytes  <- visible in memory after decryption
#
# flag1 is the only secret here.  The participant finds it in blob1's body,
# enters it as the key, and that same string decrypts blob2.

blob1 = bytearray()
blob1 += b"\xC3"   # [0]    ret
blob1 += flag1     # [1-26]

assert len(blob1) == 27, f"blob1 len={len(blob1)}"

# ---- Compute hash target for "enderman" ----
#
# Polynomial rolling hash (Horner's method, mod 2^16):
#   acc = 0
#   for b in input[0..7]: acc = (acc * 31 + b) & 0xFFFF

SECRET_INPUT = b"enderman"   # 8 bytes

def poly_hash(data):
    acc = 0
    for b in data:
        acc = (acc * 31 + b) & 0xFFFF
    return acc

HASH_TARGET = poly_hash(SECRET_INPUT)
assert HASH_TARGET == 0x8592, f"unexpected hash: 0x{HASH_TARGET:04x}"

# flag2 is stored encrypted inside blob2 body — XOR'd with the low byte of
# the correct hash (0x92).  The shellcode decrypts it in-place at runtime
# using `al` after the hash loop.  Two consequences:
#   1. flag2 is never visible as plaintext in memory, even after blob2 decrypts
#   2. skipping the `jne` in a debugger writes garbage — al was wrong
DECRYPT_KEY = HASH_TARGET & 0xFF                        # 0x92
flag2_enc   = bytes([b ^ DECRYPT_KEY for b in flag2])   # 25 bytes, all encrypted

# ---- Assemble blob2 (161 bytes) ----
#
# Memory layout:
#   [0-1]    EB 25           short jmp to offset 39
#   [2-13]   prompt          "Enter code: " (12 bytes)
#   [14-38]  flag2_enc       flag2 XOR 0x92 (25 bytes) — NOT plaintext
#   [39-160] code            122 bytes
#
# Code logic (syscall-based, position-independent):
#   write(1, prompt, 12)
#   read(0, rsp, 8)
#   acc=0; for i in 0..7: acc=(acc*31+rsp[i])&0xFFFF   <- Horner hash
#   for i in 0..23: flag2_enc[i] ^= al                 <- in-place decrypt
#   if acc == HASH_TARGET: write(1, flag2_dec, 24)
#   ret
#
# Encrypted with flag1 (25 bytes cycling key).

# jmp offset: target=39, next_ip=2 => 39-2 = 37 = 0x25
blob2 = bytearray()
blob2 += b"\xEB\x25"   # [0-1]   jmp +37 -> offset 39
blob2 += prompt        # [2-13]
blob2 += flag2_enc     # [14-38]  encrypted, not readable

code2 = bytearray()
# All offsets relative to blob2 start.  Code starts at offset 39.

# [39]  sub rsp, 32
code2 += b"\x48\x83\xEC\x20"
# [43]  lea rsi, [rip + (2-50)]   prompt@2; rip_after=50; off=2-50=-48=0xFFFFFFD0
code2 += b"\x48\x8D\x35\xD0\xFF\xFF\xFF"
# [50]  mov edi, 1
code2 += b"\xBF\x01\x00\x00\x00"
# [55]  mov edx, 12
code2 += b"\xBA\x0C\x00\x00\x00"
# [60]  mov eax, 1                (write syscall)
code2 += b"\xB8\x01\x00\x00\x00"
# [65]  syscall                   write(1, prompt, 12)
code2 += b"\x0F\x05"

# [67]  mov rsi, rsp              (read buffer = rsp)
code2 += b"\x48\x89\xE6"
# [70]  xor edi, edi              (fd = 0)
code2 += b"\x31\xFF"
# [72]  mov edx, 8
code2 += b"\xBA\x08\x00\x00\x00"
# [77]  xor eax, eax              (read syscall)
code2 += b"\x31\xC0"
# [79]  syscall                   read(0, rsp, 8)
code2 += b"\x0F\x05"

# [81]  xor eax, eax              acc = 0
code2 += b"\x31\xC0"
# [83]  xor ecx, ecx              i = 0
code2 += b"\x31\xC9"
# hash_loop @ offset 85:
# [85]  imul eax, eax, 31         (6B C0 1F)
code2 += b"\x6B\xC0\x1F"
# [88]  movzx edx, byte [rsp+rcx] (0F B6 14 0C)
code2 += b"\x0F\xB6\x14\x0C"
# [92]  add eax, edx
code2 += b"\x01\xD0"
# [94]  and eax, 0xFFFF
code2 += b"\x25\xFF\xFF\x00\x00"
# [99]  inc ecx
code2 += b"\xFF\xC1"
# [101] cmp ecx, 8
code2 += b"\x83\xF9\x08"
# [104] jl hash_loop              next_ip=106, target=85: 85-106=-21=0xEB
code2 += b"\x7C\xEB"

# After hash loop: eax = computed hash, al = low byte
#
# Decrypt flag2_enc in-place using al: flag2_enc[i] ^= al  for i in 0..23
# If hash is correct (al=0x92), flag2 becomes plaintext.
# If hash is wrong, flag2 stays garbage — RIP skip to write won't help.
#
# [106] lea rdi, [rip + (14-113)] flag2_enc@14; rip_after=113; off=14-113=-99=0xFFFFFF9D
code2 += b"\x48\x8D\x3D\x9D\xFF\xFF\xFF"
# [113] xor ecx, ecx              i = 0
code2 += b"\x31\xC9"
# decrypt_loop @ offset 115:
# [115] xor byte [rdi+rcx], al    (30 04 0F)
code2 += b"\x30\x04\x0F"
# [118] inc ecx                   (FF C1)
code2 += b"\xFF\xC1"
# [120] cmp ecx, 24               (83 F9 18)  — 24 meaningful bytes (without null)
code2 += b"\x83\xF9\x18"
# [123] jl decrypt_loop           next_ip=125, target=115: 115-125=-10=0xF6
code2 += b"\x7C\xF6"

# [125] cmp eax, HASH_TARGET
code2 += b"\x3D" + HASH_TARGET.to_bytes(4, "little")
# [130] jne skip                  next_ip=132, skip=156: 156-132=24=0x18
code2 += b"\x75\x18"

# [132] lea rsi, [rip + (14-139)] flag2 now decrypted; rip_after=139; off=14-139=-125=0xFFFFFF83
code2 += b"\x48\x8D\x35\x83\xFF\xFF\xFF"
# [139] mov edi, 1
code2 += b"\xBF\x01\x00\x00\x00"
# [144] mov edx, 24               (23 chars + newline, no null)
code2 += b"\xBA\x18\x00\x00\x00"
# [149] mov eax, 1
code2 += b"\xB8\x01\x00\x00\x00"
# [154] syscall                   write(1, flag2_dec, 24)
code2 += b"\x0F\x05"

# skip @ offset 156:
# [156] add rsp, 32
code2 += b"\x48\x83\xC4\x20"
# [160] ret
code2 += b"\xC3"

assert len(code2) == 122, f"code2 len={len(code2)}"

blob2 += code2
assert len(blob2) == 161, f"blob2 len={len(blob2)}"

# ---- Encrypt blobs ----

KEY1 = 0x42
blob1_enc = bytes([b ^ KEY1 for b in blob1])

# blob2 is encrypted with flag1 (25 bytes cycling) — flag1 is both the flag
# the participant finds in blob1 AND the key they must enter to decrypt blob2.
blob2_enc = bytes([b ^ flag1_key[i % len(flag1_key)] for i, b in enumerate(blob2)])

# ---- Write blobs.h ----

def c_array(name, data):
    chunks = [data[i:i+12] for i in range(0, len(data), 12)]
    lines  = ["    " + ", ".join(f"0x{b:02x}" for b in chunk) for chunk in chunks]
    return f"static const uint8_t {name}[] = {{\n" + ",\n".join(lines) + "\n};\n"

header  = "/* Auto-generated by generate_blobs.py — do not edit */\n"
header += "#pragma once\n#include <stdint.h>\n\n"
header += c_array("blob1_encrypted", blob1_enc)
header += f"\n#define BLOB1_KEY  0x{KEY1:02x}u\n"
header += f"#define BLOB1_SIZE {len(blob1)}u\n\n"
header += c_array("blob2_encrypted", blob2_enc)
header += f"\n#define BLOB2_SIZE {len(blob2)}u\n"

with open("blobs.h", "w") as f:
    f.write(header)

print(f"[+] blobs.h generated")
print(f"    blob1 : {len(blob1)} bytes  (key=0x{KEY1:02x})")
print(f"    blob2 : {len(blob2)} bytes  (key=flag1)")
print(f"    flag1 : {flag1[:-1].decode()}")
print(f"    flag2 : {flag2[:-2].decode()}")
