# Please ensure that you solved the challenge properly at the local.
# If things do not run smoothly, you generally won't be allowed to make another attempt.
from secret.network_util import check_client, ban_client

import sys
import os
import subprocess
import tempfile

OFFSET_PLAINTEXT = 0x4010
OFFSET_KEY = 0x4020
OFFSET_MAIN_START = 0x1169
OFFSET_MAIN_END = 0x11ed

def main():
    if not check_client():
        return

    key = os.urandom(16)
    with open("encrypt", "rb") as f:
        content = bytearray(f.read())

    # input format: hex(plaintext) i j
    try:
        plaintext_hex, i_str, j_str = input().split()
        pt = bytes.fromhex(plaintext_hex)
        assert len(pt) == 16
        i = int(i_str)
        assert 0 <= i < len(content)
        assert not OFFSET_MAIN_START <= i < OFFSET_MAIN_END
        j = int(j_str)
        assert 0 <= j < 8
    except Exception as err:
        print(err, file=sys.stderr)
        # ban_client()
        return

    # update key, plaintext, and inject the fault
    content[OFFSET_KEY:OFFSET_KEY + 16] = key
    content[OFFSET_PLAINTEXT:OFFSET_PLAINTEXT + 16] = pt
    content[i] ^= (1 << j)

    tmpfile = tempfile.NamedTemporaryFile(delete=True)
    with open(tmpfile.name, "wb") as f:
        f.write(content)
    os.chmod(tmpfile.name, 0o775)
    tmpfile.file.close()

    # execute the modified binary
    try:
        ciphertext = subprocess.check_output(tmpfile.name, timeout=1.0)
        print(ciphertext.hex())
    except Exception as err:
        print(err, file=sys.stderr)
        ban_client()
        return

    # please guess the AES key
    if bytes.fromhex(input()) == key:
        with open("secret/flag.txt") as f:
            print(f.read())
        from datetime import datetime
        print(datetime.now(), plaintext_hex, i, j, file=sys.stderr)


main()

import os
import subprocess
import tempfile
from contextlib import suppress
from multiprocessing import Pool

OFFSET_PLAINTEXT = 0x4010
OFFSET_KEY = 0x4020
OFFSET_MAIN_START = 0x1169
OFFSET_MAIN_END = 0x11ed

# chose any key/pt pair and calculate xor
key = "keyakeyakeyakeya".encode()
pt = "0123456789abcdef".encode()
kxp = kxp = bytes([k ^ p for k, p in zip(key, pt)])
# the standard_res is the output of running it once normally
standard_res = b'\xbaP\xaa\x1d?\xfa\xbfJ5\xd2\xd8&\xfd%"\xcb'


def brute_index(idx):
    content_copy = content.copy()
    i = idx // 8
    j = idx % 8
    content_copy[i] ^= (1 << j)
    tmp = tempfile.NamedTemporaryFile(prefix=str(idx), delete=True)
    with open(tmp.name, "wb") as f:
        f.write(content_copy)
    os.chmod(tmp.name, 0o775)
    tmp.file.close()

    # ignore exception with contextlib
    with suppress(Exception):
        out = subprocess.check_output(tmp.name, timeout=0.01, stderr=subprocess.DEVNULL)
        if out != standard_res:
            if key in out:
                postfix = "Bitflip in main" if OFFSET_MAIN_START <= i < OFFSET_MAIN_END else ""
                print(f"Found key @ {i}/{j} " + postfix)
            if kxp in out:
                print(f"Found key^plaintext @ {i}/{j}")


with open("encrypt", "rb") as f:
    content = bytearray(f.read())
print(f"{key=}, {pt=}, {kxp=}")

# update key, plaintext, and inject the fault
content[OFFSET_KEY:OFFSET_KEY + 16] = key
content[OFFSET_PLAINTEXT:OFFSET_PLAINTEXT + 16] = pt

with Pool() as p:
    list(p.map(brute_index, range(len(content)*8)))

from pwn import unhex, enhex, remote, process

# line we want to fault-inject: write(STDOUT_FILENO, plaintext, 16);
plain = "a" * 16  # irrelevant, just has to be printable
plain = enhex(plain.encode())  # will be 616161....
pos = 8871  # at 0x11b9 there is the relevant MOV EDX 0x10 / 0x ba 10 00 00 00 instruction that loads the 16 on stack
bit = 2  # flip the 0x10 to 0x30, so that we print 48 bytes instead of 16 chars
payload = f"{plain} {pos} {bit}"
print(payload)

proc = process(["python", "main.py"])
# proc = remote("139.162.24.230", 31340)
# proc = remote("139.162.24.230", 31339)
proc.sendline(payload.encode())
res = proc.recvline()  # the ciphertext + key + 16 bytes of garbage are returned
plain = unhex(res)  # first 16 bytes are the ciphertext
# xor with plaintext to get the key
key = bytes([plain[i] ^ ord("a") for i in range(16)])
proc.sendline(enhex(key))
print(proc.recvline())