at-container-registry/deploy/upcloud/cloudinit.go

package main

import (
	"bytes"
	_ "embed"
	"fmt"
	"strings"
	"text/template"

	"go.yaml.in/yaml/v3"
)

//go:embed systemd/appview.service.tmpl
var appviewServiceTmpl string

//go:embed systemd/hold.service.tmpl
var holdServiceTmpl string

//go:embed systemd/scanner.service.tmpl
var scannerServiceTmpl string

//go:embed configs/appview.yaml.tmpl
var appviewConfigTmpl string

//go:embed configs/hold.yaml.tmpl
var holdConfigTmpl string

//go:embed configs/scanner.yaml.tmpl
var scannerConfigTmpl string

//go:embed configs/cloudinit.sh.tmpl
var cloudInitTmpl string

// ConfigValues holds values injected into config YAML templates.
// Only truly dynamic/computed values belong here — deployment-specific
// values like client_name, owner_did, etc. are literal in the templates.
type ConfigValues struct {
	// S3 / Object Storage
	S3Endpoint  string
	S3Region    string
	S3Bucket    string
	S3AccessKey string
	S3SecretKey string

	// Infrastructure (computed from zone + config)
	Zone       string // e.g. "us-chi1"
	HoldDomain string // e.g. "us-chi1.cove.seamark.dev"
	HoldDid    string // e.g. "did:web:us-chi1.cove.seamark.dev"
	BasePath   string // e.g. "/var/lib/seamark"

	// Scanner (auto-generated shared secret)
	ScannerSecret string // hex-encoded 32-byte secret; empty disables scanning
}

// renderConfig executes a Go template with the given values.
func renderConfig(tmplStr string, vals *ConfigValues) (string, error) {
	t, err := template.New("config").Parse(tmplStr)
	if err != nil {
		return "", fmt.Errorf("parse config template: %w", err)
	}
	var buf bytes.Buffer
	if err := t.Execute(&buf, vals); err != nil {
		return "", fmt.Errorf("render config template: %w", err)
	}
	return buf.String(), nil
}

// serviceUnitParams holds values for rendering systemd service unit templates.
type serviceUnitParams struct {
	DisplayName string // e.g. "Seamark"
	User        string // e.g. "seamark"
	BinaryPath  string // e.g. "/opt/seamark/bin/seamark-appview"
	ConfigPath  string // e.g. "/etc/seamark/appview.yaml"
	DataDir     string // e.g. "/var/lib/seamark"
	ServiceName string // e.g. "seamark-appview"
}

func renderServiceUnit(tmplStr string, p serviceUnitParams) (string, error) {
	t, err := template.New("service").Parse(tmplStr)
	if err != nil {
		return "", fmt.Errorf("parse service template: %w", err)
	}
	var buf bytes.Buffer
	if err := t.Execute(&buf, p); err != nil {
		return "", fmt.Errorf("render service template: %w", err)
	}
	return buf.String(), nil
}

// scannerServiceUnitParams holds values for rendering the scanner systemd unit.
// Extends the standard fields with HoldServiceName for the After= dependency.
type scannerServiceUnitParams struct {
	DisplayName     string // e.g. "Seamark"
	User            string // e.g. "seamark"
	BinaryPath      string // e.g. "/opt/seamark/bin/seamark-scanner"
	ConfigPath      string // e.g. "/etc/seamark/scanner.yaml"
	DataDir         string // e.g. "/var/lib/seamark"
	ServiceName     string // e.g. "seamark-scanner"
	HoldServiceName string // e.g. "seamark-hold" (After= dependency)
}

func renderScannerServiceUnit(p scannerServiceUnitParams) (string, error) {
	t, err := template.New("scanner-service").Parse(scannerServiceTmpl)
	if err != nil {
		return "", fmt.Errorf("parse scanner service template: %w", err)
	}
	var buf bytes.Buffer
	if err := t.Execute(&buf, p); err != nil {
		return "", fmt.Errorf("render scanner service template: %w", err)
	}
	return buf.String(), nil
}

// generateAppviewCloudInit generates the cloud-init user-data script for the appview server.
// Sets up the OS, directories, config, and systemd unit. Binaries are deployed separately via SCP.
func generateAppviewCloudInit(cfg *InfraConfig, vals *ConfigValues) (string, error) {
	naming := cfg.Naming()

	configYAML, err := renderConfig(appviewConfigTmpl, vals)
	if err != nil {
		return "", fmt.Errorf("appview config: %w", err)
	}

	serviceUnit, err := renderServiceUnit(appviewServiceTmpl, serviceUnitParams{
		DisplayName: naming.DisplayName(),
		User:        naming.SystemUser(),
		BinaryPath:  naming.InstallDir() + "/bin/" + naming.Appview(),
		ConfigPath:  naming.AppviewConfigPath(),
		DataDir:     naming.BasePath(),
		ServiceName: naming.Appview(),
	})
	if err != nil {
		return "", fmt.Errorf("appview service unit: %w", err)
	}

	return generateCloudInit(cloudInitParams{
		BinaryName:  naming.Appview(),
		ServiceUnit: serviceUnit,
		ConfigYAML:  configYAML,
		ConfigPath:  naming.AppviewConfigPath(),
		ServiceName: naming.Appview(),
		DataDir:     naming.BasePath(),
		InstallDir:  naming.InstallDir(),
		SystemUser:  naming.SystemUser(),
		ConfigDir:   naming.ConfigDir(),
		LogFile:     naming.LogFile(),
		DisplayName: naming.DisplayName(),
	})
}

