feat: update the admin k9s command to add the connection to the bastion

Qovery · Oct 20, 2024 · 557af51 · 557af51
1 parent 9e081b5
commit 557af51
Show file tree

Hide file tree

Showing 2 changed files with 109 additions and 11 deletions.
diff --git a/cmd/admin_k9s.go b/cmd/admin_k9s.go
@@ -1,15 +1,22 @@
 package cmd
 
 import (
+	"context"
+	"fmt"
+	"net"
 	"os"
 	"os/exec"
+	"syscall"
+	"time"
 
 	"github.com/qovery/qovery-cli/pkg"
 	"github.com/qovery/qovery-cli/utils"
 	log "github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
 )
 
+var doNotConnectToBastion bool
+
 var k9sCmd = &cobra.Command{
 	Use:   "k9s",
 	Short: "Launch k9s with a cluster ID",
@@ -20,6 +27,7 @@ var k9sCmd = &cobra.Command{
 
 func init() {
 	adminCmd.AddCommand(k9sCmd)
+	k9sCmd.Flags().BoolVarP(&doNotConnectToBastion, "no-bastion", "", false, "do not connect to the bastion")
 }
 
 func launchK9s(args []string) {
@@ -30,9 +38,21 @@ func launchK9s(args []string) {
 		return
 	}
 
+	if !doNotConnectToBastion {
+		ctx, cancel := context.WithCancel(context.Background())
+		defer cancel()
+
+		sshCmd, err := setupSSHConnection(ctx)
+		if err != nil {
+			log.Errorf("Failed to kill SSH process: %v", err)
+			// continue anyway
+		}
+		defer cleanupSSHConnection(sshCmd)
+	}
+
 	clusterId := args[0]
-	vars := pkg.GetVarsByClusterId(clusterId)
-	if len(vars) == 0 {
+	vars, err := pkg.GetVarsByClusterId(clusterId)
+	if len(vars) == 0 || err != nil {
 		return
 	}
 
@@ -64,7 +84,7 @@ func launchK9s(args []string) {
 	cmd.Stdin = os.Stdin
 	cmd.Stderr = os.Stderr
 
-	err := cmd.Run()
+	err = cmd.Run()
 	if err != nil {
 		log.Error("Can't launch k9s : " + err.Error())
 	}
@@ -85,3 +105,82 @@ func checkEnv() {
 		panic("unreachable") // staticcheck false positive: https://staticcheck.io/docs/checks#SA5011
 	}
 }
+
+func setupSSHConnection(ctx context.Context) (*exec.Cmd, error) {
+	sshArgs := []string{
+		"-N", "-D", "1080",
+		"-o", "ServerAliveInterval=10",
+		"-o", "ServerAliveCountMax=3",
+		"-o", "TCPKeepAlive=yes",
+		"[email protected]",
+		"-p", "2222",
+	}
+
+	sshCmd := exec.CommandContext(ctx, "ssh", sshArgs...)
+	if err := sshCmd.Start(); err != nil {
+		return nil, fmt.Errorf("error starting SSH command: %w", err)
+	}
+
+	if err := waitForSSHConnection(ctx, "localhost:1080", 30*time.Second); err != nil {
+		err := sshCmd.Process.Kill()
+		if err != nil {
+			return nil, err
+		}
+		return nil, fmt.Errorf("error waiting for SSH connection: %w", err)
+	}
+
+	log.Info("SSH connection established successfully")
+	if err := os.Setenv("HTTPS_PROXY", "socks5://localhost:1080"); err != nil {
+		err := sshCmd.Process.Kill()
+		if err != nil {
+			return nil, err
+		}
+		return nil, fmt.Errorf("failed to set HTTPS_PROXY: %w", err)
+	}
+
+	return sshCmd, nil
+}
+
+func cleanupSSHConnection(sshCmd *exec.Cmd) {
+	if sshCmd != nil && sshCmd.Process != nil {
+		log.Info("Terminating SSH process...")
+		if err := sshCmd.Process.Signal(syscall.SIGTERM); err != nil {
+			log.Errorf("Failed to terminate SSH process: %v", err)
+			if err := sshCmd.Process.Kill(); err != nil {
+				log.Errorf("Failed to kill SSH process: %v", err)
+			}
+		}
+		_, _ = sshCmd.Process.Wait()
+		log.Info("SSH process terminated")
+	}
+
+	if err := os.Unsetenv("HTTPS_PROXY"); err != nil {
+		log.Errorf("Failed to unset HTTPS_PROXY: %v", err)
+	} else {
+		log.Info("HTTPS_PROXY has been unset")
+	}
+}
+
+func waitForSSHConnection(ctx context.Context, address string, timeout time.Duration) error {
+	ticker := time.NewTicker(time.Second)
+	defer ticker.Stop()
+
+	timeoutChan := time.After(timeout)
+
+	for {
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case <-timeoutChan:
+			return fmt.Errorf("timeout waiting for SSH connection")
+		case <-ticker.C:
+			if conn, err := net.DialTimeout("tcp", address, time.Second); err == nil {
+				err := conn.Close()
+				if err != nil {
+					return err
+				}
+				return nil
+			}
+		}
+	}
+}
diff --git a/pkg/vault.go b/pkg/vault.go
@@ -3,6 +3,7 @@ package pkg
 import (
 	b64 "encoding/base64"
 	"encoding/json"
+	"errors"
 	"os"
 
 	"github.com/hashicorp/vault/api"
@@ -29,19 +30,17 @@ func connectToVault() *api.Client {
 	return client
 }
 
-func GetVarsByClusterId(clusterID string) []utils.Var {
+func GetVarsByClusterId(clusterID string) ([]utils.Var, error) {
 	client := connectToVault()
 
 	result, err := client.Logical().Read("/official-clusters-access/data/" + clusterID)
 	if err != nil {
 		log.Error(err)
-		os.Exit(1)
-		panic("unreachable") // staticcheck false positive: https://staticcheck.io/docs/checks#SA5011
+		return nil, err
 	}
 	if result == nil {
 		log.Error("Cluster information are not found")
-		os.Exit(1)
-		panic("unreachable") // staticcheck false positive: https://staticcheck.io/docs/checks#SA5011
+		return nil, errors.New("cluster information are not found")
 	}
 
 	var vaultVars []utils.Var
@@ -57,19 +56,19 @@ func GetVarsByClusterId(clusterID string) []utils.Var {
 			jsonStr, err := json.Marshal(value)
 			if err != nil {
 				log.Error("Can't convert to json GOOGLE_CREDENTIALS")
-				return []utils.Var{}
+				return []utils.Var{}, nil
 			}
 			vaultVars = append(vaultVars, utils.Var{Key: "GOOGLE_CREDENTIALS", Value: string(jsonStr)})
 		case "kubeconfig_b64", "KUBECONFIG_b64":
 			decodedValue, encErr := b64.StdEncoding.DecodeString(value.(string))
 			if encErr != nil {
 				log.Error("Can't decode KUBECONFIG")
-				return []utils.Var{}
+				return []utils.Var{}, nil
 			}
 			filePath := utils.WriteInFile(clusterID, "kubeconfig", decodedValue)
 			vaultVars = append(vaultVars, utils.Var{Key: "KUBECONFIG", Value: filePath})
 		}
 	}
 
-	return vaultVars
+	return vaultVars, nil
 }