#!/usr/bin/env python3
# ==========================================================================
#         ____            _                     _____           _
#        / ___| _   _ ___| |_ ___ _ __ ___     |_   _|__   ___ | |___
#        \___ \| | | / __| __/ _ \ '_ ` _ \ _____| |/ _ \ / _ \| / __|
#         ___) | |_| \__ \ ||  __/ | | | | |_____| | (_) | (_) | \__ \
#        |____/ \__, |___/\__\___|_| |_| |_|     |_|\___/ \___/|_|___/
#               |___/
#                             --- System-Tools ---
#                  https://www.nntb.no/~dreibh/system-tools/
# ==========================================================================
#
# X.509 CA and Certificate Test Setup Script
# Copyright (C) 2015-2026 by Thomas Dreibholz
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
# Contact: thomas.dreibholz@gmail.com

import argparse
import sys

if sys.version_info < (3, 10):
   sys.stderr.write('ERROR: ' + sys.argv[0] + ' requires Python 3.10 or later!\n')
   sys.exit(1)

from typing            import Any, Final
from CertificateHelper import *



# ###### Main program #######################################################

# ====== Handle command-line parameters =====================================
subjectAltName   : str | None      = None
certType         : CertificateType = CertificateType.Server
revokeIfExisting : bool            = True
revoke           : bool            = False

# ====== Handle arguments ===================================================
parser : Final[argparse.ArgumentParser] = argparse.ArgumentParser(
   description     = 'X.509 CA and Certificate Test Setup Script',
   formatter_class = argparse.ArgumentDefaultsHelpFormatter
)

# Positional arguments:
parser.add_argument('ssl_directory',      help = 'The main directory for SSL certificates')
parser.add_argument('names', nargs = '+', help='One or more names for the certificates')

# Certificate Type (mutually exclusive):
type_group = parser.add_mutually_exclusive_group()
type_group.add_argument('-s', '--server', action = 'store_const', dest = 'cert_type',
                        const=CertificateType.Server, help = 'Create a server certificate')
type_group.add_argument('-c', '--client', action = 'store_const', dest = 'cert_type',
                        const=CertificateType.Client, help = 'Create a client certificate')
type_group.add_argument('-u', '--user',   action = 'store_const', dest = 'cert_type',
                        const=CertificateType.User,   help = 'Create a user certificate')
parser.set_defaults(cert_type= CertificateType.Server)

# Existing Certificate Behavior (mutually exclusive):
existing_group = parser.add_mutually_exclusive_group()
existing_group.add_argument('-x', '--revoke-if-existing', action = 'store_true', default = True,
                              help = 'Revoke the certificate if it already exists')
existing_group.add_argument('-i', '--ignore-if-existing', action = 'store_false',
                              dest = 'revoke_if_existing',
                              help = 'Ignore and do not revoke if the certificate exists')

# Other Options:
parser.add_argument('-r', '--revoke', action = 'store_true',
                     help = 'Revoke the generated certificate(s) immediately')
parser.add_argument('-S', '--san', '--subjectAltName', dest = 'subjectAltName', default=None,
                     help = 'Subject Alternative Name (subjectAltName|LOCAL|LOOKUP)')

parser.add_argument('-A', '--algorithm', type = str, default = 'RSA',
                    choices = [ 'RSA', 'EC', 'ED25519', 'ED448' ],
                    help = 'The key algorithm to use')
parser.add_argument('-C', '--curve', type=str, default=DefaultECCurve,
                    help = 'The elliptic curve to use (only applicable if algorithm is EC)')
parser.add_argument('-K', '--keylen-ca', type = int, default = DefaultCAKeyLength,
                     help = 'Key length for the CA (only applicable if algorithm is RSA)')
parser.add_argument('-k', '--keylen-cert', type = int, default = DefaultCertKeyLength,
                     help = 'Key length for the certificates (only applicable if algorithm is RSA)')

args = parser.parse_args()

keyAlgorithm         : Final[KeyAlgorithm]  = KeyAlgorithm[args.algorithm]
caKeyLength          : Final[int]           = args.keylen_ca
certKeyLength        : Final[int]           = args.keylen_cert
ecCurve              : Final[str]           = args.curve
mainDirectory        : Final[str]           = args.ssl_directory
certificatesToCreate : list[dict[str, Any]] = [ ]
name                 : str
for name in args.names:
   parsedName : str
   san        : str
   ( parsedName,  san ) = prepareSubjectAltName(args.cert_type, name, args.subjectAltName)
   certificatesToCreate.append({ 'name'             : parsedName,
                                 'type'             : args.cert_type,
                                 'subjectAltName'   : san,
                                 'revokeIfExisting' : args.revoke_if_existing })
# print(certificatesToCreate)


# ====== Create CA hierarchy, if not existing ===============================

