package main

import (
	"crypto"
	"crypto/ecdsa"
	"crypto/ed25519"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"flag"
	"fmt"
	"log"
	"math/big"
	"os"
	"path/filepath"
	"regexp"
	"time"
)

type caInfo struct {
	caDir  string
	caKey  crypto.Signer
	caCert *x509.Certificate
}

type DNSNames []string

func (dnsNames *DNSNames) String() string {
	return fmt.Sprint(*dnsNames)
}

func (dnsNames *DNSNames) Set(s string) error {
	*dnsNames = append(*dnsNames, s)
	return nil
}

var currentTime = time.Now()

func caInit(caDir string, certName string, notAfter time.Time, domainConstraint string) {
	caKey := generateKey(filepath.Join(caDir, "ca.key"))
	makeRootCert(caKey, filepath.Join(caDir, "ca.crt"), certName, notAfter, domainConstraint)
}

func getCa(caDir string) *caInfo {
	caKey := readCaKey(filepath.Join(caDir, "ca.key"))
	caCert := readCaCert(filepath.Join(caDir, "ca.crt"))

	return &caInfo{caDir, caKey, caCert}
}

func readCaKey(pemFile string) crypto.Signer {
	derBytes := readPem(pemFile, "PRIVATE KEY")
	privateKey, err := x509.ParsePKCS8PrivateKey(derBytes)
	fatalIfErr(err)

	return privateKey.(crypto.Signer)
}

func readCaCert(pemFile string) *x509.Certificate {
	derBytes := readPem(pemFile, "CERTIFICATE")
	caCert, err := x509.ParseCertificate(derBytes)
	fatalIfErr(err)

	return caCert
}

func readPem(pemFile, pemType string) []byte {
	pemData, err := os.ReadFile(pemFile)
	fatalIfErr(err)
	pemBlock, _ := pem.Decode(pemData)
	if pemBlock == nil {
		log.Fatal("unable to decode PEM")
	}
	if pemBlock.Type != pemType {
		log.Fatalf("incorrect PEM type, expected '%s'", pemType)
	}

	return pemBlock.Bytes
}

func writePem(pemFile string, derBytes []byte, pemType string) {
	f, err := os.OpenFile(pemFile, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0o660)
	fatalIfErr(err)
	defer f.Close()
	err = pem.Encode(f, &pem.Block{
		Type:  pemType,
		Bytes: derBytes,
	})
	fatalIfErr(err)
}

func generateKey(keyFile string) crypto.Signer {
	switch os.Getenv("CA_KEY_TYPE") {
	case "ECDSA":
		return generateEcDsaKey(keyFile)
	case "EdDSA":
		return generateEdDsaKey(keyFile)
	default:
		return generateEcDsaKey(keyFile)
	}
}

func generateEcDsaKey(keyFile string) crypto.Signer {
	privateKey, err := ecdsa.GenerateKey(elliptic.P384(), rand.Reader)
	fatalIfErr(err)
	derBytes, err := x509.MarshalPKCS8PrivateKey(privateKey)
	fatalIfErr(err)
	writePem(keyFile, derBytes, "PRIVATE KEY")

	return privateKey
}

func generateEdDsaKey(keyFile string) crypto.Signer {
	_, privateKey, err := ed25519.GenerateKey(rand.Reader)
	fatalIfErr(err)
	derBytes, err := x509.MarshalPKCS8PrivateKey(privateKey)
	fatalIfErr(err)
	writePem(keyFile, derBytes, "PRIVATE KEY")

	return privateKey
}

func makeRootCert(caKey crypto.Signer, certFile string, certName string, notAfter time.Time, domainConstraint string) (*x509.Certificate, error) {
	tpl := getCaTemplate(certName, notAfter)
	if domainConstraint != "" {
		tpl.PermittedDNSDomainsCritical = true
		tpl.PermittedDNSDomains = []string{domainConstraint}
	}
	derBytes, err := x509.CreateCertificate(rand.Reader, tpl, tpl, caKey.Public(), caKey)
	fatalIfErr(err)
	writePem(certFile, derBytes, "CERTIFICATE")

	return x509.ParseCertificate(derBytes)
}

func getCertSubject(commonName string, orgUnit string) pkix.Name {
	certSubject := pkix.Name{
		CommonName: commonName,
	}
	if orgUnit != "" {
		certSubject.OrganizationalUnit = []string{orgUnit}
	}

	return certSubject
}

func getCaTemplate(certName string, notAfter time.Time) *x509.Certificate {
	return &x509.Certificate{
		Subject:               getCertSubject(certName, ""),
		SerialNumber:          generateSerial(),
		NotBefore:             currentTime.Add(-5 * time.Minute),
		NotAfter:              notAfter,
		KeyUsage:              x509.KeyUsageCertSign | x509.KeyUsageCRLSign,
		BasicConstraintsValid: true,
		IsCA:                  true,
		MaxPathLenZero:        true,
	}
}

func getClientTemplate(certName string, orgUnit string, notAfter time.Time) *x509.Certificate {
	return &x509.Certificate{
		Subject:               getCertSubject(certName, orgUnit),
		SerialNumber:          generateSerial(),
		NotBefore:             currentTime.Add(-5 * time.Minute),
		NotAfter:              notAfter,
		KeyUsage:              x509.KeyUsageDigitalSignature,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth},
		BasicConstraintsValid: true,
		IsCA:                  false,
	}
}

