LokiStack/backend/services/loki_deployer.go

package services

import (
	"context"
	"fmt"
	"os"
	"os/exec"
	"strings"

	"lokideploy/models"
	"lokideploy/utils"
)

type LokiDeployer struct {
	deployment *models.LokiDeployment
	logFunc    func(level, message, detail string)
}

func NewLokiDeployer(deployment *models.LokiDeployment, logFunc func(level, message, detail string)) *LokiDeployer {
	return &LokiDeployer{deployment: deployment, logFunc: logFunc}
}

func (d *LokiDeployer) Deploy(ctx context.Context) error {
	d.log("info", fmt.Sprintf("Starting Loki deployment: %s (mode: %s)", d.deployment.Name, d.deployment.Mode), "")
	switch d.deployment.Mode {
	case "k8s-helm":
		return d.deployHelm(ctx)
	case "k8s-yaml":
		return d.deployYaml(ctx)
	case "host-binary":
		return d.deployBinary(ctx)
	default:
		return fmt.Errorf("不支持的部署模式: %s", d.deployment.Mode)
	}
}

func (d *LokiDeployer) deployHelm(ctx context.Context) error {
	if d.deployment.Cluster == nil {
		return fmt.Errorf("请选择目标集群")
	}
	kubeconfig := d.deployment.Cluster.KubeConfig
	if kubeconfig == "" && d.deployment.Cluster.KubeConfigPath != "" {
		data, _ := os.ReadFile(d.deployment.Cluster.KubeConfigPath)
		kubeconfig = string(data)
	}

	kubeFile, _ := os.CreateTemp("", "kube-*.yaml")
	defer os.Remove(kubeFile.Name())
	kubeFile.WriteString(kubeconfig)
	kubeFile.Close()

	env := append(os.Environ(), fmt.Sprintf("KUBECONFIG=%s", kubeFile.Name()))
	ns := d.deployment.Namespace
	if ns == "" {
		ns = "loki"
	}

	d.log("info", fmt.Sprintf("Creating namespace %s...", ns), "")
	exec.CommandContext(ctx, "kubectl", "create", "namespace", ns, "--dry-run=client", "-o", "yaml").CombinedOutput()

	d.log("info", "Adding Grafana Helm repo...", "")
	exec.CommandContext(ctx, "helm", "repo", "add", "grafana", "https://grafana.github.io/helm-charts").CombinedOutput()
	exec.CommandContext(ctx, "helm", "repo", "update").CombinedOutput()

	d.log("info", fmt.Sprintf("Installing Loki in %s...", ns), "")
	releaseName := d.deployment.Name
	if releaseName == "" {
		releaseName = "loki"
	}

	chart := "grafana/loki"
	if d.deployment.Replicas > 1 {
		chart = "grafana/loki-distributed"
	}

	cmd := exec.CommandContext(ctx, "helm", "upgrade", "--install", releaseName, chart,
		"--namespace", ns, "--create-namespace", "--wait", "--timeout", "10m")
	cmd.Env = env
	out, err := cmd.CombinedOutput()
	if err != nil {
		d.log("error", fmt.Sprintf("Helm install failed: %v", err), string(out))
		return err
	}
	d.log("success", "Loki installed via Helm", string(out))
	return nil
}

func (d *LokiDeployer) deployYaml(ctx context.Context) error {
	if d.deployment.Cluster == nil {
		return fmt.Errorf("请选择目标集群")
	}
	kubeconfig := d.deployment.Cluster.KubeConfig
	if kubeconfig == "" && d.deployment.Cluster.KubeConfigPath != "" {
		data, _ := os.ReadFile(d.deployment.Cluster.KubeConfigPath)
		kubeconfig = string(data)
	}

	kubeFile, _ := os.CreateTemp("", "kube-*.yaml")
	defer os.Remove(kubeFile.Name())
	kubeFile.WriteString(kubeconfig)
	kubeFile.Close()

	env := append(os.Environ(), fmt.Sprintf("KUBECONFIG=%s", kubeFile.Name()))
	ns := d.deployment.Namespace
	if ns == "" {
		ns = "loki"
	}

	manifest := d.generateManifest(ns)
	manifestFile, _ := os.CreateTemp("", "loki-*.yaml")
	defer os.Remove(manifestFile.Name())
	manifestFile.WriteString(manifest)
	manifestFile.Close()

	cmd := exec.CommandContext(ctx, "kubectl", "apply", "-f", manifestFile.Name(), "-n", ns)
	cmd.Env = env
	out, err := cmd.CombinedOutput()
	if err != nil {
		d.log("error", fmt.Sprintf("kubectl apply failed: %v", err), string(out))
		return err
	}
	d.log("success", "Loki deployed via YAML", string(out))
	return nil
}

