#!/usr/bin/python3

'''
Encrypts a message using RSA public 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_encrypt.py <public key file name> <output file name>
		python3 rsa_encrypt.py publicKey.pem ciphertext.bin
				
'''

from Crypto.Cipher import PKCS1_OAEP
from Crypto.PublicKey import RSA
import sys
import base64

message = b'We attack at dawn!'

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

	#read public key and set padding to OAEP
	print("Reading public key from",sys.argv[1])
	key = RSA.importKey(open(sys.argv[1],'rb').read()) 
	cipher = PKCS1_OAEP.new(key)   #use OAEP padding  

	#encrypt message using public key         
	ciphertext = cipher.encrypt(message)

	#write output file
	print("Writing output to",sys.argv[2])
	f = open(sys.argv[2],'wb')
	f.write(base64.b64encode(ciphertext))
	f.close()