1 files changed, 12 insertions, 8 deletions

diff --git a/cmd/client/client.go b/cmd/client/client.go
index 62f1f48..7799b87 100644
--- a/cmd/client/client.go
+++ b/cmd/client/client.go
@@ -26,19 +26,23 @@ import (
 // This should be compiled into the binary
 var clientId string
 
-func Register(root *cobra.Command) {
+func NewClientCommand() *cobra.Command {
         clientCmd := &cobra.Command{
                 Use:   "client proxy-host ssh-to-host ssh-port username",
                 Short: "Run websocket client",
-                Args:  cobra.ExactArgs(4),
+                Args:  cobra.ExactArgs(3),
                 Run: func(c *cobra.Command, args []string) {
                         cfg := app.Config{}
                         cli.MustGetConfig(c, &cfg)
-                        clientMain(cfg, args[0], args[1], args[2], args[3])
+                        clientMain(cfg, args[0], args[1], args[2])
                 },
         }
         cli.AddFlags(clientCmd, &app.Config{}, app.DefaultConfig, "client")
-        root.AddCommand(clientCmd)
+        return clientCmd
+}
+
+func Register(root *cobra.Command) {
+        root.AddCommand(NewClientCommand())
 }
 
 func generateCertificateRequest(username, host string) (ed25519.PrivateKey, []byte, error) {
@@ -181,13 +185,13 @@ func fetchOauthToken(ctx context.Context, clientId, proxyHost string) (string, e
         return tokenResponse.AccessToken, nil
 }
 
-func clientMain(cfg app.Config, proxyHost, host, port, username string) {
+func clientMain(cfg app.Config, host, port, username string) {
         log.SetOutput(os.Stderr)
 
         ctx, cancel := context.WithCancel(context.Background())
         defer cancel()
 
-        oauthToken, err := fetchOauthToken(ctx, clientId, proxyHost)
+        oauthToken, err := fetchOauthToken(ctx, clientId, cfg.ClientHost)
         if err != nil {
                 log.Fatalf("Error fetching oauth token: %s", err)
         }
@@ -197,7 +201,7 @@ func clientMain(cfg app.Config, proxyHost, host, port, username string) {
                 log.Fatalf("Error generating certificate request: %s", err)
         }
 
-        certificate, err := getCertificateFromCA(ctx, oauthToken, certRequest, proxyHost)
+        certificate, err := getCertificateFromCA(ctx, oauthToken, certRequest, cfg.ClientHost)
         if err != nil {
                 log.Fatalf("Error fetching certificate: %s", err)
         }
@@ -206,7 +210,7 @@ func clientMain(cfg app.Config, proxyHost, host, port, username string) {
                 log.Fatalf("Error adding certificate to agent: %s", err)
         }
 
-        ws, err := dialProxyHost(ctx, oauthToken, proxyHost, host, port)
+        ws, err := dialProxyHost(ctx, oauthToken, cfg.ClientHost, host, port)
         if err != nil {
                 log.Fatalf("Error dialing proxy host: %s", err)
         }