#
# 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):
    key_data  = None
    pub_key   = None
    sig_data  = None
    hh        = None
    self.hh = SHA512.new()
    self.genKeys()  

  #
  def genSig(self, hshh):
    signer = PKCS1_v1_5.new(self.key_data)
    self.sig_data  = signer.sign(hshh)
    return self.sig_data 

  #
  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')))

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

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

  #
  def getSignature(self):
    return sig_data

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


#
# TESTING
#

#
def _test():
  msg = b'Hieee, this is a test =)'

  sig = Signature()
  sig_data = sig.genSig(msg)
  print("Signature created")
  print("sig_data=={}".format(sig_data))

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


if __name__ == '__main__':
  _test()
