#! /usr/bin/env python3
# (C) Copyright 2021,2023 Hewlett Packard Enterprise Development LP
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
    importcert <SSL_service> [-f] <service_cert> [<CA_bundle>]
    importcert <SSL_service> [-f] -ca <CA_bundle>

OPTIONS
    -f              Import a certificate without prompting the user.
    -ca <CA_bundle> Allows the import of a CA bundle without importing
                    a service certificate.

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

longhelp='''importcert - imports a signed certificate and supporting certificate authorities
(CAs) for the storage system SSL services.

SYNTAX
    importcert <SSL_service> [-f] <service_cert> [<CA_bundle>]
    importcert <SSL_service> [-f] -ca <CA_bundle>

DESCRIPTION
    The importcert command allows a user to import certificates for a given
    service. The user can import a CA bundle containing the intermediate and/or
    root CAs prior to importing the service certificate. The CA bundle can also
    be imported alongside the service certificate.

AUTHORITY
    Super

OPTIONS
    -f
        Import a certificate without prompting the user.

    -ca <CA_bundle>
        Allows the import of a CA bundle without importing a service
        certificate. Note the filename "stdin" can be used to paste the
        CA bundle into the CLI.

SPECIFIERS
    <SSL_service>
        Valid service names are qw-client and qw-server.

NOTES
    Note that the qw-server service is restarted when a self-signed certificate
    is generated.

    Note the filename "stdin" can be used to paste the CA bundle and or service
    certificate into the CLI.

    Use the createcert command to create a CSR and use the showcert command to
    display the certificates.

EXAMPLES
    The following example shows how to import a signed service certificate with
    the supporting CA for the qw-server service.

    $ importcert qw-server qw-server-service.pem ca.pem

    The following example shows how to import just the supporting CAs for the
    qw-server service without importing the service certificate itself.

    $ importcert qw-server -ca ca-bundle.pem

    Now that the CA bundle has been imported, the service certificate can be
    imported:

    $ importcert qw-server qw-server-service.pem
'''

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

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

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_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', nargs='?')
p.add_argument('-f', action='store_true')
p.add_argument('-ca')
args,extra=p.parse_known_args()

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

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

if not args.service in all_services:
    exit_usage('argument SSL_service: invalid choice: %s (choose from %s)' % (args.service, ", ".join(all_services)) )

if args.ca:
    if len(extra):
        exit_usage('extra arguments when using -ca option: %s' % (" ".join(extra)))
else:
    if not len(extra):  exit_usage('insufficient arguments')
    if len(extra)>2:    exit_usage('extra arguments: %s' % (" ".join(extra[2:])))

args.cert=None
if len(extra):  args.cert=extra.pop(0)
if len(extra):  args.ca=extra.pop(0)

if args.cert=='stdin':
    print('Please paste the Certificate for %s. Once finished, please press Enter twice.' % (args.service))
    cert_pem=''
    while True:
        try:
            line=input()
        except EOFError:
            break
        if line:
            cert_pem+=line+'\n'
        else:
            break
elif args.cert:
    with open(args.cert, 'r') as file:
        cert_pem=file.read()

if args.ca=='stdin':
    print('Please paste the CA bundle for %s. Once finished, please press Enter twice.' % (args.service))
    ca_pem=''
    while True:
        try:
            line=input()
        except EOFError:
            break
        if line:
            ca_pem+=line+'\n'
        else:
            break
elif args.ca:
    with open(args.ca, 'r') as file:
        ca_pem=file.read()

if not args.f:
    print('Do you want to import these certificates(s) for %s service?' % (args.service))
    if args.cert:    print('*', args.cert, 'as the service certificate?')
    if args.ca:      print('*', args.ca, 'certificate authorities?')
    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 importing signed certificate(s) (yes/no)? ')
    while True:
        if reply=='yes':   break
        elif reply=='no':  sys.exit(0)
        else:              reply=input('Please type "yes" or "no": ')

new_cert=False
new_cas=False