func getServerTemplate(certName string, altDnsNameList DNSNames, orgUnit string, notAfter time.Time) *x509.Certificate {
	return &x509.Certificate{
		Subject:               getCertSubject(certName, orgUnit),
		SerialNumber:          generateSerial(),
		NotBefore:             currentTime.Add(-5 * time.Minute),
		NotAfter:              notAfter,
		KeyUsage:              x509.KeyUsageDigitalSignature,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		BasicConstraintsValid: true,
		IsCA:                  false,
		DNSNames:              altDnsNameList,
	}
}

func generateSerial() *big.Int {
	serialNumberLimit := new(big.Int).Lsh(big.NewInt(1), 128)
	serialNumber, err := rand.Int(rand.Reader, serialNumberLimit)
	fatalIfErr(err)

	return serialNumber
}

func sign(caInfo *caInfo, certFile string, keyFile string, certName string, tpl *x509.Certificate) *x509.Certificate {
	certKey := generateKey(keyFile)
	derBytes, err := x509.CreateCertificate(rand.Reader, tpl, caInfo.caCert, certKey.Public(), caInfo.caKey)
	fatalIfErr(err)
	writePem(certFile, derBytes, "CERTIFICATE")
	parsedCert, err := x509.ParseCertificate(derBytes)
	fatalIfErr(err)

	return parsedCert
}

func parseDateTime(dateTimeToParse *string) time.Time {
	parsedTime, err := time.Parse(time.RFC3339, *dateTimeToParse)
	fatalIfErr(err)
	if !parsedTime.After(currentTime) {
		log.Fatal("-not-after must be in the future")
	}

	return parsedTime
}

func determineNotAfter(notAfter *string, defaultNotAfter time.Time, caNotAfter time.Time) time.Time {
	var notAfterTime time.Time

	switch *notAfter {
	case "":
		notAfterTime = defaultNotAfter
	case "CA":
		notAfterTime = caNotAfter
	default:
		notAfterTime = parseDateTime(notAfter)
	}

	// make sure the certificate won't outlive the CA
	if notAfterTime.After(caNotAfter) {
		log.Fatal("-not-after can't outlive the CA")
	}

	return notAfterTime
}

func determineOutFile(caDir, outFile, certName, fileExt string) string {
	if outFile == "" {
		// when the "outFile" is not specified we use the name to be used on
		// the certificate as the file name. It MUST NOT contain any characters
		// outside the set of characters allowed for domain names to make sure
		// they can be safely stored on the file system
		validateCertName(certName)
		return filepath.Join(caDir, fmt.Sprintf("%s.%s", certName, fileExt))
	}

	if filepath.IsAbs(outFile) {
		return outFile
	}

	// file specified by their relative path are written to "caDir"
	return filepath.Join(caDir, outFile)
}

func validateCertName(certName string) {
	validCertName := regexp.MustCompile(`^[a-zA-Z0-9-.]+$`)
	if !validCertName.MatchString(certName) {
		log.Fatal("invalid 'certName' specified")
	}
}

func main() {
	var dnsNames DNSNames
	caDir := os.Getenv("CA_DIR")
	if caDir == "" {
		caDir = "."
	}
	initCa := flag.Bool("init-ca", false, "Initialize CA")
	domainConstraint := flag.String("domain-constraint", "", "DNS Domain Constraint for CA")
	serverCert := flag.Bool("server", false, "Generate a Server Certificate/Key")
	clientCert := flag.Bool("client", false, "Generate a Client Certificate/Key")
	certName := flag.String("name", "", "Common Name on the CA/Server/Client Certificate")
	flag.Var(&dnsNames, "alt-dns-name", "Server Certificate (Subject Alternative Name) DNS Name")
	orgUnit := flag.String("ou", "", "Organizational Unit on the Server/Client Certificate")
	notAfter := flag.String("not-after", "", "Limit Certificate Validity for CA/Server/Client Certificate")
	outCert := flag.String("out-crt", "", "Write the Server/Client Certificate to Specified File")
	outKey := flag.String("out-key", "", "Write the Server/Client Key to Specified File")

	flag.Usage = func() {
		flag.PrintDefaults()
	}
	flag.Parse()

	if *initCa {
		if *certName == "" {
			*certName = "Root CA"
		}
		// default to 5y
		caNotAfter := currentTime.AddDate(5, 0, 0)
		if *notAfter != "" && *notAfter != "CA" {
			caNotAfter = parseDateTime(notAfter)
		}
		caInit(caDir, *certName, caNotAfter, *domainConstraint)
		return
	}

	if *certName == "" {
		flag.Usage()
		os.Exit(1)
	}

	caInfo := getCa(caDir)

	if *serverCert && *clientCert {
		flag.Usage()
		os.Exit(1)
	}

	certFile := determineOutFile(caInfo.caDir, *outCert, *certName, "crt")
	keyFile := determineOutFile(caInfo.caDir, *outKey, *certName, "key")

	if *serverCert {
		// add the certName also to the "subjectAltName"
		dnsNames.Set(*certName)
		notAfterTime := determineNotAfter(notAfter, currentTime.AddDate(1, 0, 0), caInfo.caCert.NotAfter)
		sign(caInfo, certFile, keyFile, *certName, getServerTemplate(*certName, dnsNames, *orgUnit, notAfterTime))
		return
	}

	if *clientCert {
		notAfterTime := determineNotAfter(notAfter, currentTime.AddDate(1, 0, 0), caInfo.caCert.NotAfter)
		sign(caInfo, certFile, keyFile, *certName, getClientTemplate(*certName, *orgUnit, notAfterTime))
		return
	}

	flag.Usage()
	os.Exit(1)
}

func fatalIfErr(err error) {
	if err != nil {
		log.Fatal(err)
	}
}
