certificate_provision/copy_cert.py

#!/usr/bin/python
import paramiko
import argparse
import os
import posixpath
from subprocess import list2cmdline


def connection_from_config(name):
    """Create a connection from a Host entry in the .ssh/config file"""
    config = paramiko.config.SSHConfig()
    cp = os.path.expanduser('~/.ssh/config')
    with open(cp) as file:
        config.parse(file)
    info = config.lookup(name)
    info['username'] = info.pop('user')
    info['key_filename'] = info.pop('identityfile')
    info['port'] = int(info['port'])
    client = paramiko.client.SSHClient()
    client.load_host_keys(os.path.expanduser('~/.ssh/known_hosts'))
    client.connect(**info)
    return client


def scp(connection: paramiko.sftp_client.SFTPClient, src, dst):
    """Function to copy files from src to dst using connection"""
    connection.put(src, dst)


def exec_remote(connection: paramiko.SSHClient, command):
    """Function to execute specified command on conection"""
    if not isinstance(command, str):
        command = list2cmdline(command)
    stdin, stdout, stderr = connection.exec_command(command)
    return stdout, stderr


if __name__ == "__main__":
    parser = argparse.ArgumentParser(prefix_chars="*")
    parser.add_argument(
        'cert_path', help="location of the certificate to install")
    parser.add_argument(
        'key_path', help="localtion of the private key to install")
    parser.add_argument(
        'store_path',
        help="directory to store the private key and certificate")
    parser.add_argument('*d', '**dry', action='store_true')
    args = parser.parse_args()
    # Connect to host
    router = connection_from_config('router')
    stdout, stderr = exec_remote(router, ['ls', args.store_path])
    # Check to see if remote directory exists
    if stdout.channel.recv_exit_status() == 0:
        files = list(filter(bool, stdout.read().decode().split('\n')))
        # Delete old certificate and key
        for file in files:
            if file in ('cert.pem', 'key.pem'):
                if args.dry:
                    print("Removing {}".format(
                        posixpath.join(args.store_path, file)))
                else:
                    exec_remote(
                        router, ['rm', posixpath.join(
                            args.store_path, file)]
                    )
    else:
        # Create missing directory
        if args.dry:
            print("Creating missing dir {}".format(args.store_path))
        else:
            exec_remote(router, ['mkdir', args.store_path])
    sftp = router.open_sftp()
    # Copy certificate and key to store_path
    for src, dst in zip(
            [args.cert_path, args.key_path],
            [posixpath.join(args.store_path, "cert.pem"),
             posixpath.join(args.store_path, "key.pem")]
    ):
        if args.dry:
            print("{} -> {}".format(src, dst))
        else:
            scp(sftp, src, dst)
    sftp.close()
    router.close()