From be0bd1c9c1571a57856467d441a83746f6bf888b Mon Sep 17 00:00:00 2001
From: Tyler Fong <tyler@brev.dev>
Date: Sun, 25 Feb 2024 04:50:17 +0000
Subject: [PATCH] added --host to shell and open

---
 pkg/cmd/open/open.go   |  9 +++++++--
 pkg/cmd/shell/shell.go | 15 +++++++++++----
 2 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/pkg/cmd/open/open.go b/pkg/cmd/open/open.go
index ab871aa8..187b4c29 100644
--- a/pkg/cmd/open/open.go
+++ b/pkg/cmd/open/open.go
@@ -48,6 +48,7 @@ type OpenStore interface {
 func NewCmdOpen(t *terminal.Terminal, store OpenStore, noLoginStartStore OpenStore) *cobra.Command {
 	var waitForSetupToFinish bool
 	var directory string
+	var host bool
 
 	cmd := &cobra.Command{
 		Annotations:           map[string]string{"ssh": ""},
@@ -63,13 +64,14 @@ func NewCmdOpen(t *terminal.Terminal, store OpenStore, noLoginStartStore OpenSto
 			if waitForSetupToFinish {
 				setupDoneString = "------ Done running execs ------"
 			}
-			err := runOpenCommand(t, store, args[0], setupDoneString, directory)
+			err := runOpenCommand(t, store, args[0], setupDoneString, directory, host)
 			if err != nil {
 				return breverrors.WrapAndTrace(err)
 			}
 			return nil
 		},
 	}
+	cmd.Flags().BoolVarP(&host, "host", "", false, "ssh into the host machine instead of the container")
 	cmd.Flags().BoolVarP(&waitForSetupToFinish, "wait", "w", false, "wait for setup to finish")
 	cmd.Flags().StringVarP(&directory, "dir", "d", "", "directory to open")
 
@@ -77,7 +79,7 @@ func NewCmdOpen(t *terminal.Terminal, store OpenStore, noLoginStartStore OpenSto
 }
 
 // Fetch workspace info, then open code editor
-func runOpenCommand(t *terminal.Terminal, tstore OpenStore, wsIDOrName string, setupDoneString string, directory string) error {
+func runOpenCommand(t *terminal.Terminal, tstore OpenStore, wsIDOrName string, setupDoneString string, directory string, host bool) error {
 	// todo check if workspace is stopped and start if it if it is stopped
 	fmt.Println("finding your instance...")
 	res := refresh.RunRefreshAsync(tstore)
@@ -112,6 +114,9 @@ func runOpenCommand(t *terminal.Terminal, tstore OpenStore, wsIDOrName string, s
 	}
 
 	localIdentifier := workspace.GetLocalIdentifier()
+	if host {
+		localIdentifier += "-host"
+	}
 
 	err = res.Await()
 	if err != nil {
diff --git a/pkg/cmd/shell/shell.go b/pkg/cmd/shell/shell.go
index f964c27a..1c22d2e8 100644
--- a/pkg/cmd/shell/shell.go
+++ b/pkg/cmd/shell/shell.go
@@ -42,7 +42,7 @@ type ShellStore interface {
 func NewCmdShell(t *terminal.Terminal, store ShellStore, noLoginStartStore ShellStore) *cobra.Command {
 	var runRemoteCMD bool
 	var directory string
-
+	var host bool
 	cmd := &cobra.Command{
 		Annotations:           map[string]string{"ssh": ""},
 		Use:                   "shell",
@@ -54,20 +54,21 @@ func NewCmdShell(t *terminal.Terminal, store ShellStore, noLoginStartStore Shell
 		Args:                  cmderrors.TransformToValidationError(cmderrors.TransformToValidationError(cobra.ExactArgs(1))),
 		ValidArgsFunction:     completions.GetAllWorkspaceNameCompletionHandler(noLoginStartStore, t),
 		RunE: func(cmd *cobra.Command, args []string) error {
-			err := runShellCommand(t, store, args[0], directory)
+			err := runShellCommand(t, store, args[0], directory, host)
 			if err != nil {
 				return breverrors.WrapAndTrace(err)
 			}
 			return nil
 		},
 	}
+	cmd.Flags().BoolVarP(&host, "host", "", false, "ssh into the host machine instead of the container")
 	cmd.Flags().BoolVarP(&runRemoteCMD, "remote", "r", true, "run remote commands")
 	cmd.Flags().StringVarP(&directory, "dir", "d", "", "override directory to launch shell")
 
 	return cmd
 }
 
-func runShellCommand(t *terminal.Terminal, sstore ShellStore, workspaceNameOrID, directory string) error {
+func runShellCommand(t *terminal.Terminal, sstore ShellStore, workspaceNameOrID, directory string, host bool) error {
 	s := t.NewSpinner()
 	workspace, err := util.GetUserWorkspaceByNameOrIDErr(sstore, workspaceNameOrID)
 	if err != nil {
@@ -93,7 +94,13 @@ func runShellCommand(t *terminal.Terminal, sstore ShellStore, workspaceNameOrID,
 	if workspace.Status != "RUNNING" {
 		return breverrors.New("Workspace is not running")
 	}
-	sshName := string(workspace.GetLocalIdentifier())
+
+	localIdentifier := workspace.GetLocalIdentifier()
+	if host {
+		localIdentifier += "-host"
+	}
+
+	sshName := string(localIdentifier)
 
 	err = refreshRes.Await()
 	if err != nil {