Lfsr

from Crypto.Util.number import long_to_bytes

# The output stream from the challenge
stream_output = "..."  # Paste the 2048-bit string here

def berlekamp_massey(sequence):
    """
    Berlekamp-Massey algorithm to find the shortest LFSR that generates the sequence.
    Returns the connection polynomial coefficients (tap positions).
    """
    n = len(sequence)
    # Convert string to list of integers if needed
    if isinstance(sequence, str):
        sequence = [int(bit) for bit in sequence]
    
    # Initialize variables
    C = [0] * n  # Connection polynomial
    B = [0] * n  # Previous connection polynomial
    C[0] = 1
    B[0] = 1
    L = 0  # Current length of LFSR
    m = -1  # Last length change
    b = 1
    
    for N in range(n):
        # Calculate discrepancy
        d = sequence[N]
        for i in range(1, L + 1):
            d ^= C[i] & sequence[N - i]
        
        if d == 1:
            T = C.copy()
            # Update C
            for i in range(n):
                if i >= N - m and i - (N - m) < n:
                    C[i] ^= B[i - (N - m)]
            
            if L <= N // 2:
                L = N + 1 - L
                m = N
                B = T
                b = d
    
    # Return the tap positions (indices where C[i] = 1, excluding C[0])
    taps = [i for i in range(1, L + 1) if C[i] == 1]
    return taps, L

def reverse_lfsr(stream, taps, lfsr_length, warmup_steps=768):
    """
    Given the output stream and tap positions, recover the initial state.
    We'll build a system of linear equations and solve it.
    """
    # We need to find the initial state that, after warmup_steps + clocking,
    # produces the given stream
    
    # For simplicity, let's use a different approach:
    # We'll reconstruct the state by running forward from an assumed state
    # But actually, we need the state BEFORE warmup
    
    # Let's use Gaussian elimination on GF(2)
    # Each output bit is a linear combination of state bits
    
    # Actually, simpler approach: use the stream to find the state just before output
    # Then reverse the warmup
    
    # The state when output starts can be found from the stream and taps
    state = [0] * lfsr_length
    
    # Initialize with first lfsr_length bits of stream
    # This is the state at the START of output (after warmup)
    for i in range(min(lfsr_length, len(stream))):
        state[i] = int(stream[i])
    
    # Now reverse the warmup steps
    for _ in range(warmup_steps):
        # Reverse clock: new bit goes to front, last bit is computed
        new_bit = state[-1]
        for tap in taps:
            new_bit ^= state[lfsr_length - tap - 1]
        state = [new_bit] + state[:-1]
    
    return state

def solve_challenge(stream_string):
    """
    Main solving function
    """
    print("Step 1: Converting stream to list of integers...")
    stream = [int(bit) for bit in stream_string.strip()]
    print(f"Stream length: {len(stream)} bits")
    
    print("\nStep 2: Running Berlekamp-Massey algorithm to find tap positions...")
    taps, lfsr_length = berlekamp_massey(stream)
    print(f"LFSR length found: {lfsr_length}")
    print(f"Tap positions: {taps}")
    
    print("\nStep 3: Recovering initial state before warmup...")
    # We expect lfsr_length to be 384 (48 bytes * 8 bits)
    if lfsr_length != 384:
        print(f"Warning: Expected 384 bits, got {lfsr_length}")
    
    # Reconstruct the state at the point where output starts
    # The stream gives us the output, we need to work backwards
    
    # Better approach: the first lfsr_length bits of output determine the state
    # at that point (the state after warmup)
    state_after_warmup = stream[:lfsr_length]
    
    print("\nStep 4: Reversing the warmup phase (768 steps)...")
    initial_state = reverse_lfsr_from_state(state_after_warmup, taps, lfsr_length, 768)
    
    print("\nStep 5: Converting binary state to FLAG...")
    # Convert list of bits to bytes
    binary_string = ''.join(map(str, initial_state))
    flag_int = int(binary_string, 2)
    flag = long_to_bytes(flag_int)
    
    print(f"\nFLAG: {flag}")
    return flag

def reverse_lfsr_from_state(state, taps, length, steps):
    """
    Reverse the LFSR from a known state by going backwards.
    """
    state = list(state)  # Make a copy
    
    for _ in range(steps):
        # In reverse: 
        # - The first bit becomes the last bit (we shift right)
        # - We need to compute what the first bit was
        
        # The last bit was computed as: XOR of taps
        # In forward: new_bit = state[0] ^ state[tap1] ^ state[tap2] ^ ...
        # We're going backward, so we need to recover what was shifted out
        
        # Actually, in reverse clock:
        # - Last bit moves to first position
        # - We XOR the taps to find what should be at the end
        
        last_bit = state[-1]
        state = state[:-1]  # Remove last bit
        
        # Compute the bit that was shifted in (at position 0 originally)
        first_bit = last_bit
        for tap in taps:
            if tap < len(state):
                print(len(state) - tap)
                first_bit ^= state[len(state) - tap]
        
        state = [first_bit] + state
    
    return state

# Example usage:
# Paste your stream output here
stream_output = """
"""

# Uncomment to run:
solve_challenge(stream_output)