from mbedtls import cipher as mbedcipher
from unicrypto.symmetric import symmetricBASE, cipherMODE

class RC4(symmetricBASE):
	def __init__(self, key):
		symmetricBASE.__init__(self, key)

	def setup_cipher(self):
		if len(self.key) == 16:
			self._cipher = mbedcipher.ARC4.new(self.key, mbedcipher.MODE_ECB, b'')
		else:
			from unicrypto.backends.pure.RC4 import RC4 as _pureRC4
			self._cipher = _pureRC4(self.key)
	def encrypt(self, data):
		return self._cipher.encrypt(data)
	def decrypt(self, data):
		return self._cipher.decrypt(data)