From 80f3ee235643be7c6ae910245274d8bff96bc407 Mon Sep 17 00:00:00 2001
From: jlandowner <jlandowner8@gmail.com>
Date: Tue, 27 Apr 2021 11:55:33 +0000
Subject: [PATCH] Support kubecontext

---
 cmd/attach.go        | 2 +-
 cmd/clean.go         | 2 +-
 cmd/detach.go        | 2 +-
 cmd/list.go          | 2 +-
 cmd/root.go          | 4 +++-
 cmd/tree.go          | 2 +-
 pkg/client/client.go | 8 ++++++--
 7 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/cmd/attach.go b/cmd/attach.go
index b913064..b997c7f 100644
--- a/cmd/attach.go
+++ b/cmd/attach.go
@@ -51,7 +51,7 @@ var (
 
 		RunE: func(cmd *cobra.Command, args []string) error {
 			ctx := context.Background()
-			k8sclient, err := client.NewClient(&kubeconfigPath)
+			k8sclient, err := client.NewClient(&kubeconfigPath, &kubecontext)
 			if err != nil {
 				return fmt.Errorf("Failed to load kubeconfig %v: %v", kubeconfigPath, err.Error())
 			}
diff --git a/cmd/clean.go b/cmd/clean.go
index 01454c9..5227cf4 100644
--- a/cmd/clean.go
+++ b/cmd/clean.go
@@ -39,7 +39,7 @@ var (
 		PersistentPreRunE: c.PreRunE,
 		RunE: func(cmd *cobra.Command, args []string) error {
 			ctx := context.Background()
-			k8sclient, err := client.NewClient(&kubeconfigPath)
+			k8sclient, err := client.NewClient(&kubeconfigPath, &kubecontext)
 			if err != nil {
 				return fmt.Errorf("Failed to load kubeconfig %v: %v", kubeconfigPath, err.Error())
 			}
diff --git a/cmd/detach.go b/cmd/detach.go
index c18d25d..feb4930 100644
--- a/cmd/detach.go
+++ b/cmd/detach.go
@@ -51,7 +51,7 @@ var (
 
 		RunE: func(cmd *cobra.Command, args []string) error {
 			ctx := context.Background()
-			k8sclient, err := client.NewClient(&kubeconfigPath)
+			k8sclient, err := client.NewClient(&kubeconfigPath, &kubecontext)
 			if err != nil {
 				return fmt.Errorf("Failed to load kubeconfig %v: %v", kubeconfigPath, err.Error())
 			}
diff --git a/cmd/list.go b/cmd/list.go
index 30a55ad..b19ca00 100644
--- a/cmd/list.go
+++ b/cmd/list.go
@@ -44,7 +44,7 @@ var (
 		PersistentPreRunE: l.PreRunE,
 		RunE: func(cmd *cobra.Command, args []string) error {
 			ctx := context.Background()
-			k8sclient, err := client.NewClient(&kubeconfigPath)
+			k8sclient, err := client.NewClient(&kubeconfigPath, &kubecontext)
 			if err != nil {
 				return fmt.Errorf("Failed to load kubeconfig %v: %v", kubeconfigPath, err.Error())
 			}
diff --git a/cmd/root.go b/cmd/root.go
index 892d31b..3c567a4 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -24,6 +24,7 @@ import (
 
 var (
 	kubeconfigPath string
+	kubecontext    string
 
 	rootCmd = &cobra.Command{
 		Use:   "psp-util",
@@ -41,7 +42,8 @@ Apache License Version 2.0 Copyright 2020 jlandowner
 )
 
 func init() {
-	rootCmd.PersistentFlags().StringVar(&kubeconfigPath, "kubeconfig", "", "kube config file (default is $HOME/.kube/config)")
+	rootCmd.PersistentFlags().StringVar(&kubeconfigPath, "kubeconfig", "", "kubeconfig file path (default: $HOME/.kube/config)")
+	rootCmd.PersistentFlags().StringVar(&kubecontext, "context", "", "kube-context (default: current context)")
 }
 
 func Execute() {
diff --git a/cmd/tree.go b/cmd/tree.go
index 61cb246..144e475 100644
--- a/cmd/tree.go
+++ b/cmd/tree.go
@@ -36,7 +36,7 @@ var treeCmd = &cobra.Command{
 	Short: "View a relational tree between PSP and Subjects in cluster",
 	RunE: func(cmd *cobra.Command, args []string) error {
 		ctx := context.Background()
-		k8sclient, err := client.NewClient(&kubeconfigPath)
+		k8sclient, err := client.NewClient(&kubeconfigPath, &kubecontext)
 		if err != nil {
 			return fmt.Errorf("Failed to load kubeconfig %v: %v", kubeconfigPath, err.Error())
 		}
diff --git a/pkg/client/client.go b/pkg/client/client.go
index 5af1088..4334e59 100644
--- a/pkg/client/client.go
+++ b/pkg/client/client.go
@@ -26,15 +26,19 @@ import (
 )
 
 // NewClient returns kubernetes Clientset
-func NewClient(kubeconfigPath *string) (*kubernetes.Clientset, error) {
+func NewClient(kubeconfigPath *string, kubecontext *string) (*kubernetes.Clientset, error) {
 	if *kubeconfigPath == "" {
 		*kubeconfigPath = filepath.Join(homeDir(), ".kube", "config")
 	}
 
-	config, err := clientcmd.BuildConfigFromFlags("", *kubeconfigPath)
+	config, err := clientcmd.NewNonInteractiveDeferredLoadingClientConfig(
+		&clientcmd.ClientConfigLoadingRules{ExplicitPath: *kubeconfigPath},
+		&clientcmd.ConfigOverrides{CurrentContext: *kubecontext}).ClientConfig()
+
 	if err != nil {
 		return nil, err
 	}
+
 	return kubernetes.NewForConfig(config)
 }