package main

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha512"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"errors"
	"fmt"
	"math/big"
	"os"
	"time"
)

// GenerateLocalhost generates a new key pair for localhost signed by the given CA key pair
func GenerateLocalhost(ca *x509.Certificate, caKey *rsa.PrivateKey) ([]byte, *rsa.PrivateKey, error) {
	// generate random serial
	buf := make([]byte, 16)
	if _, err := rand.Read(buf); err != nil {
		return nil, nil, fmt.Errorf("could not generate serial: %w", err)
	}

	serial := new(big.Int)
	serial.SetBytes(buf)

	// generate key
	key, err := rsa.GenerateKey(rand.Reader, 4096)
	if err != nil {
		return nil, nil, fmt.Errorf("could not generate key: %w", err)
	}

	// generate subject key id by hashing the public key
	ski := sha512.Sum512(x509.MarshalPKCS1PublicKey(&key.PublicKey))

	tmpl := &x509.Certificate{
		SerialNumber: serial,
		Subject: pkix.Name{
			Organization: []string{"Lightspeed Systems"},
		},
		SubjectKeyId:          ski[:],
		DNSNames:              []string{"localhost"},
		NotBefore:             time.Now().Add(-time.Minute), // give a little buffer
		NotAfter:              time.Now().AddDate(1, 0, 0), // 1 year expiry
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		KeyUsage:              x509.KeyUsageDigitalSignature | x509.KeyUsageKeyEncipherment,
		BasicConstraintsValid: true,
	}

	// marshal certificate
	der, err := x509.CreateCertificate(rand.Reader, tmpl, ca, &key.PublicKey, caKey)
	if err != nil {
		return nil, nil, fmt.Errorf("could not generate certificate: %w", err)
	}

	return der, key, nil
}

// UpdateLocalhostCertificate checks if the given cert exists and is still valid and generates a new key pair with the given ca if not
func UpdateLocalhostCertificate(capath, cakeypath, certpath, certkeypath string) error {
	var (
		block *pem.Block
		cert  *x509.Certificate
	)

	// read cert
	buf, err := os.ReadFile(certpath)
	if err != nil {
		// if cert doesn't exist, skip to generation
		if errors.Is(err, os.ErrNotExist) {
			goto gen
		}
		return fmt.Errorf("could not read %s: %w", certpath, err)
	}

	// parse cert PEM data
	block, _ = pem.Decode(buf)
	if block.Type != "CERTIFICATE" {
		return fmt.Errorf("could not parse %s: expected CERTIFICATE block, got %s", certpath, block.Type)
	}

	// parse cert der to certificate
	cert, err = x509.ParseCertificate(block.Bytes)
	if err != nil {
		return fmt.Errorf("could not parse %s: %w", certpath, err)
	}

	// check if expired
	if !time.Now().After(cert.NotAfter) {
		// not expired, no work to do
		return nil
	}

gen:
	// read ca cert
	buf, err = os.ReadFile(capath)
	if err != nil {
		return fmt.Errorf("could not read CA %s: %w", capath, err)
	}

	// parse ca PEM data
	block, _ = pem.Decode(buf)
	if block.Type != "CERTIFICATE" {
		return fmt.Errorf("could not parse CA %s: expected CERTIFICATE block, got %s", capath, block.Type)
	}

	// parse ca der to certificate
	ca, err := x509.ParseCertificate(block.Bytes)
	if err != nil {
		return fmt.Errorf("could not parse CA %s: %w", capath, err)
	}

	// read ca key
	buf, err = os.ReadFile(cakeypath)
	if err != nil {
		return fmt.Errorf("could not read CA key %s: %w", cakeypath, err)
	}

	// parse ca key PEM data
	block, _ = pem.Decode(buf)
	if block.Type != "RSA PRIVATE KEY" {
		return fmt.Errorf("could not parse CA key %s: expected RSA PRIVATE KEY block, got %s", cakeypath, block.Type)
	}

	// parse ca PKCS1 key
	key, err := x509.ParsePKCS1PrivateKey(block.Bytes)
	if err != nil {
		return fmt.Errorf("could not parse CA key %s: %w", cakeypath, err)
	}

	// generate key pair
	certbuf, key, err := GenerateLocalhost(ca, key)
	if err != nil {
		return fmt.Errorf("could not generate new certificate: %w", err)
	}

	// encode key pair to PEM
	certpem := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: certbuf})
	keypem := pem.EncodeToMemory(&pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(key)})

	// write key pair to disk
	if err := os.WriteFile(certpath, certpem, 0600); err != nil {
		if err != nil {
			return fmt.Errorf("could not write new certificate %s: %w", certpath, err)
		}
	}

	if err := os.WriteFile(certkeypath, keypem, 0600); err != nil {
		if err != nil {
			return fmt.Errorf("could not write new certificate key %s: %w", certkeypath, err)
		}
	}

	return nil
}

func main() {
	err := UpdateLocalhostCertificate(
		"/usr/local/etc/ca.pem",
		"/usr/local/etc/ca_key.pem",
		"/usr/local/etc/localhost.pem",
		"/usr/local/etc/localhost_key.pem",
	)

	fmt.Println(err)
}