import os
import codecs
from Crypto.PublicKey import RSA
from Crypto.Cipher import PKCS1_OAEP, AES
from Crypto.Hash import SHA256
from Crypto.Signature import PKCS1_v1_5
from Crypto import Random
from base64 import b64decode, b64encode


BS = 16


pub_key_contentserver = '-----BEGIN PUBLIC KEY-----\nMIGfMA0GCSqGSIb3DQEBAQUAA4GNADCBiQKBgQDYI/jGk6f0LAT2Z6AAQUR7izbi\na5O4Zzaz4WtK00jr3AqbZMVeVZAs+As5RS35PY2BlCuEEza/J5XX1tlbUhGk/Nzu\nyYqlID6ILEk9kUqh1A6EAuNVrcCL174BRLy620pU5m+E61za0tIr1lU+Jhy4ikVK\niGQ+na5a5g0kuzZTHwIDAQAB\n-----END PUBLIC KEY-----'
private_key_contentserver = '-----BEGIN RSA PRIVATE KEY-----\nMIICXQIBAAKBgQDYI/jGk6f0LAT2Z6AAQUR7izbia5O4Zzaz4WtK00jr3AqbZMVe\nVZAs+As5RS35PY2BlCuEEza/J5XX1tlbUhGk/NzuyYqlID6ILEk9kUqh1A6EAuNV\nrcCL174BRLy620pU5m+E61za0tIr1lU+Jhy4ikVKiGQ+na5a5g0kuzZTHwIDAQAB\nAoGAFuf9CgKFBrYRkpGZ3d0M5nDzEJzpC1546ChoKAZrUH/B6gUMe7pirLle6yNf\nQ25YDFcJI5arsyd9VGITJ//zGZC1CC9MBL1v4Ey5Vm3OAwcn7aoiPYoZhtNUDGhE\nCKzIrKWlq5wQDfYCjL2zQnCIAq+nZWACI/X+k/bPd37FZpUCQQDdQNQ3Azaq0owx\nd2crResCW4OryvhTPipLNz8AaW2IMuDfxLUKEaTTkU3NOFRRhIddRRA9/vfcHbrw\nlgZ3Vo5NAkEA+hWX0tItz4lL4pSJTc38Tzjb9WyGQPuAINlntnDwOhtvvRYPfTLi\nLb/22btAfgQWSBICxFdO5Ze0A5Dx140VGwJARwVhWYtZh/nv8I0Ae/6EkowntwR/\nM9FXqC9CtPIiq76ROqMc7e999j/FNqPnRQeCoCjkLtJiY7DTahjuWG5bXQJBAIAD\nS5scAV0pz5FlLT+JgGzhEx729WYQF085WjB2cVGdN75Xq4gP4t0+VVKw2ltnJiyw\nI4BznKxD0l689D69Nh8CQQDLkom4Fhkl8CZQLB1cIX9PjNw+KCBkTV92dJmzdmR5\nGE+p7jO849HPS83W21AV6/eyL8VMqDyE6CBzHNr2ITJT\n-----END RSA PRIVATE KEY-----'

pub_key_licenseserver= "-----BEGIN PUBLIC KEY-----MIGfMA0GCSqGSIb3DQEBAQUAA4GNADCBiQKBgQCG54kIgf8wi+WLrmg0OLXTsMF7BsYjJ8+mmx1/Ml1Hts2oVoisf8iKbfcC0jWBc1aEibi7EIZK9AWEmgpIHPZpf2Ezt36KX7pjA7CDgZBye/bw+G1CaDGxrN8a18db8c16hNxHMTMiqPm7WAgoyekOzlCT6/rqF+3w2oYuGx1tmQIDAQAB-----END PUBLIC KEY-----"


def make_key_pair_rsa(username):
    random_generator = Random.new().read
    priv_key = RSA.generate(1024, random_generator)
    pub_key = priv_key.publickey()
    print("new key pair")
    print(priv_key.exportKey())
    print(pub_key.exportKey())
    user_smartcard_path = "/tmp/smartcards/" + str(username) + ".txt"
    file_object = open(user_smartcard_path, "w+")
    file_object.write(priv_key.exportKey().decode('utf-8'))
    file_object.write("\n")
    file_object.write(pub_key.exportKey().decode('utf-8'))
    file_object.close()
    return pub_key.exportKey().decode('utf-8')


def make_rand_key_aes():
    return os.urandom(32)


def pad(s):
    return s + (BS - len(s) % BS) * chr(BS - len(s) % BS).encode()


def unpad(s):
    return s[:-ord(s[len(s)-1:])]


def encrypt_aes(message, key):
    message = message.encode()
    raw = pad(message)
    cipher = AES.new(key, AES.MODE_ECB)
    enc = cipher.encrypt(raw)
    return b64encode(enc).decode('utf-8')


def decrypt_aes(enc, key):
    enc = b64decode(enc)
    cipher = AES.new(key, AES.MODE_ECB)
    dec = cipher.decrypt(enc)
    return dec.decode('utf-8')


def make_encrypted_key_content_server():
    random_key = make_rand_key_aes()
    key = RSA.importKey(pub_key_contentserver)
    pub_key = key.publickey()
    encrypted_key = encrypt(random_key, pub_key)
    return encrypted_key.hex()


def decrypt(ciphertext):
    key = RSA.importKey(private_key_contentserver)
    cipher = PKCS1_OAEP.new(key)
    ciphertext = codecs.decode(ciphertext, 'hex')
    return cipher.decrypt(ciphertext)


def encrypt(message):
    cipher = PKCS1_OAEP.new(pub_key_contentserver)
    return cipher.encrypt(message)


def sign(message, priv_key):
    priv_key = RSA.importKey(priv_key)
    signer = PKCS1_v1_5.new(priv_key)
    hash_gen = SHA256.new()
    hash_gen.update(message.encode())
    return b64encode(signer.sign(hash_gen)).decode('utf-8')


def verify(message, signature, pub_key):
    pub_key = pub_key.replace("-----BEGIN PUBLIC KEY-----", "")
    pub_key = pub_key.replace("-----END PUBLIC KEY-----", "")
    pub_key = pub_key.replace(" ", "")
    pub_key = ('\n'.join(pub_key[i:i + 64] for i in range(0, len(pub_key), 64)))
    pub_key = "-----BEGIN PUBLIC KEY-----\n" + pub_key + "\n-----END PUBLIC KEY-----"
    signature = b64decode(signature)
    pub_key = RSA.importKey(pub_key)
    verifier = PKCS1_v1_5.new(pub_key)
    hash_gen = SHA256.new(message.encode())
    return verifier.verify(hash_gen, signature)