import struct
from hackingscripts import util
from PIL import Image
from pyzbar.pyzbar import decode

def bit_stream_to_qr(bit_stream, qr_size=(29, 32), pix_size=10):
    img = Image.new("RGB", (qr_size[1] * pix_size, qr_size[0] * pix_size))
    pix = img.load()

    columns = ["" for r in range(qr_size[1])]
    for i, b in enumerate(bit_stream):
        columns[i % qr_size[1]] += b

    for ci, column in enumerate(columns):
        for ri, b in enumerate(column):
            color = (0,0,0) if b == "1" else (255,255,255)
            for xi in range(pix_size):
                for yi in range(pix_size):
                    pix[ci*pix_size+xi,ri*pix_size+yi] = color

    return img

def apply_scramble(data, loop_count, shift_direction, shift_type_1, shift_type_2):
    for i in range(loop_count):
        if shift_direction == 0:
            if shift_type_1 != -1:
                offset = 0x70 + shift_type_1
                value = data[offset]
                for x in range(0x1c+1):
                    data[offset] = data[offset - 4]
                    offset -= 4
                data[offset + 4] = value
            if shift_type_2 != -1:
                offset = shift_type_2 * 4
                value = data[offset]
                for x in range(3+1):
                    data[offset] = data[offset + 1]
                    offset += 1
                data[offset - 1] = value 
        elif shift_direction == 1:
            if shift_type_2 != -1:
                offset = 3 + shift_type_2 * 4
                value = data[offset]
                for x in range(3+1):
                    data[offset] = data[offset - 1]
                    offset -= 1
                data[offset + 1] = value
            if shift_type_1 != -1:
                offset = shift_type_1
                value = data[offset]
                for x in range(0x1c + 1):
                    data[offset] = data[offset + 4]
                    offset += 4
                data[offset - 4] = value
        else:
            raise
    return data

def bytes_to_bit_stream(data, qr_size=(29, 32)):
    bit_stream = ""
    for v in data:
        bit_stream += util.lpad(bin(v)[2:], 8, "0")
    bit_stream = bit_stream[0:qr_size[0] * qr_size[1]]
    return bit_stream

if __name__ == "__main__":

    with open("HV23", "rb") as f:
        prog_data = f.read()

    QR_DATA = bytearray(prog_data[0x8a3e:0x8a3e+0x70+7])
    SCRAMBLE_DATA = prog_data[0x628:0x628+0x60]
    NUM_SCRAMBLES = 0x18
    assert len(SCRAMBLE_DATA) == NUM_SCRAMBLES * 4

    scrambles = []
    for scramble in range(0, NUM_SCRAMBLES*4, 4):
        scramble_data = SCRAMBLE_DATA[scramble:scramble+4]
        shift_type_2, shift_type_1, loop_count, shift_direction = struct.unpack("bbBB", scramble_data)
        scrambles.append((loop_count, shift_direction, shift_type_1, shift_type_2))

    data = QR_DATA
    while True:
        for scramble in scrambles:
            data = apply_scramble(data, *scramble)
            bit_stream = bytes_to_bit_stream(data)

            # last 3 columns are all white? we might have a valid QR-code
            if all(bit_stream[row*32+29:row*32+32] == "000" for row in range(29)):
                img = bit_stream_to_qr(bit_stream)
                decoded_objects = decode(img)
                if decoded_objects:
                    img.save(f"qr_code.png")
                    print("[+] Flag:", decoded_objects[0].data.decode())
                    exit()