// generateHoldCloudInit generates the cloud-init user-data script for the hold server.
// When withScanner is true, a second phase is appended that creates scanner data
// directories and installs a scanner systemd service. Binaries are deployed separately via SCP.
func generateHoldCloudInit(cfg *InfraConfig, vals *ConfigValues, withScanner bool) (string, error) {
	naming := cfg.Naming()

	configYAML, err := renderConfig(holdConfigTmpl, vals)
	if err != nil {
		return "", fmt.Errorf("hold config: %w", err)
	}

	serviceUnit, err := renderServiceUnit(holdServiceTmpl, serviceUnitParams{
		DisplayName: naming.DisplayName(),
		User:        naming.SystemUser(),
		BinaryPath:  naming.InstallDir() + "/bin/" + naming.Hold(),
		ConfigPath:  naming.HoldConfigPath(),
		DataDir:     naming.BasePath(),
		ServiceName: naming.Hold(),
	})
	if err != nil {
		return "", fmt.Errorf("hold service unit: %w", err)
	}

	script, err := generateCloudInit(cloudInitParams{
		BinaryName:  naming.Hold(),
		ServiceUnit: serviceUnit,
		ConfigYAML:  configYAML,
		ConfigPath:  naming.HoldConfigPath(),
		ServiceName: naming.Hold(),
		DataDir:     naming.BasePath(),
		InstallDir:  naming.InstallDir(),
		SystemUser:  naming.SystemUser(),
		ConfigDir:   naming.ConfigDir(),
		LogFile:     naming.LogFile(),
		DisplayName: naming.DisplayName(),
	})
	if err != nil {
		return "", err
	}

	if !withScanner {
		return script, nil
	}

	// Render scanner config YAML
	scannerConfigYAML, err := renderConfig(scannerConfigTmpl, vals)
	if err != nil {
		return "", fmt.Errorf("scanner config: %w", err)
	}

	// Append scanner setup phase (no build — binary deployed via SCP)
	scannerUnit, err := renderScannerServiceUnit(scannerServiceUnitParams{
		DisplayName:     naming.DisplayName(),
		User:            naming.SystemUser(),
		BinaryPath:      naming.InstallDir() + "/bin/" + naming.Scanner(),
		ConfigPath:      naming.ScannerConfigPath(),
		DataDir:         naming.BasePath(),
		ServiceName:     naming.Scanner(),
		HoldServiceName: naming.Hold(),
	})
	if err != nil {
		return "", fmt.Errorf("scanner service unit: %w", err)
	}

	// Escape single quotes for heredoc embedding
	scannerUnit = strings.ReplaceAll(scannerUnit, "'", "'\\''")
	scannerConfigYAML = strings.ReplaceAll(scannerConfigYAML, "'", "'\\''")

	scannerPhase := fmt.Sprintf(`
# === Scanner Setup ===

# Scanner data dirs
mkdir -p %s/vulndb %s/tmp
chown -R %s:%s %s

# Scanner config
cat > %s << 'CFGEOF'
%s
CFGEOF

# Scanner systemd service
cat > /etc/systemd/system/%s.service << 'SVCEOF'
%s
SVCEOF
systemctl daemon-reload
systemctl enable %s

echo "=== Scanner setup complete ==="
`,
		naming.ScannerDataDir(), naming.ScannerDataDir(),
		naming.SystemUser(), naming.SystemUser(), naming.ScannerDataDir(),
		naming.ScannerConfigPath(),
		scannerConfigYAML,
		naming.Scanner(),
		scannerUnit,
		naming.Scanner(),
	)

	return script + scannerPhase, nil
}

type cloudInitParams struct {
	BinaryName  string
	ServiceUnit string
	ConfigYAML  string
	ConfigPath  string
	ServiceName string
	DataDir     string
	InstallDir  string
	SystemUser  string
	ConfigDir   string
	LogFile     string
	DisplayName string
}

func generateCloudInit(p cloudInitParams) (string, error) {
	// Escape single quotes in embedded content for heredoc safety
	p.ServiceUnit = strings.ReplaceAll(p.ServiceUnit, "'", "'\\''")
	p.ConfigYAML = strings.ReplaceAll(p.ConfigYAML, "'", "'\\''")

	t, err := template.New("cloudinit").Parse(cloudInitTmpl)
	if err != nil {
		return "", fmt.Errorf("parse cloudinit template: %w", err)
	}
	var buf bytes.Buffer
	if err := t.Execute(&buf, p); err != nil {
		return "", fmt.Errorf("render cloudinit template: %w", err)
	}
	return buf.String(), nil
}

