src/cmd/cos_kernel/actions/unlock.go

package actions

import (
	"context"
	_ "embed"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"time"

	"github.com/golang/glog"
	"github.com/spf13/cobra"
)

//go:embed scripts/disable_verified_boot.sh
var disableVerifiedBootScript string

//go:embed scripts/unlock_root_fs.sh
var unlockRootFSScript string

func unlockGCPInstance(ctx context.Context, args []string) error {
	glog.Info("checking gcp instance verified boot status")

	disableArgs := append(args, "--command", fmt.Sprintf("sudo bash -c '%s'", disableVerifiedBootScript))
	_, err := runGcloudComputeSSH(ctx, disableArgs, nil)


	if err != nil {
		exitErr, ok := errors.AsType[*exec.ExitError](err)
		if !ok || exitErr.ExitCode() != 10 {
			return fmt.Errorf("failed to disable verified boot on remote instance: %w", err)
		}

		glog.Info("rebooting remote instance to apply grub changes")
		rebootArgs := append(args, "--command", "sudo bash -c 'sleep 1 && echo b > /proc/sysrq-trigger' > /dev/null 2>&1 &")

		_, err = runGcloudComputeSSH(ctx, rebootArgs, &SSHOptions{
			MaxRetries: 3,
			RetryDelay: 2 * time.Second,
			Timeout:    10 * time.Second,
		})
		if err != nil {
			return fmt.Errorf("failed to trigger reboot: %w", err)
		}

		glog.Info("waiting for instance to become available")
		waitArgs := append(args, "--command", "echo ready", "--ssh-flag=-o ConnectTimeout=5")
		if _, err = runGcloudComputeSSH(ctx, waitArgs, &SSHOptions{
			MaxRetries: 5,
			RetryDelay: 10 * time.Second,
			Timeout:    30 * time.Second,
		}); err != nil {
			return fmt.Errorf("instance failed to become available: %w", err)
		}
	} else {
		glog.Info("verified boot is already disabled.")
	}

	glog.Info("unlocking gcp instance root fs")

	unlockArgs := append(args, "--command", fmt.Sprintf("sudo bash -c '%s'", unlockRootFSScript))
	_, err = runGcloudComputeSSH(ctx, unlockArgs, nil)
	if err != nil {
		exitErr, ok := errors.AsType[*exec.ExitError](err)
		if !ok || exitErr.ExitCode() != 10 {
			return fmt.Errorf("failed to unlock root fs on remote instance: %w", err)
		}
		glog.Info("root fs unlocked")
		return nil
	}

	glog.Info("root fs already unlocked")
	return nil
}

func unlockLocalInstance() error {
	glog.Info("checking verified boot status")

	_, err := executeScript(disableVerifiedBootScript)
	if err != nil {
		exitErr, ok := errors.AsType[*exec.ExitError](err)
		if !ok || exitErr.ExitCode() != 10 {
			return fmt.Errorf("failed to disable verified boot: %w", err)
		}

		glog.Info("rebooting to apply grub changes")
		if err := os.WriteFile("/proc/sysrq-trigger", []byte("b"), 0644); err != nil {
			return fmt.Errorf("failed to trigger reboot: %w", err)
		}
		return nil
	}

	glog.Info("verified boot is already disabled")
	glog.Info("unlocking local instance root fs")

	_, err = executeScript(unlockRootFSScript)
	if err != nil {
		exitErr, ok := errors.AsType[*exec.ExitError](err)
		if !ok || exitErr.ExitCode() != 10 {
			return fmt.Errorf("failed to unlock root fs: %w", err)
		}
		glog.Info("root fs unlocked")
		return nil
	}

	glog.Info("root fs already unlocked")
	return nil
}

var unlockCmd = &cobra.Command{
	Use:   "unlock",
	Short: "Unlocks a local or remote GCP instance to allow development",
	Long: `Executes a bash script to set up either a local instance or remote GCP instance for
development. If no arguments are given, then the script is run to unlock the local instance. If 
arguments are given, then the arguments are passed into gcloud compute ssh to be executed on`,
	RunE: func(cmd *cobra.Command, args []string) error {
		if len(args) == 0 {
			return unlockLocalInstance()
		} else {
			return unlockGCPInstance(cmd.Context(), args)
		}
	},
}

func init() {
	rootCmd.AddCommand(unlockCmd)
}