...

Source file src/edge-infra.dev/third_party/gke-iap/gcloud.go

Documentation: edge-infra.dev/third_party/gke-iap

     1  package main
     2  
     3  import (
     4  	"fmt"
     5  	"math/rand"
     6  	"strconv"
     7  	"strings"
     8  )
     9  
    10  type gcloud struct {
    11  	opts *Options
    12  }
    13  
    14  type gkeInstance struct {
    15  	Name string
    16  	Zone string
    17  }
    18  
    19  func (g *gcloud) SetupKubectl() error {
    20  	args := []string{
    21  		"container",
    22  		"clusters",
    23  		"get-credentials",
    24  		g.opts.ClusterName,
    25  		"--project", g.opts.ProjectID,
    26  	}
    27  
    28  	location, err := g.getGKELocation()
    29  	if err != nil {
    30  		return err
    31  	}
    32  
    33  	if len(strings.Split(location, "-")) == 3 {
    34  		args = append(args, "--zone", location)
    35  	} else {
    36  		args = append(args, "--region", location)
    37  	}
    38  
    39  	kubectlContext := fmt.Sprintf("gke_%s_%s_%s", g.opts.ProjectID, location, g.opts.ClusterName)
    40  
    41  	msg := fmt.Sprintf("Setting %s context in kubeconfig\n", kubectlContext)
    42  	_, err = Exec("gcloud", args, msg)
    43  	if err != nil {
    44  		return err
    45  	}
    46  
    47  	err = g.setupKubectlProxy(kubectlContext, location)
    48  	if err != nil {
    49  		return err
    50  	}
    51  
    52  	return nil
    53  }
    54  
    55  func (g *gcloud) setupKubectlProxy(kubectlContext, location string) error {
    56  	args := []string{
    57  		"config",
    58  		"set-cluster",
    59  		"--proxy-url", fmt.Sprintf("http://127.0.0.1:%d", g.opts.LocalPort),
    60  		fmt.Sprintf("gke_%s_%s_%s", g.opts.ProjectID, location, g.opts.ClusterName),
    61  	}
    62  
    63  	msg := fmt.Sprintf("Setting http://127.0.0.1:%d HTTP proxy for %s context\n", g.opts.LocalPort, kubectlContext)
    64  	_, err := Exec("kubectl", args, msg)
    65  	if err != nil {
    66  		return err
    67  	}
    68  
    69  	return nil
    70  }
    71  
    72  func (g *gcloud) StartTunnel() error {
    73  	instances, err := g.getGKEInstances()
    74  	if err != nil {
    75  		return err
    76  	}
    77  
    78  	instance := GetRandomGKEInstance(instances)
    79  
    80  	args := []string{
    81  		"compute",
    82  		"start-iap-tunnel",
    83  		"--project", g.opts.ProjectID,
    84  		"--zone", instance.Zone,
    85  		"--local-host-port", fmt.Sprintf("127.0.0.1:%d", g.opts.LocalPort),
    86  		instance.Name, strconv.Itoa(g.opts.InstancePort),
    87  	}
    88  
    89  	msg := fmt.Sprintf("Listening on port [%d]\n", g.opts.LocalPort)
    90  	_, err = Exec("gcloud", args, msg)
    91  	if err != nil {
    92  		return err
    93  	}
    94  
    95  	return nil
    96  }
    97  
    98  func (g *gcloud) getGKELocation() (string, error) {
    99  	args := []string{
   100  		"container",
   101  		"clusters",
   102  		"list",
   103  		"--project", g.opts.ProjectID,
   104  		"--filter", fmt.Sprintf("name=%s", g.opts.ClusterName),
   105  		"--format", "value(location)",
   106  	}
   107  
   108  	msg := fmt.Sprintf("Fetching '%s' GKE cluster location in '%s'\n", g.opts.ClusterName, g.opts.ProjectID)
   109  	stdout, err := Exec("gcloud", args, msg)
   110  	if err != nil {
   111  		return "", err
   112  	}
   113  
   114  	return strings.TrimSpace(stdout), nil
   115  }
   116  
   117  func (g *gcloud) getGKEInstances() ([]gkeInstance, error) {
   118  	args := []string{
   119  		"compute",
   120  		"instances",
   121  		"list",
   122  		"--project", g.opts.ProjectID,
   123  		"--filter", fmt.Sprintf("name~^gke-%s-", g.opts.ClusterName),
   124  		"--format", "csv[no-heading](name,zone)",
   125  	}
   126  
   127  	msg := fmt.Sprintf("Fetching instances of '%s' GKE cluster in '%s'\n", g.opts.ClusterName, g.opts.ProjectID)
   128  	stdout, err := Exec("gcloud", args, msg)
   129  	if err != nil {
   130  		return []gkeInstance{}, err
   131  	}
   132  
   133  	var instances []gkeInstance
   134  
   135  	rows := strings.TrimSpace(stdout)
   136  
   137  	for _, row := range strings.Split(rows, "\n") {
   138  		data := strings.Split(row, ",")
   139  		instances = append(instances, gkeInstance{
   140  			Name: data[0],
   141  			Zone: data[1],
   142  		})
   143  	}
   144  
   145  	return instances, nil
   146  }
   147  
   148  func GetActiveProject() string {
   149  	args := []string{
   150  		"config",
   151  		"get",
   152  		"project",
   153  	}
   154  	stdout, _ := Exec("gcloud", args, "")
   155  	return strings.TrimSpace(stdout)
   156  }
   157  
   158  func GetRandomGKEInstance(instances []gkeInstance) gkeInstance {
   159  	return instances[rand.Intn(len(instances))]
   160  }
   161  

View as plain text