HellreigN/agent/internal/registration/registration.go

package registration

import (
	"bytes"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/json"
	"encoding/pem"
	"fmt"
	"net/http"
	"os"
	"path/filepath"
)

type Certs struct {
	CACertPEM     []byte
	ClientCertPEM []byte
	ClientKeyPEM  []byte
}

type RegisterRequest struct {
	CSR   string `json:"csr"`
	Token string `json:"token"`
}

type RegisterResponse struct {
	CACert     string `json:"ca_cert"`
	ClientCert string `json:"client_cert"`
}

type TokenResponse struct {
	Token string `json:"token"`
}

type ErrorResponse struct {
	Error string `json:"error"`
}

// GenerateKeyAndCSR generates a new ECDSA private key and CSR for the agent.
func GenerateKeyAndCSR(label string) (*ecdsa.PrivateKey, []byte, error) {
	key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		return nil, nil, fmt.Errorf("generate key: %w", err)
	}

	template := x509.CertificateRequest{
		Subject: pkix.Name{
			CommonName:   label,
			Organization: []string{"HellreigN Agent"},
		},
		SignatureAlgorithm: x509.ECDSAWithSHA256,
	}

	csrDER, err := x509.CreateCertificateRequest(rand.Reader, &template, key)
	if err != nil {
		return nil, nil, fmt.Errorf("create csr: %w", err)
	}

	csrPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "CERTIFICATE REQUEST",
		Bytes: csrDER,
	})

	return key, csrPEM, nil
}

// Register sends CSR to backend and receives signed certificates.
func Register(backendURL, token string, csrPEM []byte) (*Certs, error) {
	reqBody := RegisterRequest{CSR: string(csrPEM), Token: token}
	body, err := json.Marshal(reqBody)
	if err != nil {
		return nil, err
	}

	url := fmt.Sprintf("%s/api/v1/agents/register", backendURL)
	resp, err := http.Post(url, "application/json", bytes.NewReader(body))
	if err != nil {
		return nil, fmt.Errorf("register request: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		var errResp ErrorResponse
		json.NewDecoder(resp.Body).Decode(&errResp)
		return nil, fmt.Errorf("registration failed (status %d): %s", resp.StatusCode, errResp.Error)
	}

	var regResp RegisterResponse
	if err := json.NewDecoder(resp.Body).Decode(&regResp); err != nil {
		return nil, fmt.Errorf("decode response: %w", err)
	}

	return &Certs{
		CACertPEM:     []byte(regResp.CACert),
		ClientCertPEM: []byte(regResp.ClientCert),
	}, nil
}

// SaveCerts saves CA cert, client cert, and client key to the given directory.
func SaveCerts(certDir string, certs *Certs, key *ecdsa.PrivateKey) error {
	if err := os.MkdirAll(certDir, 0700); err != nil {
		return fmt.Errorf("create cert dir: %w", err)
	}

	if err := os.WriteFile(filepath.Join(certDir, "ca.crt"), certs.CACertPEM, 0644); err != nil {
		return err
	}
	if err := os.WriteFile(filepath.Join(certDir, "client.crt"), certs.ClientCertPEM, 0644); err != nil {
		return err
	}

	keyDER, err := x509.MarshalECPrivateKey(key)
	if err != nil {
		return err
	}
	keyPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "EC PRIVATE KEY",
		Bytes: keyDER,
	})
	if err := os.WriteFile(filepath.Join(certDir, "client.key"), keyPEM, 0600); err != nil {
		return err
	}

	return nil
}

// LoadCerts loads existing certificates and key from disk.
func LoadCerts(certDir string) (*Certs, *ecdsa.PrivateKey, error) {
	caCert, err := os.ReadFile(filepath.Join(certDir, "ca.crt"))
	if err != nil {
		return nil, nil, err
	}
	clientCert, err := os.ReadFile(filepath.Join(certDir, "client.crt"))
	if err != nil {
		return nil, nil, err
	}
	clientKeyPEM, err := os.ReadFile(filepath.Join(certDir, "client.key"))
	if err != nil {
		return nil, nil, err
	}

	block, _ := pem.Decode(clientKeyPEM)
	if block == nil {
		return nil, nil, fmt.Errorf("decode client key")
	}
	key, err := x509.ParseECPrivateKey(block.Bytes)
	if err != nil {
		return nil, nil, fmt.Errorf("parse client key: %w", err)
	}

	return &Certs{
		CACertPEM:     caCert,
		ClientCertPEM: clientCert,
	}, key, nil
}

// CertsExist checks if all certificate files exist in the directory.
func CertsExist(certDir string) bool {
	files := []string{"ca.crt", "client.crt", "client.key"}
	for _, f := range files {
		if _, err := os.Stat(filepath.Join(certDir, f)); err != nil {
			return false
		}
	}
	return true
}