"""
Implementacja szyfru strumieniowego ChaCha
"""

import struct
import sys, os, binascii
from base64 import b64encode

def yield_chacha20_xor_stream(key, iv, position=0):
  # Generowanie strumienia ChaCha xor"""
  if not isinstance(position, int):
    raise TypeError
  if position & ~0xffffffff:
    raise ValueError('Argument position nie jest typu uint32.')
  if not isinstance(key, bytes):
    raise TypeError
  if not isinstance(iv, bytes):
    raise TypeError
  if len(key) != 32:
    raise ValueError
  if len(iv) != 8:
    raise ValueError

  def rotate(v, c):
    return ((v << c) & 0xffffffff) | v >> (32 - c)

  def quarter_round(x, a, b, c, d):
    x[a] = (x[a] + x[b]) & 0xffffffff
    x[d] = rotate(x[d] ^ x[a], 16)
    x[c] = (x[c] + x[d]) & 0xffffffff
    x[b] = rotate(x[b] ^ x[c], 12)
    x[a] = (x[a] + x[b]) & 0xffffffff
    x[d] = rotate(x[d] ^ x[a], 8)
    x[c] = (x[c] + x[d]) & 0xffffffff
    x[b] = rotate(x[b] ^ x[c], 7)

  ctx = [0] * 16
  ctx[:4] = (1634760805, 857760878, 2036477234, 1797285236)
  ctx[4 : 12] = struct.unpack('<8L', key)
  ctx[12] = ctx[13] = position
  ctx[14 : 16] = struct.unpack('<LL', iv)
  while 1:
    x = list(ctx)
    for i in range(10):
      quarter_round(x, 0, 4,  8, 12)
      quarter_round(x, 1, 5,  9, 13)
      quarter_round(x, 2, 6, 10, 14)
      quarter_round(x, 3, 7, 11, 15)
      quarter_round(x, 0, 5, 10, 15)
      quarter_round(x, 1, 6, 11, 12)
      quarter_round(x, 2, 7,  8, 13)
      quarter_round(x, 3, 4,  9, 14)
    for c in struct.pack('<16L', *(
        (x[i] + ctx[i]) & 0xffffffff for i in range(16))):
      yield c
    ctx[12] = (ctx[12] + 1) & 0xffffffff
    if ctx[12] == 0:
      ctx[13] = (ctx[13] + 1) & 0xffffffff

def chacha20_encrypt(data, key, iv=None, position=0):
  # Funkcja szyfrująca (odszyfrowująca) szyfrem ChaCha20
  if not isinstance(data, bytes):
    raise TypeError
  if iv is None:
    iv = b'\0' * 8
  if isinstance(key, bytes):
    if not key:
      raise ValueError('Klucz jest pustym łańcuchem znaków.')
    if len(key) < 32:
      key = (key * (32 // len(key) + 1))[:32]
    if len(key) > 32:
      raise ValueError('Klucz jest za długi.')

  return bytes(a ^ b for a, b in
      zip(data, yield_chacha20_xor_stream(key, iv, position)))

def main():
    # key = os.urandom(32)
    key = b'tajnyKlucz!!!' # 13 znaków
    print('Klucz użyty do szyfrowania: {}'.format(key))
    print()
    plaintext = b'Wszyscy zyjemy w zoltej lodzi podwodnej.'
    print('Tekst jawny: {}'.format(plaintext))
    
    iv = b'SekrWekt' # 8 znaków
    print()
    enc = chacha20_encrypt(plaintext, key, iv)
    decode_enc = b64encode(enc).decode('utf-8')
    print('Zaszyfrowana wiadomość: {}. '.format(decode_enc))
    print()
    dec = chacha20_encrypt(enc,key, iv)
    print('Odszyfrowana wiadomość: {}. '.format(dec))
    print()

if __name__ == "__main__":
    sys.exit(int(main() or 0))
