#!/usr/bin/env python3
"""Extract, decrypt, and verify The Last Bitbender's mixed-mode payload."""

from pathlib import Path
from struct import pack, unpack


MASK64 = (1 << 64) - 1
HERE = Path(__file__).resolve().parent
PE_DATA_OFFSET = 0x600
BLOB_SIZE = 0x2A2
SAMPLE_OFFSET = PE_DATA_OFFSET + BLOB_SIZE


def rol64(value: int, count: int) -> int:
    return ((value << count) | (value >> (64 - count))) & MASK64


def transform(block: bytes) -> bytes:
    """Reimplementation of all four mixed-bitness stages."""
    left, right = unpack("<QQ", block)
    left ^= 0xA6F1C0D93B5E2748
    left = (left + (left & 0xFFFFFFFF) * (right & 0xFFFFFFFF)) & MASK64
    right = rol64(right, 13)
    left ^= right
    right = rol64((right + left) & MASK64, 29)
    right = (right * 0xFF51AFD7ED558CCD) & MASK64
    left = rol64((left + right) & MASK64, 17)
    return pack("<QQ", left ^ right, (left + right) & MASK64)


def main() -> None:
    challenge = (HERE / "chall.exe").read_bytes()
    blob = bytearray(challenge[PE_DATA_OFFSET : PE_DATA_OFFSET + BLOB_SIZE])
    assert len(blob) == BLOB_SIZE

    extracted_blob = HERE / "bitbender.bin"
    if extracted_blob.exists():
        assert extracted_blob.read_bytes() == blob
    else:
        extracted_blob.write_bytes(blob)

    state = (0x46662DE2AE713EE0 * 0xD1B54A32D192ED03) & MASK64
    state = rol64(state, 17) ^ 0x6E7A5380F8318187

    for offset in range(0xD4, 0xD4 + 0xC2):
        state = (state * 0x9E6C63C6A3C4B1D1 + 0x2545F4914F6CDD1D) & MASK64
        blob[offset] ^= state >> 56

    stage2 = HERE / "stage2.bin"
    stage2.write_bytes(blob[0xD4 : 0xD4 + 0xC2])

    state32 = 0x5F3A19C7
    for offset in range(0x196, 0x196 + 0x9A):
        state32 = (state32 * 0x2C9277B5 + 0xAC564B05) & 0xFFFFFFFF
        blob[offset] ^= state32 >> 24

    stage3 = HERE / "stage3.bin"
    stage3.write_bytes(blob[0x196 : 0x196 + 0x9A])

    state = 0xB5297A4D2C1F60E9
    for offset in range(0x230, 0x230 + 0x52):
        state = (state * 0x2545F4914F6CDD1D + 0x9E6C63C6A3C4B1D1) & MASK64
        blob[offset] ^= state >> 56

    stage4 = HERE / "stage4.bin"
    stage4.write_bytes(blob[0x230 : 0x230 + 0x52])

    decrypted_blob = HERE / "bitbender.decrypted.bin"
    decrypted_blob.write_bytes(blob)

    sample = challenge[SAMPLE_OFFSET : SAMPLE_OFFSET + 0x10]
    expected = challenge[SAMPLE_OFFSET + 0x10 : SAMPLE_OFFSET + 0x20]
    assert sample.hex() == "9c41e07db2f5361a8ad30c47e961b5f2"
    assert expected.hex() == "023a3db6ab0ec7efd2babd484c91f80f"
    assert transform(sample) == expected

    print(extracted_blob)
    print(stage2)
    print(stage3)
    print(stage4)
    print(decrypted_blob)
    print("self-test transform reproduced")


if __name__ == "__main__":
    main()