def add(x, y):
	return (x + y) & 0xffffffff

def rotl(x, n):
	return (x << n) & 0xffffffff | (x >> (32 - n))

def quarterround(a, b, c, d):
	a = add(a, b)
	d ^= a
	d = rotl(d, 16)

	c = add(c, d)
	b ^= c
	b = rotl(b, 12)

	a = add(a, b)
	d ^= a
	d = rotl(d, 8)

	c = add(c, d)
	b ^= c
	b = rotl(b, 7)

	return a, b, c, d

def u32(i):
	a, b, c, d = i
	return a | (b << 8) | (c << 16) | (d << 24)

def unu32(n):
	return bytes((n & 0xff, (n >> 8) & 0xff, (n >> 16) & 0xff, n >> 24))

def chacha20_block(key, counter, nonce, do_final_add = True):
	s = [
		0x61707865, 0x3320646e, 0x79622d32, 0x6b206574,
		u32(key[:4]), u32(key[4:8]), u32(key[8:12]), u32(key[12:16]),
		u32(key[16:20]), u32(key[20:24]), u32(key[24:28]), u32(key[28:]),
		counter, u32(nonce[:4]), u32(nonce[4:8]), u32(nonce[8:])
	]
	initial = s[:]

	def qr(i, j, k, l):
		s[i], s[j], s[k], s[l] = quarterround(s[i], s[j], s[k], s[l])
	for _ in range(10):
		qr(0, 4, 8, 12)
		qr(1, 5, 9, 13)
		qr(2, 6, 10, 14)
		qr(3, 7, 11, 15)
		qr(0, 5, 10, 15)
		qr(1, 6, 11, 12)
		qr(2, 7, 8, 13)
		qr(3, 4, 9, 14)

	if do_final_add:
		s = [add(initial[i], s[i]) for i in range(len(s))]
	return b''.join(unu32(i) for i in s)

def ceildiv(a, b):
	return a // b if a % b == 0 else a // b + 1

def chacha20(key, counter, nonce, message):
	for i in range(ceildiv(len(message), 64)):
		keystream = chacha20_block(key, counter + i, nonce)
		yield from (a ^ b for a, b in zip(message[i*64:i*64+64], keystream))

def prettyprint_state(s):
	for i in range(4):
		print('%08x  %08x  %08x  %08x' % (s[i*4], s[i*4+1], s[i*4+2], s[i*4+3]))
	print()

if __name__ == '__main__':
	key = bytes.fromhex('00:01:02:03:04:05:06:07:08:09:0a:0b:0c:0d:0e:0f:10:11:12:13:14:15:16:17:18:19:1a:1b:1c:1d:1e:1f'.replace(':',''))
	nonce = bytes.fromhex('00:00:00:00:00:00:00:4a:00:00:00:00'.replace(':',''))
	message = b"Ladies and Gentlemen of the class of '99: If I could offer you only one tip for the future, sunscreen would be it."
	print(bytes(chacha20(key, 1, nonce, message)).hex() == '6e2e359a2568f98041ba0728dd0d6981e97e7aec1d4360c20a27afccfd9fae0bf91b65c5524733ab8f593dabcd62b3571639d624e65152ab8f530c359f0861d807ca0dbf500d6a6156a38e088a22b65e52bc514d16ccf806818ce91ab77937365af90bbf74a35be6b40b8eedf2785e42874d')