package main

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

const (
	caCertFile = "ca.crt"
	caKeyFile  = "ca.key"
)

// MITM manages a local CA and issues per-host leaf certificates on demand.
// The browser must trust the CA cert for HTTPS interception to work.
// On first run the CA is auto-generated and stored in the config directory.
type MITM struct {
	dir    string
	caCert *x509.Certificate
	caKey  *ecdsa.PrivateKey

	mu    sync.Mutex
	cache map[string]*tls.Certificate // hostname → leaf cert
}

func NewMITM(dir string) (*MITM, error) {
	m := &MITM{dir: dir, cache: make(map[string]*tls.Certificate)}

	certPath := filepath.Join(dir, caCertFile)
	keyPath := filepath.Join(dir, caKeyFile)

	if _, err := os.Stat(certPath); os.IsNotExist(err) {
		log.Printf("mitm: generating CA in %s", dir)
		if err := generateCA(certPath, keyPath); err != nil {
			return nil, fmt.Errorf("generate CA: %w", err)
		}
	}

	if err := m.loadCA(certPath, keyPath); err != nil {
		return nil, fmt.Errorf("load CA: %w", err)
	}

	return m, nil
}

// CACertPath returns the path to the CA cert file the browser must trust.
func (m *MITM) CACertPath() string {
	return filepath.Join(m.dir, caCertFile)
}

// TLSConfig returns a server-side TLS config that dynamically issues certs
// for each hostname presented in the TLS SNI extension.
func (m *MITM) TLSConfig(fallbackHost string) *tls.Config {
	return &tls.Config{
		GetCertificate: func(hello *tls.ClientHelloInfo) (*tls.Certificate, error) {
			host := hello.ServerName
			if host == "" {
				host = fallbackHost
			}
			return m.certForHost(host)
		},
	}
}

func (m *MITM) certForHost(hostname string) (*tls.Certificate, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	if c, ok := m.cache[hostname]; ok {
		// Reissue if the cached cert expires within the next minute.
		if time.Until(c.Leaf.NotAfter) > time.Minute {
			return c, nil
		}
	}

	cert, err := m.issueCert(hostname)
	if err != nil {
		return nil, err
	}
	m.cache[hostname] = cert
	return cert, nil
}

func (m *MITM) issueCert(hostname string) (*tls.Certificate, error) {
	key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		return nil, err
	}

	tmpl := &x509.Certificate{
		SerialNumber: randomSerial(),
		Subject: pkix.Name{
			CommonName:   hostname,
			Organization: []string{"ChromeCard Component3"},
		},
		NotBefore:             time.Now().Add(-time.Minute),
		NotAfter:              time.Now().Add(24 * time.Hour),
		KeyUsage:              x509.KeyUsageDigitalSignature,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		BasicConstraintsValid: true,
	}

	if ip := net.ParseIP(hostname); ip != nil {
		tmpl.IPAddresses = []net.IP{ip}
	} else {
		tmpl.DNSNames = []string{hostname}
	}

	certDER, err := x509.CreateCertificate(rand.Reader, tmpl, m.caCert, &key.PublicKey, m.caKey)
	if err != nil {
		return nil, err
	}
	leaf, err := x509.ParseCertificate(certDER)
	if err != nil {
		return nil, err
	}

	return &tls.Certificate{
		Certificate: [][]byte{certDER},
		PrivateKey:  key,
		Leaf:        leaf,
	}, nil
}

func (m *MITM) loadCA(certPath, keyPath string) error {
	certPEM, err := os.ReadFile(certPath)
	if err != nil {
		return err
	}
	block, _ := pem.Decode(certPEM)
	if block == nil {
		return fmt.Errorf("no PEM block in %s", certPath)
	}
	cert, err := x509.ParseCertificate(block.Bytes)
	if err != nil {
		return err
	}

	keyPEM, err := os.ReadFile(keyPath)
	if err != nil {
		return err
	}
	block, _ = pem.Decode(keyPEM)
	if block == nil {
		return fmt.Errorf("no PEM block in %s", keyPath)
	}
	key, err := x509.ParseECPrivateKey(block.Bytes)
	if err != nil {
		return err
	}

	m.caCert = cert
	m.caKey = key
	return nil
}

func generateCA(certPath, keyPath string) error {
	key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		return err
	}

	tmpl := &x509.Certificate{
		SerialNumber: randomSerial(),
		Subject: pkix.Name{
			CommonName:   "ChromeCard Component3 CA",
			Organization: []string{"ChromeCard"},
		},
		NotBefore:             time.Now().Add(-time.Minute),
		NotAfter:              time.Now().Add(10 * 365 * 24 * time.Hour),
		KeyUsage:              x509.KeyUsageCertSign | x509.KeyUsageCRLSign,
		IsCA:                  true,
		BasicConstraintsValid: true,
	}

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

	if err := writePEM(certPath, "CERTIFICATE", certDER, 0644); err != nil {
		return err
	}

	keyDER, err := x509.MarshalECPrivateKey(key)
	if err != nil {
		return err
	}
	return writePEM(keyPath, "EC PRIVATE KEY", keyDER, 0600)
}

func writePEM(path, pemType string, der []byte, perm os.FileMode) error {
	f, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, perm)
	if err != nil {
		return err
	}
	defer f.Close()
	return pem.Encode(f, &pem.Block{Type: pemType, Bytes: der})
}

func randomSerial() *big.Int {
	n, _ := rand.Int(rand.Reader, new(big.Int).Lsh(big.NewInt(1), 128))
	return n
}