if args.service=='qw-client':
    if args.cert:   # The qw-client service only allows import of intca and rootca types.
        exit_usage('a Certificate Signing Request must be created before a certificate can be imported')
    if not args.ca:
        exit_usage('insufficient arguments');

if args.service=='qw-server' and args.cert:
    if not os.access(qw_server_csr_pem, os.R_OK):   # Confirm csr file exists for qw-server.
        exit_usage('a Certificate Signing Request must be created before a certificate can be imported')

    if re.search(r'BEGIN TRUSTED CERTIFICATE', cert_pem):   # Check for import of trusted certs.
        exit_err('service certificate cannot be a trusted cert')

    # Fill a pipe with the certificate pem content.
    cert_pipe_r,cert_pipe_w=os.pipe()   # Pipe for copying data to stdin.
    pipe=os.fdopen(cert_pipe_w, 'w')
    pipe.write(cert_pem)
    pipe.close()

    # Confirm that modulus of incoming cert matches the existing csr and private key.
    p1=subprocess.run(['openssl', 'x509', '-modulus', '-noout'],
        stdin=cert_pipe_r, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True)
    if p1.returncode and p1.stderr:  exit_err(p1.stderr)
    if not re.match(r"Modulus=", p1.stdout):   # Exit if not able to get modulus.
        exit_err('unable to parse the presented certificate')
    p2=subprocess.run(['openssl', 'req', '-modulus', '-noout', '-in', qw_server_csr_pem],
        stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True)
    if p2.returncode and p2.stderr:  exit_err(p2.stderr)
    if not re.match(r"Modulus=", p2.stdout):   # Exit if not able to get modulus.
        exit_err('unable to parse the CSR')
    # Note that the first block in the qw_server_pem file is (should be) the private key.
    p3=subprocess.run(['openssl', 'rsa', '-modulus', '-noout', '-in', qw_server_pem],
        stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True)
    if p3.returncode and p3.stderr:
        if re.search(r'unable to load Private Key', p3.stderr):
            exit_err('Private Key not found, try creating a new csr with createcert')
        else:
            exit_err(p3.stderr)
    if not re.match(r"Modulus=", p3.stdout):   # Exit if not able to get modulus.
        exit_err('unable to parse the keyfile')

    if p1.stdout!=p2.stdout:   # Compare modulus of incoming cetificate and existing csr.
        exit_err('the presented certificate does not match the previously stored CSR')
    if p1.stdout!=p3.stdout:   # Compare modulus of incoming certificate and existing key.
        exit_err('the presented certificate does not match the previously stored key')

    # Fill a pipe with the certificate pem content.
    cert_pipe_r,cert_pipe_w=os.pipe()   # Pipe for copying data to stdin.
    pipe=os.fdopen(cert_pipe_w, 'w')
    pipe.write(cert_pem)
    pipe.close()

    # Confirm that the incoming cert has the Server attribute.
    p1=subprocess.run(['openssl', 'x509', '-text', '-noout',],
        stdin=cert_pipe_r, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True)
    if p1.returncode and p1.stderr:  exit_err(p1.stderr)
    sr=re.search(r"TLS Web Server Authentication", p1.stdout)
    if not sr:
        exit_err('the certificate does not contain the Server attribute')
    new_cert=True


