from Crypto.PublicKey import RSA
from Crypto.Hash import SHA256
from Crypto.Signature import pkcs1_15
import json
from base64 import b64encode, b64decode
from pprint import pprint

class ITKoin:
    def __init__ (self):
        self.pending_transactions = []
        self.unspent_transactions = []
        self.sender_inputs = []
        self.chain = []

    @staticmethod
    def generate_rsa_key(filename):
        rsakey = RSA.generate(2048)
        rsapublickey = rsakey.publickey()
        pprint(vars(rsakey))
        pprint(vars(rsapublickey))
        PEMrsakey = rsakey.export_key('PEM')
        pprint(PEMrsakey)
        PEMrsapublickey = rsapublickey.export_key('PEM')
        pprint(PEMrsapublickey)
        privatekeyfilename = filename + 'priv.pem'
        f = open(privatekeyfilename, 'wb')
        f.write(PEMrsakey)
        f.close()
        publickeyfilename = filename + 'pub.pem'
        f = open(publickeyfilename, 'wb')
        f.write(PEMrsapublickey)
        f.close()
        return

    def load_key(self, filename):
        privatekeyfilename = filename+'priv.pem'
        privatekeyfileobject = open(privatekeyfilename, 'r')
        privatekeyfilecontent = privatekeyfileobject.read()
        pprint(privatekeyfilecontent)
        rsakey = RSA.import_key(open(privatekeyfilename,'r').read())
        self.rsakey = rsakey
        pprint(vars(self.rsakey))
        return

    def load_public_key (self, filename):
        publickeyfilename = filename+'pub.pem'
        publickeyfileobject = open(publickeyfilename, 'r')
        publickeyfilecontent = publickeyfileobject.read()
        pprint(publickeyfilecontent)
        rsakey = RSA.import_key(open(publickeyfilename,'r').read())
        rsapublickey = rsakey.publickey()
        self.rsapublickey = rsapublickey
        pprint(vars(self.rsapublickey))
        return

    @staticmethod
    def create_hashobject(data):
        stringdump = json.dumps(data)
        binarydump = stringdump.encode()
        hashobject = SHA256.new(data=binarydump)

        hashhexvalue = hashobject.hexdigest()
        print(hashhexvalue)
        return hashobject

    def create_signature(self, data):
        signatureobject = pkcs1_15.new(self.rsakey)
        hashobject = self.create_hashobject(data)
        signaturevalue = signatureobject.sign(hashobject)
        print(signaturevalue)
        b64signaturevalue = b64encode(signaturevalue)
        print(b64signaturevalue)
        print(b64signaturevalue.decode())
        return b64signaturevalue.decode()

    def verify_signature(self, data, b64signaturevalue, rsapublickey):
        verifyobject = pkcs1_15.new(rsapublickey)
        hashobject = self.create_hashobject(data)
        signaturevalue = b64decode(b64signaturevalue)
        try:
            signatureerror = verifyobject.verify(hashobject, signaturevalue)
            validsignature = True
        except (ValueError, TypeError):
            validsignature = False
        return validsignature

coinTest = ITKoin()
keyFileName = "firstkey"
coinTest.generate_rsa_key(keyFileName)
coinTest.load_key(keyFileName)
coinTest.load_public_key(keyFileName)
message = 'Hello world'
sign = coinTest.create_signature(message)
signBase = b64encode(json.dumps(sign).encode())
val = coinTest.verify_signature(message, signBase, RSA.import_key(open("firstkeypub.pem",'r').read()).publickey())