// Package ca manages mitmux's root CA: generating it on first run,
// loading it on subsequent runs, and signing per-host leaf certificates
// on demand for TLS interception.
package ca

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"errors"
	"fmt"
	"math/big"
	"net"
	"os"
	"path/filepath"
	"sync"
	"time"
)

const (
	certFileName = "ca.pem"
	keyFileName  = "ca-key.pem"
)

// CA holds the root certificate and key used to sign per-host leaf
// certificates during TLS interception.
type CA struct {
	Cert    *x509.Certificate
	Key     *ecdsa.PrivateKey
	CertPEM []byte
	KeyPEM  []byte

	leafMu    sync.Mutex
	leafCache map[string]*tls.Certificate
}

// leafLifetime is kept well under the ~398 day limit modern browsers
// enforce on leaf certificates.
const leafLifetime = 300 * 24 * time.Hour

// LeafFor returns a TLS certificate for host (a DNS name or IP address,
// no port), signed by the CA. Certificates are generated once and cached
// in memory for the life of the process.
func (c *CA) LeafFor(host string) (*tls.Certificate, error) {
	c.leafMu.Lock()
	defer c.leafMu.Unlock()

	if c.leafCache == nil {
		c.leafCache = make(map[string]*tls.Certificate)
	}
	if cert, ok := c.leafCache[host]; ok {
		return cert, nil
	}

	cert, err := c.signLeaf(host)
	if err != nil {
		return nil, err
	}
	c.leafCache[host] = cert
	return cert, nil
}

func (c *CA) signLeaf(host string) (*tls.Certificate, error) {
	key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		return nil, err
	}

	serial, err := rand.Int(rand.Reader, new(big.Int).Lsh(big.NewInt(1), 128))
	if err != nil {
		return nil, err
	}

	tmpl := &x509.Certificate{
		SerialNumber: serial,
		Subject: pkix.Name{
			CommonName:   host,
			Organization: []string{"mitmux"},
		},
		NotBefore:             time.Now().Add(-time.Hour),
		NotAfter:              time.Now().Add(leafLifetime),
		KeyUsage:              x509.KeyUsageDigitalSignature,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		BasicConstraintsValid: true,
		IsCA:                  false,
	}
	if ip := net.ParseIP(host); ip != nil {
		tmpl.IPAddresses = []net.IP{ip}
	} else {
		tmpl.DNSNames = []string{host}
	}

	der, err := x509.CreateCertificate(rand.Reader, tmpl, c.Cert, &key.PublicKey, c.Key)
	if err != nil {
		return nil, err
	}

	leaf, err := x509.ParseCertificate(der)
	if err != nil {
		return nil, err
	}

	return &tls.Certificate{
		Certificate: [][]byte{der, c.Cert.Raw},
		PrivateKey:  key,
		Leaf:        leaf,
	}, nil
}

// Dir returns the directory mitmux stores its CA material in
// (XDG config dir, e.g. ~/.config/mitmux).
func Dir() (string, error) {
	cfg, err := os.UserConfigDir()
	if err != nil {
		return "", fmt.Errorf("resolve config dir: %w", err)
	}
	return filepath.Join(cfg, "mitmux"), nil
}

// EnsureCA loads the CA from dir, generating and persisting a new one if
// none exists yet.
func EnsureCA(dir string) (*CA, error) {
	certPath := filepath.Join(dir, certFileName)
	keyPath := filepath.Join(dir, keyFileName)

	certPEM, certErr := os.ReadFile(certPath)
	keyPEM, keyErr := os.ReadFile(keyPath)
	if certErr == nil && keyErr == nil {
		return load(certPEM, keyPEM)
	}
	if !errors.Is(certErr, os.ErrNotExist) && certErr != nil {
		return nil, fmt.Errorf("read %s: %w", certPath, certErr)
	}
	if !errors.Is(keyErr, os.ErrNotExist) && keyErr != nil {
		return nil, fmt.Errorf("read %s: %w", keyPath, keyErr)
	}

	ca, err := generate()
	if err != nil {
		return nil, fmt.Errorf("generate CA: %w", err)
	}
	if err := persist(dir, ca); err != nil {
		return nil, fmt.Errorf("persist CA: %w", err)
	}
	return ca, nil
}

func generate() (*CA, error) {
	key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		return nil, err
	}

	serial, err := rand.Int(rand.Reader, new(big.Int).Lsh(big.NewInt(1), 128))
	if err != nil {
		return nil, err
	}

	tmpl := &x509.Certificate{
		SerialNumber: serial,
		Subject: pkix.Name{
			CommonName:   "mitmux local CA",
			Organization: []string{"mitmux"},
		},
		NotBefore:             time.Now().Add(-time.Hour),
		NotAfter:              time.Now().AddDate(10, 0, 0),
		KeyUsage:              x509.KeyUsageCertSign | x509.KeyUsageDigitalSignature | x509.KeyUsageCRLSign,
		BasicConstraintsValid: true,
		IsCA:                  true,
		MaxPathLenZero:        true,
	}

	der, err := x509.CreateCertificate(rand.Reader, tmpl, tmpl, &key.PublicKey, key)
	if err != nil {
		return nil, err
	}
	cert, err := x509.ParseCertificate(der)
	if err != nil {
		return nil, err
	}

	keyDER, err := x509.MarshalECPrivateKey(key)
	if err != nil {
		return nil, err
	}

	return &CA{
		Cert:    cert,
		Key:     key,
		CertPEM: pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: der}),
		KeyPEM:  pem.EncodeToMemory(&pem.Block{Type: "EC PRIVATE KEY", Bytes: keyDER}),
	}, nil
}

func load(certPEM, keyPEM []byte) (*CA, error) {
	certBlock, _ := pem.Decode(certPEM)
	if certBlock == nil {
		return nil, errors.New("no PEM block found in CA certificate")
	}
	cert, err := x509.ParseCertificate(certBlock.Bytes)
	if err != nil {
		return nil, fmt.Errorf("parse CA certificate: %w", err)
	}

	keyBlock, _ := pem.Decode(keyPEM)
	if keyBlock == nil {
		return nil, errors.New("no PEM block found in CA key")
	}
	key, err := x509.ParseECPrivateKey(keyBlock.Bytes)
	if err != nil {
		return nil, fmt.Errorf("parse CA key: %w", err)
	}

	return &CA{Cert: cert, Key: key, CertPEM: certPEM, KeyPEM: keyPEM}, nil
}

func persist(dir string, ca *CA) error {
	if err := os.MkdirAll(dir, 0o700); err != nil {
		return err
	}
	if err := os.WriteFile(filepath.Join(dir, certFileName), ca.CertPEM, 0o644); err != nil {
		return err
	}
	if err := os.WriteFile(filepath.Join(dir, keyFileName), ca.KeyPEM, 0o600); err != nil {
		return err
	}
	return nil
}