#!/usr/bin/python3
'''
Signs a message using RSA private key (possibly generated with rsa_key_gen.py)

Author: Tim Pierson, Dartmouth CS55, Winter 2021
				Based Wenliang Du https://github.com/kevin-w-du/BookCode/tree/master/Public_Key_Cryptography
Usage: 	python3 rsa_sign.py <private key file name> <password> <signature file name>
			where password is the password used during key generation (from rsa_key_gen.py)
		python3 rsa_sign.py private.pem cs55 signature.bin

		NOTE: make sure the key length used when generating private.pem is longer than the length of a SHA digest!
				
'''

from Crypto.Signature import pss
from Crypto.Hash import SHA256
from Crypto.PublicKey import RSA
import sys

message = b'An important message' #same message used in rsa_veryify.py

if __name__ == '__main__':
	#usage check: make sure we got a filename to read
	if len(sys.argv) != 4:
		print("Usage: python3 rsa_sign.py <private key file name> <password> <signature file name")
		print("Did you provide these?")
		exit()

	print("Reading private key",sys.argv[1])
	key_pem = open(sys.argv[1]).read()
	key = RSA.import_key(key_pem, passphrase=sys.argv[2])

	print("Message to sign is:",message)
	print("Hashing message")
	h = SHA256.new(message)

	print("Signing hash")
	signer = pss.new(key)                    
	signature = signer.sign(h) 

	print("Writing signature to",sys.argv[3])             
	open(sys.argv[3], 'wb').write(signature)