import os

FLAG_OFFSET = REDACTED
FLAG = REDACTED

BURGER_LAYERS = ["Bun (1)", "Lettuce (2)", "Tomato (3)", "Cheese (4)",
                 "Patty (5)", "Onion (6)", "Pickle (7)", "Mayonnaise (8)",
                 "Egg (9)", "Avocado (10)"]

def print_towers(towers):
    for i, tower in enumerate(towers):
        print(f"Tower {i + 1}: {[BURGER_LAYERS[layer - 1] for layer in tower]}")
    print()

def move_disk(towers, from_tower, to_tower):
    if not towers[from_tower]:
        print("Invalid move: Source plate is empty.")
        return False
    if towers[to_tower] and towers[to_tower][-1] < towers[from_tower][-1]:
        print("Invalid move: Cannot place larger burger layer on smaller burger layer.")
        return False
    towers[to_tower].append(towers[from_tower].pop())
    return True

def towers_of_hanoi():
    num_disks = int(input("Enter the number of burger layers (1-10): "))
    while num_disks < 1 or num_disks > 10:
        print("Please enter a number between 1 and 10 inclusive.")
        num_disks = int(input("Enter the number of burger layers (1-10): "))
    towers = [list(range(num_disks, 0, -1)), [], []]

    print_towers(towers)

    while len(towers[2]) != num_disks:
        try:
            from_tower = int(input("Move from plate (1-3): ")) - 1
            to_tower = int(input("Move to plate (1-3): ")) - 1
            if from_tower not in range(3) or to_tower not in range(3):
                print("Invalid input. Please enter a number between 1 and 3.")
                continue
            if move_disk(towers, from_tower, to_tower):
                print_towers(towers)
        except ValueError:
            print("Invalid input. Please enter a valid number.")

    print(f"Congratulations! You solved the puzzle. The last {num_disks} bits of the offset used to encrypt the flag are {bin(FLAG_OFFSET & ((1 << (num_disks)) - 1))}.")
    if num_disks >= 9:
        print("Waoh, that is a big sandwich! The offset only has 6 bits though...")

from pwn import *
import re

def get_offset_bits(num_disks, conn):
    # Function to solve Towers of Hanoi and get offset bits
    conn.sendline(b"1")
    conn.sendline(str(num_disks).encode())

    def hanoi_moves(n, source=1, auxiliary=2, target=3):
        if n > 0:
            yield from hanoi_moves(n - 1, source, target, auxiliary)
            yield (source, target)
            yield from hanoi_moves(n - 1, auxiliary, source, target)

    for from_tower, to_tower in hanoi_moves(num_disks):
        conn.sendline(str(from_tower).encode())
        conn.sendline(str(to_tower).encode())

    while True:
        line = conn.recvline().decode()
        if "bits of the offset" in line:
            return re.search(r"0b[01]+", line).group(0)[2:]

def key_gen_func(i: int) -> int:
    # Gray code sequence
    return i ^ (i >> 1)

def decrypt(encrypted_hex: str, offset: int) -> bytes:
    encrypted = bytes.fromhex(encrypted_hex)
    decrypted = bytearray()

    for i, char in enumerate(encrypted):
        decrypted.append(char ^ (key_gen_func(i + offset) & 0xFF))

    return bytes(decrypted)

def solve():
    # Connect to local Python script
    conn = process(["python3", "partial_source.py"])

    # Get offset bits and decrypt flag
    bits = get_offset_bits(6, conn)
    offset = int(bits, 2)

    conn.recvuntil(b"Select an option: ")
    conn.sendline(b"3")
    enc_flag = conn.recvline().strip().decode()

    # Decrypt and print flag
    flag = decrypt(enc_flag, offset)
    print(f"Flag: {flag.decode()}")

if __name__ == "__main__":
    solve()