Implement cloud-init module

windows-agent/internal/cloudinit/cloudinit.go

@@ -0,0 +1,204 @@
+// Package cloudinit has some helpers to set up cloud-init configuration.
+package cloudinit
+
+import (
+	"bytes"
+	"context"
+	"errors"
+	"fmt"
+	"io/fs"
+	"os"
+	"path/filepath"
+	"strings"
+
+	"github.com/canonical/ubuntu-pro-for-windows/windows-agent/internal/config"
+	"github.com/ubuntu/decorate"
+	"gopkg.in/ini.v1"
+	"gopkg.in/yaml.v3"
+)
+
+// Config is a configuration provider for ProToken and the Landscape config.
+type Config interface {
+	Subscription(context.Context) (string, config.Source, error)
+	LandscapeClientConfig(context.Context) (string, config.Source, error)
+}
+
+// WriteAgentData writes the agent's cloud-init data file.
+func WriteAgentData(ctx context.Context, conf *config.Config) (err error) {
+	defer decorate.OnError(&err, "could not create distro-specific cloud-init file")
+
+	cloudInit, err := marshalConfig(ctx, conf)
+	if err != nil {
+		return err
+	}
+
+	dir, err := rootPath()
+	if err != nil {
+		return fmt.Errorf("could not compute path: %v", err)
+	}
+
+	if err := os.MkdirAll(filepath.Dir(dir), 0400); err != nil {
+		return fmt.Errorf("could not create directory: %v", err)
+	}
+
+	path := filepath.Join(dir, "agent.yaml")
+	tmp := path + ".tmp"
+
+	if err := os.WriteFile(tmp, cloudInit, 0400); err != nil {
+		return fmt.Errorf("could not write: %v", err)
+	}
+
+	if err := os.Rename(tmp, path); err != nil {
+		return fmt.Errorf("could not rename: %v", err)
+	}
+
+	return nil
+}
+
+// WriteDistroData writes cloud-init user data to be used for a distro in particular.
+func WriteDistroData(distroName string, source string, cloudInit string) (err error) {
+	defer decorate.OnError(&err, "could not create distro-specific cloud-init file")
+
+	path, err := distroPath(distroName, source)
+	if err != nil {
+		return fmt.Errorf("could not compute path: %v", err)
+	}
+
+	if err := os.MkdirAll(filepath.Dir(path), 0400); err != nil {
+		return fmt.Errorf("could not create directory: %v", err)
+	}
+
+	tmp := path + ".tmp"
+	if err := os.WriteFile(tmp, []byte(cloudInit), 0400); err != nil {
+		return fmt.Errorf("could not write: %v", err)
+	}
+
+	if err := os.Rename(tmp, path); err != nil {
+		_ = os.Remove(tmp)
+		return fmt.Errorf("could not rename: %v", err)
+	}
+
+	return nil
+}
+
+// RemoveDistroData removes cloud-init user data to be used for a distro in particular.
+//
+// No error is returned if the data did not exist.
+func RemoveDistroData(distroName string, source string) (err error) {
+	defer decorate.OnError(&err, "could not remove distro-specific cloud-init file")
+
+	path, err := distroPath(distroName, source)
+	if err != nil {
+		return fmt.Errorf("could not compute path: %v", err)
+	}
+
+	err = os.Remove(path)
+	if errors.Is(err, fs.ErrNotExist) {
+		return nil
+	}
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+func rootPath() (string, error) {
+	homeDir := os.Getenv("USERPROFILE")
+	if homeDir == "" {
+		return "", errors.New("environment variable USERPROFILE is empty")
+	}
+
+	return filepath.Join(homeDir, ".ubuntupro"), nil
+}
+
+func distroPath(distroName string, source string) (string, error) {
+	root, err := rootPath()
+	if err != nil {
+		return "", err
+	}
+
+	dir := filepath.Join(root, ".ubuntupro", source)
+	base := fmt.Sprintf("%s.%s", distroName, ".user-data")
+	return filepath.Join(dir, base), nil
+}
+
+func marshalConfig(ctx context.Context, conf Config) ([]byte, error) {
+	w := &bytes.Buffer{}
+
+	if _, err := fmt.Fprintln(w, "# cloud-init"); err != nil {
+		return nil, fmt.Errorf("could not write # cloud-init stenza: %v", err)
+	}
+
+	if _, err := fmt.Fprintln(w, "# This file was generated automatically and must not be edited"); err != nil {
+		return nil, fmt.Errorf("could not write warning message: %v", err)
+	}
+
+	contents := make(map[string]interface{})
+
+	if err := ubuntuAdvantageModule(ctx, conf, contents); err != nil {
+		return nil, err
+	}
+
+	if err := landscapeModule(ctx, conf, contents); err != nil {
+		return nil, err
+	}
+
+	out, err := yaml.Marshal(contents)
+	if err != nil {
+		return nil, fmt.Errorf("could not Marshal user data as a YAML: %v", err)
+	}
+
+	if _, err := w.Write(out); err != nil {
+		return nil, fmt.Errorf("could not write config body: %v", err)
+	}
+
+	return w.Bytes(), nil
+}
+
+func ubuntuAdvantageModule(ctx context.Context, c Config, out map[string]interface{}) error {
+	token, src, err := c.Subscription(ctx)
+	if err != nil {
+		return err
+	}
+	if src == config.SourceNone {
+		return nil
+	}
+
+	type uaModule struct {
+		Token string `yaml:"token"`
+	}
+
+	out["ubuntu_advantage"] = uaModule{Token: token}
+	return nil
+}
+
+func landscapeModule(ctx context.Context, c Config, out map[string]interface{}) error {
+	conf, src, err := c.LandscapeClientConfig(ctx)
+	if err != nil {
+		return err
+	}
+	if src == config.SourceNone {
+		return nil
+	}
+
+	var landcapeModule struct {
+		Client map[string]string `yaml:"client"`
+	}
+
+	f, err := ini.Load(strings.NewReader(conf))
+	if err != nil {
+		return fmt.Errorf("could not load Landscape configuration file")
+	}
+
+	section, err := f.GetSection("client")
+	if err != nil {
+		return nil // Empty section
+	}
+
+	for _, keyName := range section.KeyStrings() {
+		landcapeModule.Client[keyName] = section.Key(keyName).String()
+	}
+
+	out["landscape"] = landcapeModule
+	return nil
+}