from utils.singleton import singleton
from utils import common, file_op
from config import constant
import os, base64, rsa

PADDING = b'{'
BLOCK_SIZE = 16

# one-liner to sufficiently pad the text to be encrypted
pad = lambda s: s + (BLOCK_SIZE - len(s) % BLOCK_SIZE) * PADDING

@singleton
class Ksecurity():
	def __init__(self):
		self.aes = None
		self.iv = None
	
	def on_initializing(self, *args, **kwargs):
		(self.pubkey, self.privkey) = rsa.newkeys(1024)
		self.server_pubkey = rsa.PublicKey.load_pkcs1(self.read_server_publickey())

		return True
	
	def read_server_publickey(self):
		return file_op.cat(os.path.join(common.get_work_dir(), constant.SERVER_PUBLIC_KEY), "r")
		
	def rsa_long_encrypt(self, msg, length = 100):
		msg = msg.encode("utf8")
		res = []

		for i in range(0, len(msg), length):
			res.append(rsa.encrypt(msg[i : i + length], self.server_pubkey))

		return b"".join(res)

	def rsa_long_decrypt(self, crypto, length = 128):
		res = []

		for i in range(0, len(crypto), length):
			res.append(rsa.decrypt(crypto[i : i + length], self.privkey))

		return b"".join(res)

	def get_pubkey(self):
		return self.pubkey.save_pkcs1().decode()

	def swap_publickey_with_server(self, socket):
		response = {
			"cmd_id" : "10000",
			"key" : self.get_pubkey()
		}

		socket.response(response)
		
	def set_aes_iv(self, aes, iv):
		if aes and iv:
			self.aes = aes
			self.iv = iv
		
	def can_aes_encrypt(self):
		return self.aes and self.iv

	def reset_aes(self):
		self.aes = None
		self.iv = None

	def aes_encrypt(self, data):
		count = 0
		encrypt = []

		if common.is_python2x():
			for i in data:
				encrypt.append(chr(ord(i) ^ ord(self.aes[count % len(self.aes)])))
				count += 1
			
			return b"".join(encrypt)
		else:
			for i in data:
				encrypt.append(i ^ self.aes[count % len(self.aes)])
				count += 1

			return bytes(encrypt)
				
		#aes_obj_enc = AES.new(self.aes, AES.MODE_CBC, self.iv)
		#return aes_obj_enc.encrypt(pad(data))
		
	def aes_decrypt(self, encrypt):
		return self.aes_encrypt(encrypt)
		#aes_obj_enc = AES.new(self.aes, AES.MODE_CBC, self.iv)
		#return aes_obj_enc.decrypt(encrypt).rstrip(PADDING)