From e1d9b6653e031f28f296cb76c7ae096ffd445435 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Anders=20F=20Bj=C3=B6rklund?= <anders.f.bjorklund@gmail.com>
Date: Sun, 7 Jan 2024 22:27:26 +0100
Subject: [PATCH] Add external support for provision scripts
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Anders F Björklund <anders.f.bjorklund@gmail.com>
---
 pkg/hostagent/hostagent.go |  5 +++
 pkg/hostagent/provision.go | 69 ++++++++++++++++++++++++++++++++++++++
 pkg/hostagent/sudo.go      | 43 ++++++++++++++++++++++++
 3 files changed, 117 insertions(+)
 create mode 100644 pkg/hostagent/provision.go
 create mode 100644 pkg/hostagent/sudo.go

diff --git a/pkg/hostagent/hostagent.go b/pkg/hostagent/hostagent.go
index 184ba03ff248..b4cbe33ac226 100644
--- a/pkg/hostagent/hostagent.go
+++ b/pkg/hostagent/hostagent.go
@@ -517,6 +517,11 @@ sudo chown -R "${USER}" /run/host-services`
 			return errors.Join(unlockErrs...)
 		})
 	}
+	if *a.y.VMType == limayaml.EXT {
+		if err := a.runProvisionScripts(); err != nil {
+			return err
+		}
+	}
 	if !*a.y.Plain {
 		go a.watchGuestAgentEvents(ctx)
 	}
diff --git a/pkg/hostagent/provision.go b/pkg/hostagent/provision.go
new file mode 100644
index 000000000000..438d06b1eac7
--- /dev/null
+++ b/pkg/hostagent/provision.go
@@ -0,0 +1,69 @@
+package hostagent
+
+import (
+	"errors"
+	"fmt"
+
+	"github.com/lima-vm/lima/pkg/limayaml"
+	"github.com/lima-vm/sshocker/pkg/ssh"
+	"github.com/sirupsen/logrus"
+)
+
+func (a *HostAgent) runProvisionScripts() error {
+	var errs []error
+
+	for i, f := range a.y.Provision {
+		switch f.Mode {
+		case limayaml.ProvisionModeSystem, limayaml.ProvisionModeUser:
+			logrus.Infof("Running %s provision %d of %d", f.Mode, i+1, len(a.y.Provision))
+			err := a.waitForProvision(
+				provision{
+					description: fmt.Sprintf("provision.%s/%08d", f.Mode, i),
+					sudo:        f.Mode == limayaml.ProvisionModeSystem,
+					script:      f.Script,
+				})
+			if err != nil {
+				errs = append(errs, err)
+			}
+		case limayaml.ProvisionModeDependency, limayaml.ProvisionModeBoot:
+			logrus.Infof("Skipping %s provision %d of %d", f.Mode, i+1, len(a.y.Provision))
+			continue
+		default:
+			return fmt.Errorf("unknown provision mode %q", f.Mode)
+		}
+	}
+	return errors.Join(errs...)
+}
+
+func (a *HostAgent) waitForProvision(p provision) error {
+	if p.sudo {
+		return a.waitForSystemProvision(p)
+	}
+	return a.waitForUserProvision(p)
+}
+
+func (a *HostAgent) waitForSystemProvision(p provision) error {
+	logrus.Debugf("executing script %q", p.description)
+	stdout, stderr, err := sudoExecuteScript(a.instSSHAddress, a.sshLocalPort, a.sshConfig, p.script, p.description)
+	logrus.Debugf("stdout=%q, stderr=%q, err=%v", stdout, stderr, err)
+	if err != nil {
+		return fmt.Errorf("stdout=%q, stderr=%q: %w", stdout, stderr, err)
+	}
+	return nil
+}
+
+func (a *HostAgent) waitForUserProvision(p provision) error {
+	logrus.Debugf("executing script %q", p.description)
+	stdout, stderr, err := ssh.ExecuteScript(a.instSSHAddress, a.sshLocalPort, a.sshConfig, p.script, p.description)
+	logrus.Debugf("stdout=%q, stderr=%q, err=%v", stdout, stderr, err)
+	if err != nil {
+		return fmt.Errorf("stdout=%q, stderr=%q: %w", stdout, stderr, err)
+	}
+	return nil
+}
+
+type provision struct {
+	description string
+	script      string
+	sudo        bool
+}
diff --git a/pkg/hostagent/sudo.go b/pkg/hostagent/sudo.go
new file mode 100644
index 000000000000..cae95c0c7e64
--- /dev/null
+++ b/pkg/hostagent/sudo.go
@@ -0,0 +1,43 @@
+package hostagent
+
+import (
+	"bytes"
+	"errors"
+	"fmt"
+	"os/exec"
+	"strconv"
+	"strings"
+
+	"github.com/lima-vm/sshocker/pkg/ssh"
+	"github.com/sirupsen/logrus"
+)
+
+// sudoExecuteScript executes the given script (as root) on the remote host via stdin.
+// Returns stdout and stderr.
+//
+// scriptName is used only for readability of error strings.
+func sudoExecuteScript(host string, port int, c *ssh.SSHConfig, script, scriptName string) (stdout, stderr string, err error) {
+	if c == nil {
+		return "", "", errors.New("got nil SSHConfig")
+	}
+	interpreter, err := ssh.ParseScriptInterpreter(script)
+	if err != nil {
+		return "", "", err
+	}
+	sshBinary := c.Binary()
+	sshArgs := c.Args()
+	if port != 0 {
+		sshArgs = append(sshArgs, "-p", strconv.Itoa(port))
+	}
+	sshArgs = append(sshArgs, host, "--", "sudo", interpreter)
+	sshCmd := exec.Command(sshBinary, sshArgs...)
+	sshCmd.Stdin = strings.NewReader(script)
+	var buf bytes.Buffer
+	sshCmd.Stderr = &buf
+	logrus.Debugf("executing ssh for script %q: %s %v", scriptName, sshCmd.Path, sshCmd.Args)
+	out, err := sshCmd.Output()
+	if err != nil {
+		return string(out), buf.String(), fmt.Errorf("failed to execute script %q: stdout=%q, stderr=%q: %w", scriptName, string(out), buf.String(), err)
+	}
+	return string(out), buf.String(), nil
+}