#
#
#
import logging
from Crypto.Signature import PKCS1_v1_5
from Crypto.PublicKey import RSA
from Crypto.Hash import SHA512

class Signature:
  
  signature = None
  key_data  = None
  pub_key   = None
  priv_key  = None

  #
  def __init__(self):
    self.genKeys()  

  #
  def genSig(self, bin_data):
    hshh = SHA512.new()
    hshh.update(bin_data)
    signer = PKCS1_v1_5.new(self.priv_key)
    self.signature = signer.sign(hshh)
    return self.signature

  #
  def verifySig(self, bin_data, bin_sig):
    hshh = SHA512.new()
    hshh.update(bin_data)
    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  = RSA.importKey( self.key_data.publickey().exportKey('DER'))
    self.priv_key = RSA.importKey( self.key_data.exportKey('DER'))
    logging.debug(self.pub_key.exportKey('PEM'))
    logging.debug(self.priv_key.exportKey('PEM'))

  #
  def getKeyPEM(self):
    return self.key_data.exportKey('PEM')

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