shell.go

  1package shell
  2
  3import (
  4	"bufio"
  5	"context"
  6	"fmt"
  7	"io"
  8	"os"
  9	"runtime"
 10	"strings"
 11
 12	"github.com/anmitsu/go-shlex"
 13	"github.com/charmbracelet/soft-serve/cmd"
 14	"github.com/charmbracelet/soft-serve/pkg/shell"
 15	sshcmd "github.com/charmbracelet/soft-serve/pkg/ssh/cmd"
 16	"github.com/charmbracelet/soft-serve/pkg/sshutils"
 17	"github.com/charmbracelet/ssh"
 18	"github.com/mattn/go-tty"
 19	"github.com/muesli/termenv"
 20	"github.com/spf13/cobra"
 21)
 22
 23var (
 24	commandString string
 25
 26	// Command is a login shell command.
 27	Command = &cobra.Command{
 28		Use:                "shell",
 29		SilenceUsage:       true,
 30		PersistentPreRunE:  cmd.InitBackendContext,
 31		PersistentPostRunE: cmd.CloseDBContext,
 32		Args:               cobra.NoArgs,
 33		RunE: func(cmd *cobra.Command, _ []string) error {
 34			args, err := shlex.Split(commandString, true)
 35			if err != nil {
 36				return err
 37			}
 38
 39			return runShell(cmd, args)
 40		},
 41	}
 42)
 43
 44func init() {
 45	Command.CompletionOptions.DisableDefaultCmd = true
 46	Command.SetUsageTemplate(sshcmd.UsageTemplate)
 47	Command.SetUsageFunc(sshcmd.UsageFunc)
 48	Command.Flags().StringVarP(&commandString, "", "c", "", "Command to run")
 49}
 50
 51func runShell(cmd *cobra.Command, args []string) error {
 52	ctx := cmd.Context()
 53	sshTty, isInteractive := os.LookupEnv("SSH_TTY")
 54	sshUserAuth := os.Getenv("SSH_USER_AUTH")
 55
 56	var ak string
 57	if sshUserAuth != "" {
 58		f, err := os.Open(sshUserAuth)
 59		if err != nil {
 60			return fmt.Errorf("could not open SSH_USER_AUTH file: %w", err)
 61		}
 62
 63		ak = parseSSHUserAuth(f)
 64		f.Close() // nolint: errcheck
 65	}
 66
 67	if err := os.Setenv("SOFT_SERVE_PUBLIC_KEY", ak); err != nil {
 68		return fmt.Errorf("could not set SOFT_SERVE_PUBLIC_KEY: %w", err)
 69	}
 70
 71	var pk ssh.PublicKey
 72	var err error
 73	if ak != "" {
 74		pk, _, err = sshutils.ParseAuthorizedKey(ak)
 75		if err != nil {
 76			return fmt.Errorf("could not parse authorized key: %w", err)
 77		}
 78	}
 79
 80	// We need a public key in the context even if it's nil
 81	ctx = context.WithValue(ctx, ssh.ContextKeyPublicKey, pk)
 82
 83	in, out, er := os.Stdin, os.Stdout, os.Stderr
 84	if isInteractive {
 85		switch runtime.GOOS {
 86		case "windows":
 87			tty, err := tty.Open()
 88			if err != nil {
 89				return fmt.Errorf("could not open tty: %w", err)
 90			}
 91
 92			in = tty.Input()
 93			out = tty.Output()
 94			er = tty.Output()
 95		default:
 96			var err error
 97			in, err = os.Open(sshTty)
 98			if err != nil {
 99				return fmt.Errorf("could not open input tty: %w", err)
100			}
101
102			out, err = os.OpenFile(sshTty, os.O_WRONLY, 0)
103			if err != nil {
104				return fmt.Errorf("could not open output tty: %w", err)
105			}
106			er = out
107		}
108	}
109
110	c := shell.Command(ctx, osEnv, isInteractive)
111
112	c.SetArgs(args)
113	c.SetIn(in)
114	c.SetOut(out)
115	c.SetErr(er)
116	c.SetContext(ctx)
117
118	return c.ExecuteContext(ctx)
119}
120
121func parseSSHUserAuth(r io.Reader) string {
122	scanner := bufio.NewScanner(r)
123	for scanner.Scan() {
124		line := scanner.Text()
125		if line == "" {
126			continue
127		}
128
129		if strings.HasPrefix(line, "publickey ") {
130			return strings.TrimPrefix(line, "publickey ")
131		}
132	}
133
134	return ""
135}
136
137var osEnv = &osEnviron{}
138
139type osEnviron struct{}
140
141var _ termenv.Environ = &osEnviron{}
142
143// Environ implements termenv.Environ.
144func (*osEnviron) Environ() []string {
145	return os.Environ()
146}
147
148// Getenv implements termenv.Environ.
149func (*osEnviron) Getenv(key string) string {
150	return os.Getenv(key)
151}