#-*- coding: utf-8 -*-
# (c) 2015 Anders Andersen
# See http://www.cs.uit.no/~aa/dist/tools/py/COPYING for details

from sys import argv
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.hazmat.primitives import serialization

# argv[1]: RSA key file name (creates)
# argv[2]: password (optional)
if len(argv) > 1:
    passwd = None
    if len(argv) > 2:
        passwd = argv[2]

    # Generate key and extract public part
    key = rsa.generate_private_key(
        public_exponent=65537,
        key_size=2048,
        backend=default_backend())
    pub = key.public_key()

    # Password encrypted private key or not
    if passwd:
        alg = serialization.BestAvailableEncryption(passwd.encode())
    else:
        alg = serialization.NoEncryption()

    # Write private key and public key to file
    fkey = open(argv[1], "wb")
    fpub = open(argv[1] + ".pub", "wb")
    fkey.write(
        key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.PKCS8,
            encryption_algorithm=alg))
    fpub.write(
        pub.public_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PublicFormat.SubjectPublicKeyInfo))
    fkey.close()
    fpub.close()