import hashlib
import time

def solve_flag(target_blob: bytes):
    """
    Reverses the 8-stage hashing loop by walking backwards from the 7th shift
    down to the 0th shift, recovering 1 bit of the flag per iteration.
    """
    # 1. Setup Known Constraints
    known_prefix = b"dach2026{"
    known_suffix = b"}"
    total_len = 29
    
    # Calculate indices for the 19 unknown characters
    unknown_indices = list(range(len(known_prefix), total_len - len(known_suffix)))
    
    # Pre-calculate the bitwise NOT of the known parts
    full_known_state = bytearray(total_len)
    for i, b in enumerate(known_prefix):
        full_known_state[i] = ~b & 0xFF
    full_known_state[total_len - 1] = ~known_suffix[0] & 0xFF
    
    # Split the target blob into the 8 independent target hashes
    target_hashes = set(target_blob[i*32 : (i+1)*32] for i in range(8))
    
    # This array will accumulate our recovered bits for the 19 unknown bytes
    # By the end, it will contain the fully unshifted, inverted bytes.
    recovered_unknowns = [0] * len(unknown_indices)
    
    print("[*] Starting Bit-by-Bit Backtracking Attack...")
    start_time = time.time()
    
    # 2. Walk backwards from Shift 7 down to Shift 0
    for shift_amount in range(7, -1, -1):
        print(f"\n[*] Cracking Shift Layer {shift_amount} (Guessing bit {shift_amount} of payload)...")
        
        # Pre-build the known parts of the state for this specific shift layer
        base_state = bytearray(total_len)
        for i in range(total_len):
            if i < len(known_prefix) or i == total_len - 1:
                base_state[i] = full_known_state[i] >> shift_amount
                
        match_found = False
        
        # Brute-force the next bit for all 19 unknown bytes (2^19 combinations = 524,288)
        # This takes a few seconds per layer in pure Python.
        for guess in range(1 << len(unknown_indices)):
            test_state = bytearray(base_state)
            
            # Inject our guessed bits into the test state
            for i, idx in enumerate(unknown_indices):
                guess_bit = (guess >> i) & 1
                test_state[idx] = (recovered_unknowns[i] << 1) | guess_bit
                
            # Hash and check against the pool of valid target hashes
            if hashlib.sha256(test_state).digest() in target_hashes:
                print(f"    [+] Match found! Extracted bits: {bin(guess)[2:].zfill(19)}")
                
                # Commit the guessed bits to our recovered array
                for i in range(len(unknown_indices)):
                    recovered_unknowns[i] = (recovered_unknowns[i] << 1) | ((guess >> i) & 1)
                
                match_found = True
                break
                
        if not match_found:
            print(f"    [-] CRITICAL FAILURE: Could not find a valid bit permutation at layer {shift_amount}.")
            print("        Double check the dragonfly.bin data and prefix.")
            return None

    # 3. Final Assembly
    # We now have the complete, unshifted inverted state.
    final_inverted_state = bytearray(full_known_state)
    for i, idx in enumerate(unknown_indices):
        final_inverted_state[idx] = recovered_unknowns[i]
        
    # Un-invert it to reveal the plaintext flag
    final_flag = bytes(~b & 0xFF for b in final_inverted_state)
    
    elapsed = time.time() - start_time
    print(f"\n[+] Exploit completed in {elapsed:.2f} seconds.")
    print(f"[!] RECOVERED FLAG: {final_flag.decode('ascii', errors='ignore')}")
    
    return final_flag


# ==========================================
# Execution / Test Mock
# ==========================================
if __name__ == "__main__":
    with open("./romfs/dragonfly.bin", "rb") as f:
         target_blob = f.read()
    solve_flag(target_blob)