gc-infra/iac/cluster/utils/ansible-configuration.go

package utils

import (
	"bytes"
	"fmt"
	"regexp"
	"strings"
	"text/template"

	"github.com/pulumi/pulumi-command/sdk/go/command/remote"
	"github.com/pulumi/pulumi-hcloud/sdk/go/hcloud"
	"github.com/pulumi/pulumi/sdk/v3/go/pulumi"
)

type SwarmJoinTokens struct {
	ManagerToken string
	WorkerToken  string
}

type ServerInfo struct {
	Name pulumi.StringOutput
	IP   pulumi.StringOutput
}

func InstallAnsibleDependencies(ctx *pulumi.Context, connArgs remote.ConnectionArgs, uniqueness string) error {
	_, err := remote.NewCommand(ctx, strings.Join([]string{uniqueness, "Install Ansible Dependencies"}, ": "),
		&remote.CommandArgs{
			Connection: connArgs,
			Create:     pulumi.String("apt-get update && apt-get install -y python3-pip python3-jsondiff"),
		})
	if err != nil {
		return err
	}
	return nil
}

func InitDockerSwarm(ctx *pulumi.Context, connArgs remote.ConnectionArgs, advertiseAddr pulumi.StringOutput) (pulumi.Output, error) {
	var tokens SwarmJoinTokens

	fullCommand := advertiseAddr.ApplyT(func(addr string) *string {
		initCommand := fmt.Sprintf("docker swarm init --advertise-addr %s", addr)
		fullCommand := strings.Join([]string{initCommand, "echo \"Manager Token: $(docker swarm join-token -q manager)\"", "echo \"Worker Token: $(docker swarm join-token -q worker)\""}, " && ")
		return &fullCommand
	}).(pulumi.StringPtrOutput)

	out, err := remote.NewCommand(ctx, "Init docker swarm",
		&remote.CommandArgs{
			Connection: connArgs,
			Create:     fullCommand,
		})
	if err != nil {
		return pulumi.StringOutput{}, err
	}

	return out.Stdout.ApplyT(func(output string) SwarmJoinTokens {
		searchWorker := "Worker Token: "
		patternWorker := regexp.MustCompile(searchWorker + `(\S+)`)
		searchManager := "Manager Token: "
		patternManager := regexp.MustCompile(searchManager + `(\S+)`)

		matches := patternWorker.FindStringSubmatch(output)
		if len(matches) > 1 {
			extracted := matches[1]
			tokens.WorkerToken = extracted
		}
		matches = patternManager.FindStringSubmatch(output)
		if len(matches) > 1 {
			extracted := matches[1]
			tokens.ManagerToken = extracted
		}
		return tokens
	}), nil
}

func CreateAnsibleInventory(managerNodes, workerNodes []*hcloud.Server) (pulumi.Output, error) {
	serverInfos := toServerInfo(managerNodes)
	return pulumi.All(pulumi.ToOutput(serverInfos)).ApplyT(func(results []interface{}) (string, error) {
		var serverInfos = results[0].([]ServerInfo)
		// var workerSlice = results[1].([]*hcloud.Server)

		serverData := make(map[string][]ServerInfo)

		for _, s := range serverInfos {
			serverData["Manager"] = append(serverData["Manager"], ServerInfo{
				Name: s.Name,
				IP:   s.IP,
			})
		}
		// for _, result := range workerSlice {
		// 	server := result.(map[string]interface{})
		// 	serverData["Worker"] = append(serverData["Worker"], ServerInfo{
		// 		Name: server["name"].(string),
		// 		IP:   server["ipv4_address"].(string),
		// 	})
		// }
		fmt.Println(serverData["Manager"])
		fmt.Println(results[0])
		return generateInventoryFile(serverData)
	}).(pulumi.Output), nil
}

func toServerInfo(server []*hcloud.Server) pulumi.ArrayOutput {
	serverInfo := []ServerInfo{}
	for _, s := range server {
		serverInfo = append(serverInfo, ServerInfo{
			Name: s.Name,
			IP:   s.Ipv4Address,
		})
	}
	return pulumi.All(serverInfo).ApplyT(func(args []interface{}) []interface{} {
		var serverInfo []interface{}

		for _, s := range args {
			val := s.(map[string]interface{})
			serverInfo = append(serverInfo, map[string]interface{}{
				"Name": val["Name"].(string),
				"IP":   val["IP"].(string),
			})
		}
		return serverInfo
	}).(pulumi.ArrayOutput)
}

func generateInventoryFile(inventory map[string][]ServerInfo) (string, error) {
	const inventoryTmpl = `
[all]
{{ range .Manager }}
{{ .Name }} ansible_host={{ .IP }} ansible_connection=ssh ansible_user=root ansible_ssh_private_key_file=../infra-base/private_key
{{ end }}
{{ range .Worker }}
{{ .Name }} ansible_host={{ .IP }} ansible_connection=ssh ansible_user=root ansible_ssh_private_key_file=../infra-base/private_key
{{ end }}

[manager]
{{ range .Manager }}
{{ .Name }} ansible_host={{ .IP }} ansible_connection=ssh ansible_user=root ansible_ssh_private_key_file=../infra-base/private_key
{{ end }}

[worker]
{{ range .Worker }}
{{ .Name }} ansible_host={{ .IP }} ansible_connection=ssh ansible_user=root ansible_ssh_private_key_file=../infra-base/private_key
{{ end }}
	`
	tmpl, err := template.New("inventory").Parse(inventoryTmpl)
	if err != nil {
		return "", err
	}
	var buf bytes.Buffer
	err = tmpl.Execute(&buf, inventory)
	if err != nil {
		return "", err
	}
	return buf.String(), nil
}