#! /usr/bin/env python3
# (C) Copyright 2021,2023 Hewlett Packard Enterprise Development LP
# Author: Cole Hlava
# removecert
import sys, subprocess, os.path, argparse, re

if sys.version_info < (3,6):
    exit(sys.argv[0] + ": error: Python version below 3.6 not supported.")

usage = '''

SYNTAX
    removecert {all|<SSL_service_name>} [-type <typename>]

OPTIONS
    -f            Skips the user prompt
    -type         Controls the types of certificates removed

Try "removecert -h" for detailed information for options
 
'''

longhelp = '''removecert - Removes SSL certificates from the storage system.

SYNTAX
    removecert {all|<SSL_service_name>} [-type <typename>]

DESCRIPTION
    The removecert command is used to remove certificates that are no longer
    trusted. In most cases it is better to overwrite the offending certificate
    with importcert. The user specifies which service to have its certificates
    removed. The removal can be limited to a specific type.

AUTHORITY
    Super

OPTIONS
    -f
        Skips the prompt warning the user of which certificates will be removed
        and which services will be restarted.

    -type <typename>
        Allows the user to limit the removal to a specific type.
        Valid types are csr, cert, intca, and rootca.

SPECIFIERS
    <SSL_service>
        Valid service names are qw-client and qw-server.
        The user may also specify all, which will remove certificates for all
        services.

EXAMPLES
    The following example shows how to remove all certificates for the qw-client.

        # removecert qw-client

    The following example shows how to remove just the root Certificate
    Authority for the qw-client.

        # removecert qw-client -type rootca

'''

def exit_usage(msg):
    p.print_usage()
    exit_err(msg)

def exit_err(msg):
    sys.exit(os.path.basename(sys.argv[0]) + ' error: ' + msg)

def confirmRemoval(certList):
    if certList:
        print('The following certificates will be removed:')
    t_maxLength = 3
    for cert in certList:
        if len(cert[0]) > t_maxLength:
            t_maxLength = len(cert[0])
    print('Service   ', end='')
    print('Type'.ljust(t_maxLength + 1), end='')
    print('Fingerprint')
    for cert in certList:
        if cert[0] != 'key':    # Don't output the 'key' type of certificate, since that is internal only.
            print(args.service + ' ', end='')
            print(cert[0].ljust(t_maxLength + 1), end='')
            print(cert[1])

    if certList:
        print('Also the following services will be restarted if currently running:')
        print('  qw-client: manages communications with Quorum Witness clients')
        print('  qw-server: Quorum Witness server application')

    reply=input('Continue removing certificate(s) (yes/no)? ')
    while True:
        if reply=='yes':   break
        elif reply=='no':  sys.exit(0)
        else:              reply=input('Please type "yes" or "no": ')


# Maps certificate -> (type, fingerprint, cert)
def decodeCertificates(pemFile):
    certsMap = []

    try:
        with open(pemFile, 'r') as file:
            qw_pem=file.read()         # Read in the entire file contents.
    except FileNotFoundError:
        return certsMap

    # Find the private key, if any.
    sr=re.search(r'(?sm)(-----BEGIN (RSA )?PRIVATE KEY-----.*?-----END (RSA )?PRIVATE KEY-----(\n)?)', qw_pem)
    if sr:  certsMap.append(('key', '--', sr.group(1)))

    # Iterate through each certificate from the pem file and map certificate to its type and fingerprint.
    for sc in re.finditer(r'(?sm)(-----BEGIN CERTIFICATE-----.*?-----END CERTIFICATE-----(\n)?)', qw_pem):
        cert=sc.group(1)
        pipe_r, pipe_w = os.pipe()   # Pipe for copying data to stdin.
        pipe = os.fdopen(pipe_w, 'w')
        pipe.write(cert)
        pipe.close()

        p1=subprocess.run(['openssl', 'x509', '-fingerprint', '-sha256', '-noout', '-text'], stdin=pipe_r,
            stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True)
        if p1.returncode and p1.stderr:  exit_err(p1.stderr)

        fingerprint=re.search(r"Fingerprint=(.*)", p1.stdout).group(1).replace(':', '').lower()[:40]
        sri=re.search(r"Issuer: (.*)", p1.stdout).group(1)
        srs=re.search(r"Subject: (.*)", p1.stdout).group(1)
        sr=re.search(r"CA:(TRUE|FALSE)", p1.stdout)

        serviceType = ''
        if not sr or sr.group(1) == 'FALSE':
            serviceType = 'cert'
        else:
            if sri == srs:
                serviceType = 'rootca'
            else:
                serviceType = 'intca'
        #print(f'Service type = {serviceType}')
        certsMap.append((serviceType, fingerprint, cert))
    return certsMap

