#!/usr/bin/env python3
# MIT License
#
# Copyright (c) 2026 Sławomir Koszewski
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

"""simple-ca — minimal CA for cloud-router IKEv2 PKI.

Directory layout:
  /etc/swanctl/x509ca/ca.pem              CA certificate (strongSwan trust anchor)
  /etc/swanctl/private/ca.key             CA private key
  /etc/swanctl/x509/server.pem            server certificate
  /etc/swanctl/private/server.key         server private key
  /etc/cloud-router/pki/{name}_cert.pem   user/road-warrior certificate
  /etc/cloud-router/pki/{name}_key.pem    user/road-warrior private key
  /etc/cloud-router/pki/{name}.pfx        PKCS#12 bundle for client distribution
"""

import argparse
import subprocess
import sys
from pathlib import Path

CA_CERT     = Path('/etc/swanctl/x509ca/ca.pem')
CA_KEY      = Path('/etc/swanctl/private/ca.key')
SERVER_CERT = Path('/etc/swanctl/x509/server.pem')
SERVER_KEY  = Path('/etc/swanctl/private/server.key')
PKI_DIR     = Path('/etc/cloud-router/pki')


def _run(*args, stdin=None):
    result = subprocess.run(list(args), input=stdin, capture_output=True)
    if result.returncode != 0:
        sys.stderr.buffer.write(result.stderr)
        sys.exit(1)
    return result.stdout


def _san_prefix(s):
    if '@' in s:
        return f'email:{s}'
    parts = s.split('.')
    if len(parts) == 4 and all(p.isdigit() and 0 <= int(p) <= 255 for p in parts):
        return f'IP:{s}'
    return f'DNS:{s}'


def cmd_make_ca(args):
    if CA_CERT.exists() and CA_KEY.exists():
        print('Root CA already exists, skipping.')
        return
    print('Generating root CA certificate and key...')
    _run(
        'openssl', 'req',
        '-x509',
        '-newkey', 'rsa:4096',
        '-keyout', str(CA_KEY),
        '-out', str(CA_CERT),
        '-days', str(args.days),
        '-noenc',
        '-subj', f'/CN={args.name}',
        '-text',
        '-addext', 'basicConstraints=critical,CA:TRUE,pathlen:1',
        '-addext', 'keyUsage=critical,keyCertSign,cRLSign',
    )
    CA_KEY.chmod(0o600)
    print(f'Root CA created.')
    print(f'  Certificate : {CA_CERT}')
    print(f'  Private key : {CA_KEY}')


def cmd_make_cert(args):
    name = args.subject.split('.')[0]
    if args.type == 'server':
        cert_path = SERVER_CERT
        key_path  = SERVER_KEY
    else:
        cert_path = PKI_DIR / f'{name}_cert.pem'
        key_path  = PKI_DIR / f'{name}_key.pem'

    if cert_path.exists() and key_path.exists():
        print(f'Certificate already exists at {cert_path}, skipping.')
        return

    if not CA_CERT.exists() or not CA_KEY.exists():
        print('ERROR: Root CA not found. Run: simple-ca make-ca <name>', file=sys.stderr)
        sys.exit(1)

    req_args = [
        'openssl', 'req',
        '-newkey', 'rsa:4096',
        '-keyout', str(key_path),
        '-noenc',
        '-subj', f'/CN={args.subject}',
        '-addext', 'basicConstraints=critical,CA:FALSE',
    ]

    if args.type == 'server':
        sans = [f'DNS:{args.subject}'] + [_san_prefix(s) for s in args.san]
        req_args += [
            '-addext', 'keyUsage=critical,digitalSignature,keyEncipherment',
            '-addext', 'extendedKeyUsage=serverAuth,clientAuth',
            '-addext', f'subjectAltName={",".join(sans)}',
        ]
    else:
        sans = [_san_prefix(s) for s in args.san]
        req_args += [
            '-addext', 'keyUsage=critical,digitalSignature,nonRepudiation',
            '-addext', 'extendedKeyUsage=clientAuth,emailProtection',
        ]
        if sans:
            req_args += ['-addext', f'subjectAltName={",".join(sans)}']

    x509_args = [
        'openssl', 'x509',
        '-req',
        '-CA', str(CA_CERT),
        '-CAkey', str(CA_KEY),
        '-copy_extensions', 'copyall',
        '-days', str(args.days),
        '-text',
        '-out', str(cert_path),
    ]

    csr = _run(*req_args)
    _run(*x509_args, stdin=csr)
    key_path.chmod(0o600)
    print('Certificate created.')
    print(f'  Certificate : {cert_path}')
    print(f'  Private key : {key_path}')


def cmd_make_pfx(args):
    name      = args.name
    cert_path = PKI_DIR / f'{name}_cert.pem'
    key_path  = PKI_DIR / f'{name}_key.pem'
    pfx_path  = PKI_DIR / f'{name}.pfx'

    if not cert_path.exists() or not key_path.exists():
        print(f'ERROR: Certificate or key not found for {name!r} in {PKI_DIR}.', file=sys.stderr)
        sys.exit(1)

    if pfx_path.exists():
        print(f'ERROR: {pfx_path} already exists.', file=sys.stderr)
        sys.exit(1)

    password = args.password or 'changeit'
    _run(
        'openssl', 'pkcs12',
        '-export',
        '-out', str(pfx_path),
        '-inkey', str(key_path),
        '-in', str(cert_path),
        '-certfile', str(CA_CERT),
        '-password', f'pass:{password}',
    )
    print(f'PKCS#12 created: {pfx_path}')


def main():
    ap = argparse.ArgumentParser(
        prog='simple-ca',
        description='Minimal CA for cloud-router IKEv2 PKI.',
    )
    sub = ap.add_subparsers(dest='command', required=True)

    p_ca = sub.add_parser('make-ca', help='Create root CA')
    p_ca.add_argument('name', help='CA common name')
    p_ca.add_argument('--days', type=int, default=3650, metavar='N')
    p_ca.set_defaults(func=cmd_make_ca)

    p_cert = sub.add_parser('make-cert', help='Issue a certificate')
    p_cert.add_argument('subject', help='Subject CN (FQDN for server, username for user)')
    p_cert.add_argument('san', nargs='*', help='Additional SANs: IP, DNS, or email')
    p_cert.add_argument('--type', choices=['server', 'user'], default='server',
                        help='Certificate type (default: server)')
    p_cert.add_argument('--days', type=int, default=365, metavar='N')
    p_cert.set_defaults(func=cmd_make_cert)

    p_pfx = sub.add_parser('make-pfx', help='Export PKCS#12 bundle for a user certificate')
    p_pfx.add_argument('name', help='Certificate name (without extension)')
    p_pfx.add_argument('--password', metavar='PASS',
                       help='Export password (default: changeit)')
    p_pfx.set_defaults(func=cmd_make_pfx)

    args = ap.parse_args()
    args.func(args)


if __name__ == '__main__':
    main()
