import numpy as np

def RC4Stream(n, key):
	output = np.zeros(n, dtype=int)
	K = np.array(key, dtype=int)
	keylength = K.shape[0]
	sizeS = 256

	# SETUP
	S = np.zeros(sizeS, dtype=int)
	T = np.zeros(sizeS, dtype=int)
	for i in range(sizeS):
		S[i] = i
		T[i] = K[i % keylength]

	# INITIAL PERMUTATION
	j = 0
	for i in range(sizeS):
		j = (j + S[i] + T[i]) % sizeS
		S_temp = S[j]
		S[j] = S[i]
		S[i] = S_temp

	# STREAM GENERATION
	i, j = 0, 0
	for k in range(n):
		i = (i+1) % sizeS
		j = (j + S[i]) % sizeS
		# Swap S[i] and S[j]
		S_temp = S[j]
		S[j] = S[i]
		S[i] = S_temp
		# Save output value
		output[k] = S[(S[i] + S[j]) % sizeS]

	return output

key = [1, 2, 3, 4, 5]
print(RC4Stream(10, key))