if args.ca:
    cas_dict={}
    cas_list=[]
    n=1;

    if re.search(r'BEGIN TRUSTED CERTIFICATE', ca_pem):   # Check for import of trusted certs.
        exit_err('service certificate cannot be a trusted cert')

    # Iterate over each certificate in the certificate bundle.
    for sr in re.finditer(r'(?sm)(-----BEGIN CERTIFICATE-----.*?-----END CERTIFICATE-----(\n)?)', ca_pem):
        if sr:
            curr_cert=sr.group(1)
        else:
            # Exit with error if no certificates were found in the bundle.
            if n==1:    exit_err('Unable to parse the presented CA bundle.')
            # Otherwise break out of the loop if there are no more certificates to process.
            break

        pipe_r,pipe_w=os.pipe()   # Pipe for copying data to stdin.
        pipe=os.fdopen(pipe_w, 'w')
        pipe.write(curr_cert)
        pipe.close()
        # This process uses openssl to format the certificate in human readable format.
        p2=subprocess.run(['openssl', 'x509', '-noout', '-text'], stdin=pipe_r,
            stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True)
        if p2.returncode and p2.stderr:  exit_err(p2.stderr)

        sri=re.search(r"Issuer: (.*)", p2.stdout).group(1)
        srs=re.search(r"Subject: (.*)", p2.stdout).group(1)
        sr=re.search(r"CA:(TRUE|FALSE)", p2.stdout)
        if not sr or sr.group(1)=='FALSE':   # If there is no CA attribute, or it is FALSE, then type is cert.
            exit_err('service certificates not allowed in import CA bundle')
        else:
            if sri==srs:
                cert_type='rootca'
            else:
                cert_type='intca'
        if (args.service=='qw-client') and (cert_type!='rootca'):
            exit_err('the qw-client certificate must be a trusted root Certificate Authority; other certificate types cannot be imported')

        if not curr_cert in cas_dict:    # Detect and flatten duplicate CAs in the incoming CA bundle.
            cas_dict[curr_cert]=cert_type
            cas_list.append(curr_cert)     # Preserve order of initial unique cert.

        n+=1

    if len(cas_list):  new_cas=True


if args.service=='qw-client' and new_cas:    # The qw-client service never has new service cert to process, only CAs.
    # For qw-client service, overwrite the cert content to the qw_client_pem file location.
    with open(qw_client_pem, 'w+') as file:
        for ca in cas_list:
            file.write(ca)
    restart=True

if args.service=='qw-server' and (new_cert or new_cas):
    old_key_pem=''
    old_cert_pem=''
    old_cas_pem=''
    # May need to bring forward some existing service certs, and private key, from the server cert.pem file.
    if os.path.isfile(qw_server_pem):
        with open(qw_server_pem, 'r') as file:
            server_pem=file.read()         # Read in the entire file contents.

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

        # Find the first certificate, and check if it is a non-CA certificate.
        sc=re.search(r'(?sm)(-----BEGIN CERTIFICATE-----.*?-----END CERTIFICATE-----(\n)?)', server_pem)
        if sc:
            pipe_r,pipe_w=os.pipe()   # Pipe for copying data to stdin.
            pipe=os.fdopen(pipe_w, 'w')
            pipe.write(sc.group(1))
            pipe.close()
            # This process uses openssl to format the certificate in human readable format.
            p1=subprocess.run(['openssl', 'x509', '-noout', '-text'], stdin=pipe_r,
                stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True)
            if p1.returncode and p1.stderr:  exit_err(p1.stderr)
            sr=re.search(r"CA:(TRUE|FALSE)", p1.stdout)
            if not sr or sr.group(1)=='FALSE':   # If there is no CA attribute, or it is FALSE, then type is cert.
                # The first certificate is a non-CA cert, so keep track of it separately.
                old_cert_pem=sc.group(1)
                old_cas_pem=server_pem[sc.span(1)[1]:]
            else:
                old_cas_pem=sc.group(1)+server_pem[sc.span(1)[1]:]

    # Reassemble and write the new content to the qw_server_pem file location.
    with open(qw_server_pem, 'w+') as file:
        file.write(old_key_pem+'\n')               # Save back the old private key, if any.
        if new_cert:
            file.write(cert_pem+'\n')              # If there is a new signed cert to import, save it.
        else:
            file.write(old_cert_pem+'\n')          # Otherwise save back the old signed cert, if any.
        if new_cas:
            file.write("\n".join(cas_list)+'\n')   # If there are new CA certs to import, save them.
        else:
            file.write(old_cas_pem+'\n')           # Otherwise save back the old CA certs, if any.

    if new_cert:
        os.unlink(qw_server_csr_pem)    # Since we have written a new service cert, the csr pem should now be deleted.
    restart=True

if restart:
    # 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)
