aman-sys/aman/tools/KeyPairCreator.py

#!/usr/bin/env python

import argparse
from datetime import datetime
import os
import sys
from typing import Tuple
import zmq.auth

# @brief Creates a new keypair for ZMQ encryption
# @param[in] directory The location where to store the keys
# @return The public and private key tuple
def KeyPairCreator(directory: str, server: bool) -> Tuple[str, str]:
    if not server:
        target = 'client'
    else:
        target = 'server'

    public, private = zmq.auth.create_certificates(directory, target)
    return public, private

def str2bool(value):
    if isinstance(value, bool):
        return value
    elif value.lower() in ('yes', 'true', 't', 'y', '1'):
        return True
    elif value.lower() in ('no', 'false', 'f', 'n', '0'):
        return False
    else:
        raise argparse.ArgumentTypeError('Boolean value expected')

def findIdentificationKey(path, publicKey : bool):
    if True == publicKey:
        identifier = 'public-key = '
    else:
        identifier = 'secret-key = '

    with open(path) as file:
        key = ''

        for line in file:
            if identifier in line:
                elements = line.split('=')
                for idx in range(1, len(elements)):
                    if 0 == len(key):
                        key = elements[idx][2:-1]
                        key = key + elements[idx][-1]
                    else:
                        key = key + '=' + elements[idx]

                return key[0:-2]

    return None

if __name__ == '__main__':
    # create the commandline parser
    parser = argparse.ArgumentParser(description='Create a new key-value pair')
    parser.add_argument('--directory', type=str, help='Directory where to store the key pair')
    parser.add_argument('--publickey', nargs='?', type=str, default=os.getcwd(), help='Full path to the public key of the server')
    parser.add_argument('--server', default=False, action='store_true', help="Creates server key pair")
    args = parser.parse_args()

    # validate the arguments
    if False == args.server and not os.path.exists(args.publickey):
        sys.stderr.write('The public key of the server cannot be found')
        sys.exit(-1)

    # create the directory if it does not exist
    if not os.path.exists(args.directory):
        os.makedirs(args.directory)

    # create the keys
    _, private = KeyPairCreator(args.directory, args.server)

    if False == args.server:
        publicServer  = findIdentificationKey(args.publickey, True)
        publicClient  = findIdentificationKey(private, True)
        privateClient = findIdentificationKey(private, False)

        if None == publicServer:
            sys.stderr.write('The public key of the server cannot be found in the defined file')
            sys.exit(-1)
        if None == publicClient:
            sys.stderr.write('Unable to extract the created public key')
            sys.exit(-1)
        if None == privateClient:
            sys.stderr.write('Unable to extract the created private key')
            sys.exit(-1)

        # rename keys
        timestamp = str(datetime.now(tz=None))
        timestamp = timestamp.replace(' ', '_')
        timestamp = timestamp.replace(':', '-')
        os.rename(os.path.join(args.directory, 'client.key'), os.path.join(args.directory, timestamp + '.key'))
        os.rename(os.path.join(args.directory, 'client.key_secret'), os.path.join(args.directory, timestamp + '.key_secret'))

        print(publicServer)
        print(publicClient)
        print(privateClient)