#
# This Class Provides Crypto Functions
#
import logging
from Crypto.Signature import PKCS1_v1_5
from Crypto.PublicKey import RSA
from Crypto.Hash import SHA512

class Signature:

  #
  def __init__(self):
    self.key_data  = None
    self.pub_key   = None
    self.sig_data  = None
    self.hash_data = None
    self.hh        = SHA512.new()

  #
  def genSig(self):
    signer = PKCS1_v1_5.new(self.key_data)
    self.sig_data  = signer.sign(self.hh)

  #
  def verifySig(self, hshh, bin_sig):
    verifier = PKCS1_v1_5.new(self.pub_key)
    return verifier.verify(hshh, bin_sig)

  #
  def genKeys(self):
    logging.info("Generating public keys...")
    self.key_data = RSA.generate(4096)
    self.pub_key  = self.key_data.publickey()
    logging.debug("public key==\n{}".format(self.pub_key.exportKey('PEM')))
    logging.debug("private key (and pub, other info)==\n{}".format(self.key_data.exportKey('PEM')))

  #
  def importPubKey(self, bin_pubkey):
    self.pub_key = RSA.importKey(bin_pubkey)


  #
  def importPvtKey(self, bin_pvtkey):
    self.key_data = RSA.importKey(bin_pvtkey)


  #
  # Private key also includes public key data...SO ANNOYING!
  #
  def getPvtKeyPEM(self):
    return self.key_data.exportKey('PEM')

  #
  def getPubKeyPEM(self):
    return self.pub_key.exportKey('PEM')

  #
  def getPvtKeyDER(self):
    return self.key_data.exportKey('DER')

  #
  def getPubKeyDER(self):
    return self.pub_key.exportKey('DER')

  #
  def getSignature(self):
    return sig_data

  #
  def getNewHashHandle(self):
    self.hh = SHA512.new()
    return self.hh

  #
  def hashBin(self, bin_data):
    hshh = SHA512.new()
    hshh.update(bin_data)
    img_hash = hshh.digest()
    return img_hash


#
# TESTING
#

#
def _test():
  #logging.basicConfig(level=logging.DEBUG)

  msg = b'Hieee, this is a test =)'
  
  sig = Signature()
  sig.hh.update(msg)
  sig.hash_data = sig.hh.digest()
  sig.genSig()
  print("Hash created:\n{}".format(sig.hash_data))
  print("Size=={}".format(str(len(sig.hash_data))))
  print("Signature created:\n{}".format(sig.sig_data))
  print("Size=={}".format(str(len(sig.sig_data))))

  isVerified = sig.verifySig(sig.hh, sig.sig_data)
  print("isVerified=={}".format(str(isVerified)))


if __name__ == '__main__':
  _test()