LFSR python script for Assignment 2 of Applied Cryptography in QUT

question2.py

def xor(a, b):
    """
    Helper function to perform XOR operation on binary strings
    Source: https://www.tutorialspoint.com/tuple-xor-operation-in-python
    """

    return ''.join(str(int(x) ^ int(y)) for x, y in zip(a, b))


def berlekamp_massey(keystream):
    """
    Performs the Berlekamp-Massey algorithm to determine the LFSR configuration from a given keystream.

    Parameters:
    keystream (str): Binary input.

    Returns:
    tuple: A tuple returns with LFSR (lfsr_len) and the coefficients of the LFSR (lfsr_coeff).
    lfsr_len (int): Length of the LFSR configuration.
    lfsr_coeff (list): Coefficients representing the LFSR configuration for generating the keystream.

    Source:
    - https://github.com/bozhu/BMA
    - https://gist.github.com/StuartGordonReid/a514ed478d42eca49568
    - https://github.com/thewhiteninja/lfsr-berlekamp-massey/blob/master/berlekamp_massey.py
    """

    n = len(keystream) # Get the length of the keystream
    # Initialise with 1 followed by n of zeros.
    b = [1] + [0] * n
    c = [1] + [0] * n
    l = 0
    
    for i in range(n): # Loop through each binary
        d = 0
        for j in range(l + 1):
            # Cannot use xor() because it's bitwise on individual binary digits
            d ^= c[j] * int(keystream[i - j]) 
        
        # if the bit is 1, go to the branching codition
        if d == 1:
            t = c.copy()
            for j in range(n - l):
                c[l + j] ^= b[j]
            # Chcek if the length sould be updated.
            if l <= i // 2:
                l = i + 1 - l
                b = t
    
    lfsr_len = l
    lfsr_coeff = c[:l + 1]
    return lfsr_len, lfsr_coeff


def generate_keystream(lfsr_len, lfsr_coeff, keystream_prefix):

    """
    Generate the ciphertext using the recovered LFSR configuration
    
    Parameters:
    lfsr_len (int): Length of the LFSR configuration
    lfsr_coeff (list): Coefficients of the LFSR
    keystream_prefix (str): Partial known keystream
    
    Returns:
    str: Ciphertext generated using the provided LFSR configuration

    Source:
    - https://gist.github.com/Nabiljain/72bb87f0c3762a59a305c17079cb85ed
    """

    # Initialise LFSR with the known keystream prefix
    lfsr = list(map(int, keystream_prefix[:lfsr_len]))
    keystream = keystream_prefix

    # Generate additions binary for the keystream until it is 48 characters long.
    # Total length should be 48 (ciphertext length) - 32(cipher) = 16
    while len(keystream) < len(keystream_prefix) + 16: 
        # To complete 48-bit keystream calculate the feedback by summing the LFSR with coefficient.
        feedback = sum(lfsr[i] * lfsr_coeff[i] for i in range(lfsr_len)) % 2
        # Applend the first element of the LFSR to the keystream
        keystream += str(lfsr[0])
        # Have to shift it to left because LFSR.
        lfsr = lfsr[1:] + [feedback]
    return keystream


# Given values
p_prime = '11110111011011000010001011011110'
c_prime = '100001011100000100101001101111111001000111100111'
c = '111000110010011111111010010100101010110010001011'

# Find the first 32 bits of the keystream
keystream_32 = xor(p_prime, c_prime)
print(f"Keystream 32: {keystream_32}")

# Recover the LFSR configuration using Berlekamp-Massey algorithm
lfsr_len, lfsr_coeff = berlekamp_massey(keystream_32)
print(f"LFSR Length Found: {lfsr_len}\nLFSR Coefficients Found: {lfsr_coeff}")

# Generate the remaining 16 bits of the keystream using the recovered LFSR
keystream_48 = generate_keystream(lfsr_len, lfsr_coeff, keystream_32)
keystream_16 = keystream_48[32:]

# Find the remaining 16 bits of the plaintext
plaintext_16 = xor(keystream_16, c[32:])

# Print the answers
plaintext_32 = xor(keystream_32, c[:32])

# Verification: Encrypt the recovered plaintext and compare with the given ciphertext
recovered_plaintext = plaintext_32 + plaintext_16
recovered_ciphertext = xor(recovered_plaintext, keystream_48)

print("\nVerification:")
if recovered_ciphertext == c:
    print("Success! The recovered plaintext is correct.")
    print(f"Recovered Plaintext (32 bits): {plaintext_32}")
    print(f"Recovered Plaintext (16 bits): {plaintext_16}")
    print(f"Recovered Ciphertext: {recovered_ciphertext}")
    print(f"2a {plaintext_32}")
    print(f"2b {plaintext_16}")
else:
    print("Error: The recovered plaintext does not match the given ciphertext.")