def unlinkpem(filename):
    try:
        os.unlink(filename)
    except FileNotFoundError:  # Don't return an error if the file does not exist.
        pass


# File locations.
qw_server_csr_pem='/usr/local/etc/csr.pem'
qw_server_pem='/usr/local/etc/cert.pem'
qw_client_pem='/usr/local/etc/cacert.pem'

all_services=['qw-client', 'qw-server', 'all']
all_types=['csr', 'cert', 'intca', 'rootca']

p = argparse.ArgumentParser(usage=usage, add_help=False)
p.add_argument('-h', '--help', action='store_true')
p.add_argument('service', metavar='SSL_service_name', nargs='?')
p.add_argument('-type')
p.add_argument('-f', action='store_true')
args = p.parse_args()

if args.help:
    print(longhelp)
    sys.exit(0)

if not args.service:
    exit_usage('insufficient arguments')

if args.service not in all_services:
    exit_err(f'Invalid service(s): {args.service}')

if not args.type:    # If not set, then set default types list to all types.
    types=all_types
else:
    types=args.type.split(',')
    bad_choices = list(set(types) - set(all_types))        # Find invalid certificate types.
    if bad_choices:
        exit_usage('argument -type: invalid choice: %s (choose from %s)' % (bad_choices[0], ", ".join(all_types)) )


# Find qw-client certificates to remove.
clientCertsToBeRemoved = []
clientCertsToWriteBack = ''
if args.service in ['qw-client', 'all']:
    clientCerts = decodeCertificates(qw_client_pem)
    for item in clientCerts:
        if item[0] in types:    # If item type matches types arguments, then append to remove list.
            clientCertsToBeRemoved.append(item)
        else:                   # Otherwise item certificate should be written back to pem file.
            clientCertsToWriteBack += item[2] + '\n'

# Find qw-server certificates to remove.
serverCertsToBeRemoved = []
serverCertsToWriteBack = ''
removeCsr = False
if args.service in ['qw-server', 'all']:
    # Special case handling for csr type.
    if 'csr' in types:
        types.append('key')    # If 'csr' type is specified for removal, then also remove the 'key' type.

        # This process outputs the csr in der format in order to get the fingerprint in the next step.
        p1=subprocess.Popen(['openssl', 'req', '-outform', 'der', '-in', qw_server_csr_pem],
            stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        # This process uses openssl to get the fingerprint for the csr.
        p2=subprocess.run(['openssl', 'dgst', '-sha256',], stdin=p1.stdout,
            stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True)
        if p2.returncode and p2.stderr:  exit_err(p2.stderr)

        fingerprint=p2.stdout.split(" ")[1][:40]
        serverCertsToBeRemoved = [('csr', fingerprint, '--')]  # Add csr to list of certs to be removed.
        removeCsr = True    # Set flag to remove the csr file after user confirmation.

    serverCerts = decodeCertificates(qw_server_pem)
    for item in serverCerts:
        if item[0] in types:    # If item type matches types arguments, then append to remove list.
            serverCertsToBeRemoved.append(item)
        else:                   # Otherwise item certificate should be written back to pem file.
            serverCertsToWriteBack += item[2] + '\n'


if clientCertsToBeRemoved or serverCertsToBeRemoved:
    # Confirm removal if the -f flag was not specified.
    if not args.f:
        confirmRemoval(clientCertsToBeRemoved + serverCertsToBeRemoved)

    if len(clientCertsToBeRemoved):
        if clientCertsToWriteBack == '':    # Nothing to write back, so unlink the pem file.
            unlinkpem(qw_client_pem)
        else:                               # There is something to write back.
            with open(qw_client_pem, 'w') as f:
                f.write(clientCertsToWriteBack)

    if len(serverCertsToBeRemoved):
        if serverCertsToWriteBack == '':    # Nothing to write back, so unlink the pem file.
            unlinkpem(qw_server_pem)
        else:                               # There is something to write back.
            with open(qw_server_pem, 'w') as f:
                f.write(serverCertsToWriteBack)

    if removeCsr:
        unlinkpem(qw_server_csr_pem)

    print('Certificates removed.')
    # Restart qwserv.service, as needed.
    p1=subprocess.run(['systemctl', 'restart', 'qwserv.service'],
        stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True)
    if p1.returncode and p1.stderr:  exit_err(p1.stderr)
else:
    sys.exit('There are no certificates for services(s): ' + args.service)

