import sys
from struct import pack, unpack

# The Salsa20/12 quarter-round function
def quarter_round(y0, y1, y2, y3):
    z1 = y1 ^ ((y0 + y3) % (1 << 32) << 7) % (1 << 32)
    z2 = y2 ^ ((z1 + y0) % (1 << 32) << 9) % (1 << 32)
    z3 = y3 ^ ((z2 + z1) % (1 << 32) << 13) % (1 << 32)
    z0 = y0 ^ ((z3 + z2) % (1 << 32) << 18) % (1 << 32)
    return z0, z1, z2, z3

# The Salsa20/12 row-round function
def row_round(y):
    z = [0] * 16
    z[0], z[1], z[2], z[3] = quarter_round(y[0], y[1], y[2], y[3])
    z[5], z[6], z[7], z[4] = quarter_round(y[5], y[6], y[7], y[4])
    z[10], z[11], z[8], z[9] = quarter_round(y[10], y[11], y[8], y[9])
    z[15], z[12], z[13], z[14] = quarter_round(y[15], y[12], y[13], y[14])
    return z

# The Salsa20/12 column-round function
def column_round(x):
    y = [0] * 16
    y[0], y[4], y[8], y[12] = quarter_round(x[0], x[4], x[8], x[12])
    y[5], y[9], y[13], y[1] = quarter_round(x[5], x[9], x[13], x[1])
    y[10], y[14], y[2], y[6] = quarter_round(x[10], x[14], x[2], x[6])
    y[15], y[3], y[7], y[11] = quarter_round(x[15], x[3], x[7], x[11])
    return y

# The Salsa20/12 double-round function
def double_round(x):
    return row_round(column_round(x))

# The Salsa20/12 hash function
def salsa20_12_hash(x):
    z = list(x)  # Convert to list to allow modifications
    for i in range(6):  # 12 rounds --> 6 double-rounds
        z = double_round(z)
    return [(z[i] + x[i]) & 0xffffffff for i in range(16)]

# The Salsa20/12 expansion function
def salsa20_12_expand(key, nonce, block_index):
    # Constants defined by the Salsa20 algorithm
    sigma = b'expand 32-byte k' if len(key) == 32 else b'expand 16-byte k'
    # Split the key for the 16 byte version using the sigma constant
    if len(key) == 16:
        k0, k1 = key[:16], key[:16]
    else:
        k0, k1 = key[:16], key[16:]
    # Set up the block index (nonce expansion)
    n = pack('<Q', block_index)
    
    # Salsa20 state composition
    return (sigma[:4] + k0 + sigma[4:8] + nonce + n + sigma[8:12] + k1 + sigma[12:])

# Salsa20/12 encryption/decryption (identical operations due to XOR)
def salsa20_12_crypt(key, nonce, text):
    block_size = 64  # 64 bytes
    encrypted_text = b''
    num_blocks = (len(text) + block_size - 1) // block_size  # Calculate the number of blocks

    assert len(nonce) == 8, "Nonce must be 64 bits (8 bytes)"

    for block_count in range(num_blocks):
        block_start = block_count * block_size
        block_end = block_start + block_size
        block = text[block_start:block_end].ljust(block_size, b'\x00')

        # Expand key, nonce, and counter to the Salsa20 state
        state = salsa20_12_expand(key, nonce, block_count)
        keystream = salsa20_12_hash(unpack('<16I', state))

        # Serialize keystream back to bytes and truncate to the original block size
        keystream_bytes = b''.join(pack('<I', word) for word in keystream)[:block_size]

        # XOR the block with the keystream
        block_encrypted = bytes(a ^ b for a, b in zip(block, keystream_bytes[:len(block)]))
        encrypted_text += block_encrypted[:len(text[block_start:block_end])]

    return encrypted_text.hex()


if __name__ == '__main__':
    key_length = int(input("Enter the key length in bits (128, or 256): "))
    # Note: Salsa20 does not support 64 bit keys so removed that option.
    key_hex = input("Enter the key (as a hex string): ")
    nonce_hex = input("Enter the nonce (as a hex string - 16 hex characters): ")
    text_hex = input("Enter the text (as a hex string for encryption/decryption): ")

    if key_length not in [128, 256]:
        print("Invalid key length! Supported key lengths are 128, and 256 bits.")
        sys.exit(1)

    key = bytes.fromhex(key_hex)
    nonce = bytes.fromhex(nonce_hex)
    text = bytes.fromhex(text_hex)

    # Ensure that the key and nonce have the correct length
    if len(key) not in {16, 32}:  # Key must be 128 or 256 bits (16 or 32 bytes)
        print(f"Key must be {key_length} bits long (hex string of length {key_length // 4}).")
        sys.exit(1)
    if len(nonce) != 8:  # Nonce must be 64 bits (8 bytes)
        print("Nonce must be 64 bits long (hex string of length 16).")
        sys.exit(1)

    # Encrypt or decrypt the text
    result = salsa20_12_crypt(key, nonce, text)
    print("Result:", result)

This implementation must support key sizes of 128-bit, 256-bit, and in-addition, non-standard 64-bit. The constant string ”expand 08-byte k” is used for 64-bit key, and key is repeated 4 times to fill the initial state in expansion function. Here is a formal definition of the non-standard 64- bit key Salsa20 expansion function:   Define α0 = (101, 120, 112, 97), α1 = (110, 100, 32, 48), α2 =
(56, 45, 98, 121), α3 = (116, 101, 32, 107). If k is 8-byte sequences and n is 16-byte sequences then
Salsa20k(n) = Salsa20(α0, k, k, α1, n, α2, k, k, α3)
Your program must accept the following inputs from the command-line: the key length(in bits), the key represented as a hexadecimal string, nonce(IV) represented as a hexadecimal string, and the text string to be encrypted or decrypted also in hexadecimal format. For encryption, this will be plaintext. For decryption, it will be ciphertext. The program’s output should be the resulting ciphertext when performing encryption or the original plaintext when performing decryption. They are in the hexadecimal format.  A command-line example:
#Input:
./your_prog 128 "deadbeefdeadbeefdeadbeefdeadbeef" \ "1234567890abcdef" "546869736973706c61696e74657874"
#Outpt:
"a1c7720e1abadb96e5a2600d0ce028"
In this example, 128 represents the key size in bits, ”deadbeefdeadbeefdeadbeefdeadbeef” is the hexadecimal string key, ”1234567890abcdef” is the nonce (IV) in hexadecimal string format, and the final argument is the input text, ”546869736973706c61696e74657874” hexadecimal format of ASCII string ”Thisisplaintext”, which can be either plaintext (for encryption) or ciphertext (for decryption). We assume the block number in the encryption and decryption always starts from 0 in this assignment. Eventually, your program should output the resulting(ciphertext/plaintext) string to the command-line. Recommended implementation steps:
•	Read the specification thoroughly. Pay special attention to the algorithm’s core functions.
•	Implement the quarterround function first. This function is the basic building block of the Salsa20 cipher.
•	Then, implement the rowround and columnround function on top of that quarterround function.
 
•	Implement the doubleround function based on rowround and columnround.
•	Implement the Salsa20 hash function based on doubleround function.
•	Implement the expansion function for different key sizes. Pay attention to the different values used for the different keys.
•	Implement the encryption/decryption that can handle the command-line inputs correctly.
•	Test your implementation correctness of each function using the input/output inside the specification, make sure its correctness and move to implement the next function.
•	Implement the standard Salsa20/20 in the specification, make sure it works as expected, then you can change it to Salsa20/12 without much effort.
•	Then make the necessary modification to support the non-standard 64-bit key.