# ===========================================================================
# Hierarchy to be created:
# TestLevel1
#    - TestLevel2
#       - TestIntermediate1
#          - TestIntermediate2
#             - TestLeaf
#                - Servers ...
#                - Clients ...
#                - Users ...
# ===========================================================================


# ====== Create Test CAs ====================================================
globalCRLFileName : Final[str] = "TestGlobal.crl"

# Create Test Root Level-1 CA certificate (self-signed):
TestLevel1 : Final[CA] = \
   CA(mainDirectory, 'TestLevel1',
      parentCA  = None,
      subject   = '/C=NO/ST=Oslo/L=Oslo/O=Simula Research Laboratory/streetAddress=Kristian Augusts gate 23/postalCode=0164/CN=Test Level-1 CA Certificate',
      certType  = CertificateType.RootCA,
      keyAlgorithm = keyAlgorithm, keyLength = caKeyLength, ecCurve = ecCurve,
      globalCRLFileName = globalCRLFileName)

# Create Test Root Level-2 CA certificate (signed by Test Root Level-1):
TestLevel2 : Final[CA] = \
   CA(mainDirectory, 'TestLevel2',
      parentCA  = TestLevel1,
      subject   = '/C=NO/ST=Oslo/L=Oslo/O=Simula Research Laboratory/streetAddress=Kristian Augusts gate 23/postalCode=0164/CN=Test Level-2 CA Certificate',
      certType  = CertificateType.IntermediateCA,
      keyAlgorithm = keyAlgorithm, keyLength = caKeyLength, ecCurve = ecCurve,
      globalCRLFileName = globalCRLFileName)

# Create Test Intermediate CA certificate (signed by Test Root Level-2):
TestIntermediate1 : Final[CA] = \
   CA(mainDirectory, 'TestIntermediate1',
      parentCA  = TestLevel2,
      subject   = '/C=NO/ST=Oslo/L=Oslo/O=Simula Metropolitan Centre for Digital Engineering (SimulaMet)/streetAddress=Stensberggata 27/postalCode=0170/CN=*',
      certType  = CertificateType.IntermediateCA,
      keyAlgorithm = keyAlgorithm, keyLength = caKeyLength, ecCurve = ecCurve,
      globalCRLFileName = globalCRLFileName)

# Create Test Intermediate CA certificate (signed by Test Root Level-2):
TestIntermediate2 : Final[CA] = \
   CA(mainDirectory, 'TestIntermediate2',
      parentCA  = TestIntermediate1,
      subject   = '/C=NO/ST=Oslo/L=Oslo/O=Simula Metropolitan Centre for Digital Engineering (SimulaMet)/OU=Centre for Resilient Networks and Applications (CRNA)/streetAddress=Stensberggata 27/postalCode=0170/CN=*',
      certType  = CertificateType.IntermediateCA,
      keyAlgorithm = keyAlgorithm, keyLength = caKeyLength, ecCurve = ecCurve,
      globalCRLFileName = globalCRLFileName)

# Create Test Leaf CA certificate (signed by Test Intermediate):
TestLeaf : Final[CA] = \
   CA(mainDirectory, 'TestLeaf',
      parentCA  = TestIntermediate2,
      subject   = '/C=NO/ST=Oslo/L=Oslo/O=Simula Metropolitan Centre for Digital Engineering (SimulaMet)/OU=SimulaMet Interoperability Lab (SMIL)/streetAddress=Stensberggata 27/postalCode=0170/CN=*/',
      certType  = CertificateType.LeafCA,
      keyAlgorithm = keyAlgorithm, keyLength = caKeyLength, ecCurve = ecCurve,
      globalCRLFileName = globalCRLFileName)


# ====== Create Test Servers ================================================
defaultSubjectWithoutCN = '/C=NO/ST=Oslo/L=Oslo/O=Simula Metropolitan Centre for Digital Engineering (SimulaMet)/OU=SimulaMet Interoperability Lab (SMIL)/streetAddress=Stensberggata 27/postalCode=0170'

certificates        : dict[str, Certificate] = {}
certificateToCreate : dict[str, Any]
for certificateToCreate in certificatesToCreate:
   # print(certificateToCreate)

   certificateName = certificateToCreate['name']
   subjectWithoutCN : str
   if certificateToCreate['type'] != CertificateType.User:
      subjectWithoutCN = defaultSubjectWithoutCN
   else:
      subjectWithoutCN = defaultSubjectWithoutCN + prepareUserSubject(certificateName)
   certificates[certificateName] = \
      Certificate(mainDirectory, certificateName, TestLeaf,
                  subjectWithoutCN, certificateToCreate['subjectAltName'],
                  certificateToCreate['type'],
                  keyAlgorithm = keyAlgorithm, keyLength = certKeyLength, ecCurve = ecCurve,
                  revokeIfExisting = revokeIfExisting)

   if revoke:
      certificates[certificateName].revoke()


# ====== Finally, force regeneration of the CRL to ensure it is up-to-date ==
TestLevel1.generateCRL()
