From 81b8610dff579b1516f6fe579eff7e3811a9f332 Mon Sep 17 00:00:00 2001
From: Kristoffer Dalby <kradalby@kradalby.no>
Date: Fri, 29 Oct 2021 17:15:52 +0000
Subject: [PATCH] Add helper function to setup grpc client for cli

---
 cmd/headscale/cli/utils.go | 58 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 58 insertions(+)

diff --git a/cmd/headscale/cli/utils.go b/cmd/headscale/cli/utils.go
index f2053c93..829f4c68 100644
--- a/cmd/headscale/cli/utils.go
+++ b/cmd/headscale/cli/utils.go
@@ -281,6 +281,64 @@ func getHeadscaleApp() (*headscale.Headscale, error) {
 	return h, nil
 }
 
+func getHeadscaleGRPCClient() (apiV1.HeadscaleServiceClient, *grpc.ClientConn) {
+	grpcOptions := []grpc.DialOption{
+		// TODO(kradalby): Make configurable
+		grpc.WithTimeout(5 * time.Second),
+		grpc.WithBlock(),
+	}
+
+	address := os.Getenv("HEADSCALE_ADDRESS")
+
+	// If the address is not set, we assume that we are on the server hosting headscale.
+	if address == "" {
+		log.Debug().Msgf("HEADSCALE_ADDRESS environment is not set, connecting to localhost.")
+
+		cfg := getHeadscaleConfig()
+
+		_, port, _ := net.SplitHostPort(cfg.Addr)
+
+		address = "127.0.0.1" + ":" + port
+
+		grpcOptions = append(grpcOptions, grpc.WithInsecure())
+	}
+
+	// If we are not connecting to a local server, require an API key for authentication
+	if !headscale.IsLocalhost(address) {
+		apiKey := os.Getenv("HEADSCALE_API_KEY")
+		if apiKey == "" {
+			log.Fatal().Msgf("HEADSCALE_API_KEY environment variable needs to be set.")
+		}
+		grpcOptions = append(grpcOptions,
+			grpc.WithPerRPCCredentials(tokenAuth{
+				token: apiKey,
+			}),
+		)
+
+		insecureStr := os.Getenv("HEADSCALE_INSECURE")
+		if insecureStr != "" {
+			insecure, err := strconv.ParseBool(insecureStr)
+			if err != nil {
+				log.Fatal().Err(err).Msgf("Failed to parse HEADSCALE_INSECURE: %v", err)
+			}
+
+			if insecure {
+				grpcOptions = append(grpcOptions, grpc.WithInsecure())
+			}
+		}
+	}
+
+	log.Trace().Caller().Str("address", address).Msg("Connecting via gRPC")
+	conn, err := grpc.Dial(address, grpcOptions...)
+	if err != nil {
+		log.Fatal().Err(err).Msgf("Could not connect: %v", err)
+	}
+
+	client := apiV1.NewHeadscaleServiceClient(conn)
+
+	return client, conn
+}
+
 func JsonOutput(result interface{}, errResult error, outputFormat string) {
 	var j []byte
 	var err error