From c7a623b26b621be0bda10a7a94f4e4aaf48fe53a Mon Sep 17 00:00:00 2001
From: r33drichards <xxmrmau5@gmail.com>
Date: Wed, 3 Nov 2021 21:36:47 +0000
Subject: [PATCH] add port flag to link command

---
 pkg/cmd/link/link.go           | 30 +++++++++++++++++++-----------
 pkg/portforward/portforward.go |  4 ++--
 2 files changed, 21 insertions(+), 13 deletions(-)

diff --git a/pkg/cmd/link/link.go b/pkg/cmd/link/link.go
index f0fd033f..1d8d31aa 100644
--- a/pkg/cmd/link/link.go
+++ b/pkg/cmd/link/link.go
@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"os"
+	"strings"
 
 	"github.com/brevdev/brev-cli/pkg/brev_api"
 	"github.com/brevdev/brev-cli/pkg/config"
@@ -18,12 +19,13 @@ import (
 	"github.com/spf13/cobra"
 )
 
-var sshLinkLong = "Enable a local ssh tunnel, setup private key auth, and give connection string"
-
-var sshLinkExample = "brev link <ws_name>"
+var (
+	Port           string
+	sshLinkLong    = "Enable a local ssh tunnel, setup private key auth, and give connection string"
+	sshLinkExample = "brev link <ws_name>"
+)
 
 func getWorkspaceNames() []string {
-
 	activeOrg, err := brev_api.GetActiveOrgContext()
 	if err != nil {
 		return nil
@@ -38,7 +40,7 @@ func getWorkspaceNames() []string {
 		return nil
 	}
 
-	var wsNames []string;
+	var wsNames []string
 	for _, w := range wss {
 		wsNames = append(wsNames, w.Name)
 	}
@@ -77,15 +79,21 @@ func NewCmdLink(t *terminal.Terminal) *cobra.Command {
 				t.Errprint(err, "")
 			}
 			sshPrivateKeyFilePath := files.GetSSHPrivateKeyFilePath()
-			cmdutil.CheckErr(opts.Complete(cmd, t, args))
+			if Port == "" {
+				Port = "2222:22"
+			}
+			cmdutil.CheckErr(opts.Complete(cmd, t, args, Port))
 			t.Printf("SSH Private Key: %s\n", sshPrivateKeyFilePath)
 			t.Printf(t.Green("\n\t1. Add SSH Key:\n"))
 			t.Printf(t.Yellow("\t\tssh-add %s\n", sshPrivateKeyFilePath))
 			t.Printf(t.Green("\t2. Connect to workspace:\n"))
-			t.Printf(t.Yellow("\t\tssh -p 2222 brev@0.0.0.0\n\n"))
+			localPort := strings.Split(Port, ":")[0]
+			t.Printf(t.Yellow("\t\tssh -p %s brev@0.0.0.0\n\n", localPort))
 			cmdutil.CheckErr(opts.RunPortforward())
 		},
 	}
+	cmd.Flags().StringVarP(&Port, "port", "p", "", "port forward flag describe me better")
+
 	return cmd
 }
 
@@ -105,7 +113,7 @@ func (d WorkspaceResolver) GetWorkspaceByID(id string) (*brev_api.AllWorkspaceDa
 		return nil, err
 	}
 
-	return &brev_api.AllWorkspaceData{WorkspaceMetaData: *wmeta, Workspace:*w}, nil
+	return &brev_api.AllWorkspaceData{WorkspaceMetaData: *wmeta, Workspace: *w}, nil
 }
 
 // This function will be long and messy, it's entirely built to check random error cases
@@ -132,7 +140,7 @@ func (d WorkspaceResolver) GetWorkspaceByName(name string) (*brev_api.AllWorkspa
 				if err != nil {
 					return nil, err
 				}
-				return &brev_api.AllWorkspaceData{WorkspaceMetaData: *wmeta, Workspace:w}, nil
+				return &brev_api.AllWorkspaceData{WorkspaceMetaData: *wmeta, Workspace: w}, nil
 			}
 		}
 		// if there wasn't a workspace in the org, check all the orgs
@@ -150,13 +158,13 @@ func (d WorkspaceResolver) GetWorkspaceByName(name string) (*brev_api.AllWorkspa
 		}
 
 		for _, w := range workspaces {
-			if w.Name==name {
+			if w.Name == name {
 				// Assemble full object
 				wmeta, err := c.GetWorkspaceMetaData(w.ID)
 				if err != nil {
 					return nil, err
 				}
-				return &brev_api.AllWorkspaceData{WorkspaceMetaData: *wmeta, Workspace:w}, nil
+				return &brev_api.AllWorkspaceData{WorkspaceMetaData: *wmeta, Workspace: w}, nil
 			}
 		}
 	}
diff --git a/pkg/portforward/portforward.go b/pkg/portforward/portforward.go
index a94cc206..fc57098b 100644
--- a/pkg/portforward/portforward.go
+++ b/pkg/portforward/portforward.go
@@ -49,7 +49,7 @@ func NewPortForwardOptions(portForwardHelpers k8s.K8sClient, workspaceResolver R
 	}
 }
 
-func (o *PortForwardOptions) Complete(cmd *cobra.Command, t *terminal.Terminal, args []string) error {
+func (o *PortForwardOptions) Complete(cmd *cobra.Command, t *terminal.Terminal, args []string, Port string) error {
 	workspaceIDOrName := args[0]
 
 	workspace, err := o.ResourceResolver.GetWorkspaceByID(workspaceIDOrName)
@@ -74,7 +74,7 @@ func (o *PortForwardOptions) Complete(cmd *cobra.Command, t *terminal.Terminal,
 	}
 
 	o.Address = []string{"localhost"}
-	o.Ports = []string{"2222:22"} // TODO override from args
+	o.Ports = []string{Port} // TODO override from args
 
 	o.StopChannel = make(chan struct{}, 1)
 	o.ReadyChannel = make(chan struct{})