func (d *LokiDeployer) deployBinary(ctx context.Context) error {
	if d.deployment.Host == nil {
		return fmt.Errorf("请选择目标主机")
	}
	host := d.deployment.Host
	password := host.Password
	if strings.HasPrefix(password, "ENC:") {
		password, _ = DecryptCred(password)
	}

	client, err := NewSSHClient(host.HostIP, host.Port, host.User, password)
	if err != nil {
		return err
	}
	defer client.Close()

	d.log("info", "SSH connected, downloading Loki...", "")
	arch, _, _ := client.Execute(ctx, "uname -m")
	lokiArch := strings.TrimSpace(arch)
	switch lokiArch {
	case "x86_64":
		lokiArch = "amd64"
	case "aarch64":
		lokiArch = "arm64"
	}

	version := d.deployment.Version
	if version == "" {
		version = "3.0.0"
	}

	client.Execute(ctx, "mkdir -p /opt/loki /var/lib/loki /etc/loki")

	downloadURL := fmt.Sprintf("https://github.com/grafana/loki/releases/download/v%s/loki-linux-%s.zip", version, lokiArch)
	cmd := fmt.Sprintf("cd /tmp && curl -sL -o loki.zip '%s' && unzip -o loki.zip -d /opt/loki/ && chmod +x /opt/loki/loki-linux-%s", downloadURL, lokiArch)
	client.Execute(ctx, cmd)
	client.Execute(ctx, "ln -sf /opt/loki/loki-linux-* /usr/local/bin/loki 2>/dev/null || true")

	// Upload config
	config := d.generateLokiConfig()
	client.Execute(ctx, fmt.Sprintf("cat > /etc/loki/loki.yaml << 'EOF'\n%s\nEOF", config))

	// Systemd service
	service := `[Unit]
Description=Grafana Loki
Wants=network-online.target
After=network-online.target

[Service]
Type=simple
ExecStart=/usr/local/bin/loki -config.file=/etc/loki/loki.yaml
Restart=always

[Install]
WantedBy=multi-user.target
`
	client.Execute(ctx, fmt.Sprintf("cat > /etc/systemd/system/loki.service << 'EOF'\n%s\nEOF", service))
	client.Execute(ctx, "systemctl daemon-reload && systemctl enable loki && systemctl restart loki")

	d.log("success", fmt.Sprintf("Loki deployed on %s", host.HostIP), "")
	return nil
}

func (d *LokiDeployer) generateManifest(ns string) string {
	return fmt.Sprintf(`apiVersion: v1
kind: Namespace
metadata:
  name: %s
---
apiVersion: apps/v1
kind: Deployment
metadata:
  name: loki
  namespace: %s
spec:
  replicas: 1
  selector:
    matchLabels:
      app: loki
  template:
    metadata:
      labels:
        app: loki
    spec:
      containers:
      - name: loki
        image: grafana/loki:%s
        ports:
        - containerPort: 3100
---
apiVersion: v1
kind: Service
metadata:
  name: loki
  namespace: %s
spec:
  ports:
  - port: 3100
  selector:
    app: loki
`, ns, ns, d.deployment.Version, ns)
}

func (d *LokiDeployer) generateLokiConfig() string {
	return fmt.Sprintf(`auth_enabled: false
server:
  http_listen_port: 3100
limits_config:
  retention_period: %dh
schema_config:
  configs:
  - from: 2020-10-24
    store: tsdb
    object_store: filesystem
    schema: v13
    index:
      prefix: index_
      period: 24h
`, d.deployment.RetentionDays*24)
}

func (d *LokiDeployer) log(level, message, detail string) {
	if d.logFunc != nil {
		d.logFunc(level, message, detail)
	}
}

func CleanupLoki(ctx context.Context, deployment *models.LokiDeployment) {
	if deployment.Mode == "k8s-helm" && deployment.Cluster != nil {
		kubeconfig := deployment.Cluster.KubeConfig
		if kubeconfig == "" && deployment.Cluster.KubeConfigPath != "" {
			data, _ := os.ReadFile(deployment.Cluster.KubeConfigPath)
			kubeconfig = string(data)
		}
		kubeFile, _ := os.CreateTemp("", "kube-*.yaml")
		defer os.Remove(kubeFile.Name())
		os.WriteFile(kubeFile.Name(), []byte(kubeconfig), 0600)
		env := append(os.Environ(), fmt.Sprintf("KUBECONFIG=%s", kubeFile.Name()))
		cmd := exec.CommandContext(ctx, "helm", "uninstall", deployment.Name, "--namespace", deployment.Namespace)
		cmd.Env = env
		cmd.CombinedOutput()
	}
	if deployment.Mode == "host-binary" && deployment.Host != nil {
		host := deployment.Host
		password := host.Password
		if strings.HasPrefix(password, "ENC:") {
			password, _ = DecryptCred(password)
		}
		client, _ := NewSSHClient(host.HostIP, host.Port, host.User, password)
		if client != nil {
			defer client.Close()
			client.Execute(ctx, "systemctl stop loki && systemctl disable loki && rm -f /etc/systemd/system/loki.service /usr/local/bin/loki")
		}
	}
}

func DecryptCred(encrypted string) (string, error) {
	if strings.HasPrefix(encrypted, "ENC:") {
		return utils.Decrypt(encrypted[4:])
	}
	return encrypted, nil
}