from salsa20 import Salsa20_keystream, Salsa20_xor


def hexdump(data):
    print(' '.join([
        "{0:08x}".format(int.from_bytes(data[4*i:4*i+4], byteorder='little'))
        for i in range(len(data) // 4)
    ]))


def to_bytes(arr):
    return b''.join(map(lambda x: x.to_bytes(4, byteorder='big'), arr))


def print_keystream(nonce, key, block):
    stream = \
        Salsa20_keystream(
            length=64 * (block + 1),
            nonce=to_bytes(nonce),
            key=to_bytes(key),
        )

    plaintext = "cryptography is awesome"
    ciphertext = \
        Salsa20_xor(
            b'\x00'*(block * 64) + plaintext.encode(),
            nonce=to_bytes(nonce),
            key=to_bytes(key),
        )

    hexdump(stream[-64:])
    hexdump(ciphertext[-len(plaintext.encode()):])


print('reference implementation: libsodium (NaCl)')
print('generating keystream for: ')
print()
print('1) example from salsa20 specification: nonce: (3,1,4,1,5,9,2,6), key: (1,2,3,...), block: 7')
print_keystream(
    nonce=[0x03010401, 0x05090206],
    key=[0x01020304, 0x05060708, 0x090a0b0c, 0x0d0e0f10, 0x11121314, 0x15161718, 0x191a1b1c, 0x1d1e1f20],
    block=7
)

print()
print('2) homework from crypto class: nonce: (0x00c0ffee, 0x00c0ffee), key: (0xffffff...ffff), block: 0')
print_keystream(
    nonce=[0x00c0ffee, 0x00c0ffee],
    key=[0xffffffff, 0xffffffff, 0xffffffff, 0xffffffff, 0xffffffff, 0xffffffff, 0xffffffff, 0xffffffff],
    block=0
)