def rc4algo(msg, key):
	# key scheduling
	s = [i for i in range(256)]
	k = [ord(key[i]) for i in range(len(key))]
	# print (k)
	j = 0
	for i in range(256):
		j = (j + s[i] + k[i % len(key)]) % 256
		# print (j)
		s[i], s[j] = s[j], s[i]

	#pseudo-random generation
	r = []
	i = 0
	j = 0
	for x in range(len(msg)):
		i = (i + 1) % 256
		j = (j + s[i]) % 256
		s[i], s[j] = s[j], s[i]
		r.append(s[(s[i] + s[j]) % 256])

	#XOR
	c = [ord(msg[x]) ^ r[x] for x in range(len(msg))]
	ctext = ''
	for x in range(len(msg)):
		ctext += chr(c[x])

	return ctext


def main():
	print('Enter the message: ', end='')
	msg = input()
	print('Enter the key: ', end='')
	key = input()

	ch = 0
	while (ch != 3):
		print(
		    '1. To Encrypt\n2. To Decrypt\n3. To Exit\nEnter the choice: ',
		    end='')
		ch = int(input())

		msg = rc4algo(msg, key)
		if (ch == 1):
			print('Encrypted message: ', msg)
		elif (ch == 2):
			print('Decrypted message: ', msg)


main()