Convert routes command to use gRPC

cmd/headscale/cli/routes.go

@@ -1,26 +1,37 @@
 package cli
 
 import (
+	"context"
 	"fmt"
 	"log"
-	"strings"
+	"strconv"
+	"time"
 
+	v1 "github.com/juanfont/headscale/gen/go/headscale/v1"
 	"github.com/pterm/pterm"
 	"github.com/spf13/cobra"
 )
 
 func init() {
 	rootCmd.AddCommand(routesCmd)
-	routesCmd.PersistentFlags().StringP("namespace", "n", "", "Namespace")
+
-	err := routesCmd.MarkPersistentFlagRequired("namespace")
+	listRoutesCmd.Flags().Uint64P("identifier", "i", 0, "Node identifier (ID)")
+	err := listRoutesCmd.MarkFlagRequired("identifier")
+	if err != nil {
+		log.Fatalf(err.Error())
+	}
+	routesCmd.AddCommand(listRoutesCmd)
+
+	enableRouteCmd.Flags().StringSliceP("route", "r", []string{}, "List (or repeated flags) of routes to enable")
+	enableRouteCmd.Flags().Uint64P("identifier", "i", 0, "Node identifier (ID)")
+	err = enableRouteCmd.MarkFlagRequired("identifier")
 	if err != nil {
 		log.Fatalf(err.Error())
 	}
 
-	enableRouteCmd.Flags().BoolP("all", "a", false, "Enable all routes advertised by the node")
-
-	routesCmd.AddCommand(listRoutesCmd)
 	routesCmd.AddCommand(enableRouteCmd)
+
+	nodeCmd.AddCommand(routesCmd)
 }
 
 var routesCmd = &cobra.Command{
@@ -29,119 +40,128 @@ var routesCmd = &cobra.Command{
 }
 
 var listRoutesCmd = &cobra.Command{
-	Use:   "list NODE",
+	Use:   "list",
-	Short: "List the routes exposed by this node",
+	Short: "List routes advertised and enabled by a given node",
-	Args: func(cmd *cobra.Command, args []string) error {
-		if len(args) < 1 {
-			return fmt.Errorf("Missing parameters")
-		}
-		return nil
-	},
 	Run: func(cmd *cobra.Command, args []string) {
-		n, err := cmd.Flags().GetString("namespace")
+		output, _ := cmd.Flags().GetString("output")
-		if err != nil {
-			log.Fatalf("Error getting namespace: %s", err)
-		}
-		o, _ := cmd.Flags().GetString("output")
 
-		h, err := getHeadscaleApp()
+		machineId, err := cmd.Flags().GetUint64("identifier")
 		if err != nil {
-			log.Fatalf("Error initializing: %s", err)
+			ErrorOutput(err, fmt.Sprintf("Error getting machine id from flag: %s", err), output)
-		}
-
-		availableRoutes, err := h.GetAdvertisedNodeRoutes(n, args[0])
-		if err != nil {
-			fmt.Println(err)
 			return
 		}
 
-		if strings.HasPrefix(o, "json") {
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
-			// TODO: Add enable/disabled information to this interface
+		defer cancel()
-			JsonOutput(availableRoutes, err, o)
+
+		client, conn := getHeadscaleGRPCClient(ctx)
+		defer conn.Close()
+
+		request := &v1.GetMachineRouteRequest{
+			MachineId: machineId,
+		}
+
+		response, err := client.GetMachineRoute(ctx, request)
+		if err != nil {
+			ErrorOutput(err, fmt.Sprintf("Cannot get nodes: %s", err), output)
 			return
 		}
 
-		d := h.RoutesToPtables(n, args[0], *availableRoutes)
+		if output != "" {
+			SuccessOutput(response.Routes, "", output)
+			return
+		}
+
+		d := routesToPtables(response.Routes)
+		if err != nil {
+			ErrorOutput(err, fmt.Sprintf("Error converting to table: %s", err), output)
+			return
+		}
 
 		err = pterm.DefaultTable.WithHasHeader().WithData(d).Render()
 		if err != nil {
-			log.Fatal(err)
+			ErrorOutput(err, fmt.Sprintf("Failed to render pterm table: %s", err), output)
+			return
 		}
 	},
 }
 
 var enableRouteCmd = &cobra.Command{
-	Use:   "enable node-name route",
+	Use:   "enable",
-	Short: "Allows exposing a route declared by this node to the rest of the nodes",
+	Short: "Set the enabled routes for a given node",
-	Args: func(cmd *cobra.Command, args []string) error {
+	Long: `This command will take a list of routes that will _replace_ 
-		all, err := cmd.Flags().GetBool("all")
+the current set of routes on a given node.
-		if err != nil {
+If you would like to disable a route, simply run the command again, but 
-			log.Fatalf("Error getting namespace: %s", err)
+omit the route you do not want to enable.
-		}
+	`,
-
-		if all {
-			if len(args) < 1 {
-				return fmt.Errorf("Missing parameters")
-			}
-			return nil
-		} else {
-			if len(args) < 2 {
-				return fmt.Errorf("Missing parameters")
-			}
-			return nil
-		}
-	},
 	Run: func(cmd *cobra.Command, args []string) {
-		n, err := cmd.Flags().GetString("namespace")
+		output, _ := cmd.Flags().GetString("output")
+		machineId, err := cmd.Flags().GetUint64("identifier")
 		if err != nil {
-			log.Fatalf("Error getting namespace: %s", err)
+			ErrorOutput(err, fmt.Sprintf("Error getting machine id from flag: %s", err), output)
+			return
 		}
 
-		o, _ := cmd.Flags().GetString("output")
+		routes, err := cmd.Flags().GetStringSlice("route")
-
-		all, err := cmd.Flags().GetBool("all")
 		if err != nil {
-			log.Fatalf("Error getting namespace: %s", err)
+			ErrorOutput(err, fmt.Sprintf("Error getting routes from flag: %s", err), output)
+			return
 		}
 
-		h, err := getHeadscaleApp()
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
-		if err != nil {
+		defer cancel()
-			log.Fatalf("Error initializing: %s", err)
+
+		client, conn := getHeadscaleGRPCClient(ctx)
+		defer conn.Close()
+
+		request := &v1.EnableMachineRoutesRequest{
+			MachineId: machineId,
+			Routes:    routes,
 		}
 
-		if all {
+		response, err := client.EnableMachineRoutes(ctx, request)
-			availableRoutes, err := h.GetAdvertisedNodeRoutes(n, args[0])
+		if err != nil {
-			if err != nil {
+			ErrorOutput(err, fmt.Sprintf("Cannot register machine: %s\n", err), output)
-				fmt.Println(err)
+			return
-				return
+		}
-			}
 
-			for _, availableRoute := range *availableRoutes {
+		if output != "" {
-				err = h.EnableNodeRoute(n, args[0], availableRoute.String())
+			SuccessOutput(response.Routes, "", output)
-				if err != nil {
+			return
-					fmt.Println(err)
+		}
-					return
-				}
 
-				if strings.HasPrefix(o, "json") {
+		d := routesToPtables(response.Routes)
-					JsonOutput(availableRoute, err, o)
+		if err != nil {
-				} else {
+			ErrorOutput(err, fmt.Sprintf("Error converting to table: %s", err), output)
-					fmt.Printf("Enabled route %s\n", availableRoute)
+			return
-				}
+		}
-			}
-		} else {
-			err = h.EnableNodeRoute(n, args[0], args[1])
 
-			if strings.HasPrefix(o, "json") {
+		err = pterm.DefaultTable.WithHasHeader().WithData(d).Render()
-				JsonOutput(args[1], err, o)
+		if err != nil {
-				return
+			ErrorOutput(err, fmt.Sprintf("Failed to render pterm table: %s", err), output)
-			}
+			return
-
-			if err != nil {
-				fmt.Println(err)
-				return
-			}
-			fmt.Printf("Enabled route %s\n", args[1])
 		}
 	},
 }
+
+// routesToPtables converts the list of routes to a nice table
+func routesToPtables(routes *v1.Routes) pterm.TableData {
+	d := pterm.TableData{{"Route", "Enabled"}}
+
+	for _, route := range routes.GetAdvertisedRoutes() {
+		enabled := isStringInSlice(routes.EnabledRoutes, route)
+
+		d = append(d, []string{route, strconv.FormatBool(enabled)})
+	}
+	return d
+}
+
+func isStringInSlice(strs []string, s string) bool {
+	for _, s2 := range strs {
+		if s == s2 {
+			return true
+		}
+	}
+
+	return false
+}