// syncServiceUnit compares a rendered systemd service unit against what's on
// the server. If they differ, it writes the new unit file. Returns true if the
// unit was updated (caller should daemon-reload before restart).
func syncServiceUnit(name, ip, serviceName, renderedUnit string) (bool, error) {
	unitPath := "/etc/systemd/system/" + serviceName + ".service"

	remote, err := runSSH(ip, fmt.Sprintf("cat %s 2>/dev/null || echo '__MISSING__'", unitPath), false)
	if err != nil {
		fmt.Printf("  service unit sync: could not reach %s (%v)\n", name, err)
		return false, nil
	}
	remote = strings.TrimSpace(remote)
	rendered := strings.TrimSpace(renderedUnit)

	if remote == "__MISSING__" {
		fmt.Printf("  service unit: %s not found (cloud-init will handle it)\n", name)
		return false, nil
	}

	if remote == rendered {
		fmt.Printf("  service unit: %s up to date\n", name)
		return false, nil
	}

	// Write the updated unit file
	script := fmt.Sprintf("cat > %s << 'SVCEOF'\n%s\nSVCEOF", unitPath, rendered)
	if _, err := runSSH(ip, script, false); err != nil {
		return false, fmt.Errorf("write service unit: %w", err)
	}
	fmt.Printf("  service unit: %s updated\n", name)
	return true, nil
}

// syncConfigKeys fetches the existing config from a server and merges in any
// missing keys from the rendered template. Existing values are never overwritten.
func syncConfigKeys(name, ip, configPath, templateYAML string) error {
	remote, err := runSSH(ip, fmt.Sprintf("cat %s 2>/dev/null || echo '__MISSING__'", configPath), false)
	if err != nil {
		fmt.Printf("  config sync: could not reach %s (%v)\n", name, err)
		return nil
	}
	remote = strings.TrimSpace(remote)

	if remote == "__MISSING__" {
		fmt.Printf("  config sync: %s not yet created (cloud-init will handle it)\n", name)
		return nil
	}

	// Parse both into yaml.Node trees
	var templateDoc yaml.Node
	if err := yaml.Unmarshal([]byte(templateYAML), &templateDoc); err != nil {
		return fmt.Errorf("parse template yaml: %w", err)
	}
	var existingDoc yaml.Node
	if err := yaml.Unmarshal([]byte(remote), &existingDoc); err != nil {
		return fmt.Errorf("parse remote yaml: %w", err)
	}

	// Unwrap document nodes to get the root mapping
	templateRoot := unwrapDocNode(&templateDoc)
	existingRoot := unwrapDocNode(&existingDoc)
	if templateRoot == nil || existingRoot == nil {
		fmt.Printf("  config sync: %s skipped (unexpected YAML structure)\n", name)
		return nil
	}

	added := mergeYAMLNodes(templateRoot, existingRoot)
	if !added {
		fmt.Printf("  config sync: %s up to date\n", name)
		return nil
	}

	// Marshal the modified tree back
	merged, err := yaml.Marshal(&existingDoc)
	if err != nil {
		return fmt.Errorf("marshal merged yaml: %w", err)
	}

	// Write back to server
	script := fmt.Sprintf("cat > %s << 'CFGEOF'\n%sCFGEOF", configPath, string(merged))
	if _, err := runSSH(ip, script, false); err != nil {
		return fmt.Errorf("write merged config: %w", err)
	}
	fmt.Printf("  config sync: %s updated with new keys\n", name)
	return nil
}

// unwrapDocNode returns the root mapping node, unwrapping a DocumentNode wrapper if present.
func unwrapDocNode(n *yaml.Node) *yaml.Node {
	if n.Kind == yaml.DocumentNode && len(n.Content) > 0 {
		return n.Content[0]
	}
	if n.Kind == yaml.MappingNode {
		return n
	}
	return nil
}

// mergeYAMLNodes recursively adds keys from base into existing that are not
// already present. Existing values are never overwritten. Returns true if any
// new keys were added.
func mergeYAMLNodes(base, existing *yaml.Node) bool {
	if base.Kind != yaml.MappingNode || existing.Kind != yaml.MappingNode {
		return false
	}

	added := false
	for i := 0; i+1 < len(base.Content); i += 2 {
		baseKey := base.Content[i]
		baseVal := base.Content[i+1]

		// Look for this key in existing
		found := false
		for j := 0; j+1 < len(existing.Content); j += 2 {
			if existing.Content[j].Value == baseKey.Value {
				found = true
				// If both are mappings, recurse to merge sub-keys
				if baseVal.Kind == yaml.MappingNode && existing.Content[j+1].Kind == yaml.MappingNode {
					if mergeYAMLNodes(baseVal, existing.Content[j+1]) {
						added = true
					}
				}
				break
			}
		}

		if !found {
			// Append the missing key+value pair
			existing.Content = append(existing.Content, baseKey, baseVal)
			added = true
		}
	}

	return added
}