import sys

from cryptography.hazmat.primitives import serialization, hashes
from cryptography.hazmat.primitives.asymmetric import rsa, padding


def generate_key_pair(pub_name, priv_name, key_size, passwd=None):

	private_key = rsa.generate_private_key(
		public_exponent=65537,
		key_size=int(key_size)
	)
	public_key = private_key.public_key()

	with open(pub_name, 'wb') as f:
		f.write(public_key.public_bytes(
			encoding=serialization.Encoding.PEM,
			format=serialization.PublicFormat.SubjectPublicKeyInfo
		))

	if not passwd:
		with open(priv_name, 'wb') as f:
			f.write(private_key.private_bytes(
				encoding=serialization.Encoding.PEM,
				format=serialization.PrivateFormat.TraditionalOpenSSL,
				encryption_algorithm=serialization.NoEncryption()
			))
	else:
		with open(priv_name, 'wb') as f:
			f.write(private_key.private_bytes(
				encoding=serialization.Encoding.PEM,
				format=serialization.PrivateFormat.TraditionalOpenSSL,
				encryption_algorithm=serialization.BestAvailableEncryption(passwd.encode())
			))