mirror of
https://github.com/juanfont/headscale.git
synced 2025-03-01 00:21:18 +01:00
2dfd42f80c
This commit is trying to DRY up the initiation of the gRPC client in each command: It renames the function to CLI instead of GRPC as it actually set up a CLI client, not a generic grpc client It also moves the configuration of address, timeout (which is now consistent) and api to use Viper, allowing users to set it via env vars and configuration file
362 lines
9.3 KiB
Go
362 lines
9.3 KiB
Go
package cli
|
|
|
|
import (
|
|
"fmt"
|
|
"log"
|
|
"strconv"
|
|
"time"
|
|
|
|
survey "github.com/AlecAivazis/survey/v2"
|
|
v1 "github.com/juanfont/headscale/gen/go/headscale/v1"
|
|
"github.com/pterm/pterm"
|
|
"github.com/spf13/cobra"
|
|
"tailscale.com/tailcfg"
|
|
"tailscale.com/types/wgkey"
|
|
)
|
|
|
|
func init() {
|
|
rootCmd.AddCommand(nodeCmd)
|
|
listNodesCmd.Flags().StringP("namespace", "n", "", "Filter by namespace")
|
|
nodeCmd.AddCommand(listNodesCmd)
|
|
|
|
registerNodeCmd.Flags().StringP("namespace", "n", "", "Namespace")
|
|
err := registerNodeCmd.MarkFlagRequired("namespace")
|
|
if err != nil {
|
|
log.Fatalf(err.Error())
|
|
}
|
|
registerNodeCmd.Flags().StringP("key", "k", "", "Key")
|
|
err = registerNodeCmd.MarkFlagRequired("key")
|
|
if err != nil {
|
|
log.Fatalf(err.Error())
|
|
}
|
|
nodeCmd.AddCommand(registerNodeCmd)
|
|
|
|
deleteNodeCmd.Flags().IntP("identifier", "i", 0, "Node identifier (ID)")
|
|
err = deleteNodeCmd.MarkFlagRequired("identifier")
|
|
if err != nil {
|
|
log.Fatalf(err.Error())
|
|
}
|
|
nodeCmd.AddCommand(deleteNodeCmd)
|
|
|
|
shareMachineCmd.Flags().StringP("namespace", "n", "", "Namespace")
|
|
err = shareMachineCmd.MarkFlagRequired("namespace")
|
|
if err != nil {
|
|
log.Fatalf(err.Error())
|
|
}
|
|
shareMachineCmd.Flags().IntP("identifier", "i", 0, "Node identifier (ID)")
|
|
err = shareMachineCmd.MarkFlagRequired("identifier")
|
|
if err != nil {
|
|
log.Fatalf(err.Error())
|
|
}
|
|
nodeCmd.AddCommand(shareMachineCmd)
|
|
|
|
unshareMachineCmd.Flags().StringP("namespace", "n", "", "Namespace")
|
|
err = unshareMachineCmd.MarkFlagRequired("namespace")
|
|
if err != nil {
|
|
log.Fatalf(err.Error())
|
|
}
|
|
unshareMachineCmd.Flags().IntP("identifier", "i", 0, "Node identifier (ID)")
|
|
err = unshareMachineCmd.MarkFlagRequired("identifier")
|
|
if err != nil {
|
|
log.Fatalf(err.Error())
|
|
}
|
|
nodeCmd.AddCommand(unshareMachineCmd)
|
|
}
|
|
|
|
var nodeCmd = &cobra.Command{
|
|
Use: "nodes",
|
|
Short: "Manage the nodes of Headscale",
|
|
}
|
|
|
|
var registerNodeCmd = &cobra.Command{
|
|
Use: "register",
|
|
Short: "Registers a machine to your network",
|
|
Run: func(cmd *cobra.Command, args []string) {
|
|
output, _ := cmd.Flags().GetString("output")
|
|
namespace, err := cmd.Flags().GetString("namespace")
|
|
if err != nil {
|
|
ErrorOutput(err, fmt.Sprintf("Error getting namespace: %s", err), output)
|
|
return
|
|
}
|
|
|
|
ctx, client, conn, cancel := getHeadscaleCLIClient()
|
|
defer cancel()
|
|
defer conn.Close()
|
|
|
|
machineKey, err := cmd.Flags().GetString("key")
|
|
if err != nil {
|
|
ErrorOutput(err, fmt.Sprintf("Error getting machine key from flag: %s", err), output)
|
|
return
|
|
}
|
|
|
|
request := &v1.RegisterMachineRequest{
|
|
Key: machineKey,
|
|
Namespace: namespace,
|
|
}
|
|
|
|
response, err := client.RegisterMachine(ctx, request)
|
|
if err != nil {
|
|
ErrorOutput(err, fmt.Sprintf("Cannot register machine: %s\n", err), output)
|
|
return
|
|
}
|
|
|
|
SuccessOutput(response.Machine, "Machine register", output)
|
|
},
|
|
}
|
|
|
|
var listNodesCmd = &cobra.Command{
|
|
Use: "list",
|
|
Short: "List nodes",
|
|
Run: func(cmd *cobra.Command, args []string) {
|
|
output, _ := cmd.Flags().GetString("output")
|
|
namespace, err := cmd.Flags().GetString("namespace")
|
|
if err != nil {
|
|
ErrorOutput(err, fmt.Sprintf("Error getting namespace: %s", err), output)
|
|
return
|
|
}
|
|
|
|
ctx, client, conn, cancel := getHeadscaleCLIClient()
|
|
defer cancel()
|
|
defer conn.Close()
|
|
|
|
request := &v1.ListMachinesRequest{
|
|
Namespace: namespace,
|
|
}
|
|
|
|
response, err := client.ListMachines(ctx, request)
|
|
if err != nil {
|
|
ErrorOutput(err, fmt.Sprintf("Cannot get nodes: %s", err), output)
|
|
return
|
|
}
|
|
|
|
if output != "" {
|
|
SuccessOutput(response.Machines, "", output)
|
|
return
|
|
}
|
|
|
|
d, err := nodesToPtables(namespace, response.Machines)
|
|
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 {
|
|
ErrorOutput(err, fmt.Sprintf("Failed to render pterm table: %s", err), output)
|
|
return
|
|
}
|
|
},
|
|
}
|
|
|
|
var deleteNodeCmd = &cobra.Command{
|
|
Use: "delete",
|
|
Short: "Delete a node",
|
|
Run: func(cmd *cobra.Command, args []string) {
|
|
output, _ := cmd.Flags().GetString("output")
|
|
|
|
id, err := cmd.Flags().GetInt("identifier")
|
|
if err != nil {
|
|
ErrorOutput(err, fmt.Sprintf("Error converting ID to integer: %s", err), output)
|
|
return
|
|
}
|
|
|
|
ctx, client, conn, cancel := getHeadscaleCLIClient()
|
|
defer cancel()
|
|
defer conn.Close()
|
|
|
|
getRequest := &v1.GetMachineRequest{
|
|
MachineId: uint64(id),
|
|
}
|
|
|
|
getResponse, err := client.GetMachine(ctx, getRequest)
|
|
if err != nil {
|
|
ErrorOutput(err, fmt.Sprintf("Error getting node node: %s", err), output)
|
|
return
|
|
}
|
|
|
|
deleteRequest := &v1.DeleteMachineRequest{
|
|
MachineId: uint64(id),
|
|
}
|
|
|
|
confirm := false
|
|
force, _ := cmd.Flags().GetBool("force")
|
|
if !force {
|
|
prompt := &survey.Confirm{
|
|
Message: fmt.Sprintf("Do you want to remove the node %s?", getResponse.GetMachine().Name),
|
|
}
|
|
err = survey.AskOne(prompt, &confirm)
|
|
if err != nil {
|
|
return
|
|
}
|
|
}
|
|
|
|
if confirm || force {
|
|
response, err := client.DeleteMachine(ctx, deleteRequest)
|
|
if output != "" {
|
|
SuccessOutput(response, "", output)
|
|
return
|
|
}
|
|
if err != nil {
|
|
ErrorOutput(err, fmt.Sprintf("Error deleting node: %s", err), output)
|
|
return
|
|
}
|
|
SuccessOutput(map[string]string{"Result": "Node deleted"}, "Node deleted", output)
|
|
} else {
|
|
SuccessOutput(map[string]string{"Result": "Node not deleted"}, "Node not deleted", output)
|
|
}
|
|
},
|
|
}
|
|
|
|
func sharingWorker(
|
|
cmd *cobra.Command,
|
|
args []string,
|
|
) (string, *v1.Machine, *v1.Namespace, error) {
|
|
output, _ := cmd.Flags().GetString("output")
|
|
namespaceStr, err := cmd.Flags().GetString("namespace")
|
|
if err != nil {
|
|
ErrorOutput(err, fmt.Sprintf("Error getting namespace: %s", err), output)
|
|
return "", nil, nil, err
|
|
}
|
|
|
|
ctx, client, conn, cancel := getHeadscaleCLIClient()
|
|
defer cancel()
|
|
defer conn.Close()
|
|
|
|
id, err := cmd.Flags().GetInt("identifier")
|
|
if err != nil {
|
|
ErrorOutput(err, fmt.Sprintf("Error converting ID to integer: %s", err), output)
|
|
return "", nil, nil, err
|
|
}
|
|
|
|
machineRequest := &v1.GetMachineRequest{
|
|
MachineId: uint64(id),
|
|
}
|
|
|
|
machineResponse, err := client.GetMachine(ctx, machineRequest)
|
|
if err != nil {
|
|
ErrorOutput(err, fmt.Sprintf("Error getting node node: %s", err), output)
|
|
return "", nil, nil, err
|
|
}
|
|
|
|
namespaceRequest := &v1.GetNamespaceRequest{
|
|
Name: namespaceStr,
|
|
}
|
|
|
|
namespaceResponse, err := client.GetNamespace(ctx, namespaceRequest)
|
|
if err != nil {
|
|
ErrorOutput(err, fmt.Sprintf("Error getting node node: %s", err), output)
|
|
return "", nil, nil, err
|
|
}
|
|
|
|
return output, machineResponse.GetMachine(), namespaceResponse.GetNamespace(), nil
|
|
}
|
|
|
|
var shareMachineCmd = &cobra.Command{
|
|
Use: "share",
|
|
Short: "Shares a node from the current namespace to the specified one",
|
|
Run: func(cmd *cobra.Command, args []string) {
|
|
output, machine, namespace, err := sharingWorker(cmd, args)
|
|
if err != nil {
|
|
ErrorOutput(err, fmt.Sprintf("Failed to fetch namespace or machine: %s", err), output)
|
|
return
|
|
}
|
|
|
|
ctx, client, conn, cancel := getHeadscaleCLIClient()
|
|
defer cancel()
|
|
defer conn.Close()
|
|
|
|
request := &v1.ShareMachineRequest{
|
|
MachineId: machine.Id,
|
|
Namespace: namespace.Name,
|
|
}
|
|
|
|
response, err := client.ShareMachine(ctx, request)
|
|
if err != nil {
|
|
ErrorOutput(err, fmt.Sprintf("Error sharing node: %s", err), output)
|
|
return
|
|
}
|
|
|
|
SuccessOutput(response.Machine, "Node shared", output)
|
|
},
|
|
}
|
|
|
|
var unshareMachineCmd = &cobra.Command{
|
|
Use: "unshare",
|
|
Short: "Unshares a node from the specified namespace",
|
|
Run: func(cmd *cobra.Command, args []string) {
|
|
output, machine, namespace, err := sharingWorker(cmd, args)
|
|
if err != nil {
|
|
ErrorOutput(err, fmt.Sprintf("Failed to fetch namespace or machine: %s", err), output)
|
|
return
|
|
}
|
|
|
|
ctx, client, conn, cancel := getHeadscaleCLIClient()
|
|
defer cancel()
|
|
defer conn.Close()
|
|
|
|
request := &v1.UnshareMachineRequest{
|
|
MachineId: machine.Id,
|
|
Namespace: namespace.Name,
|
|
}
|
|
|
|
response, err := client.UnshareMachine(ctx, request)
|
|
if err != nil {
|
|
ErrorOutput(err, fmt.Sprintf("Error unsharing node: %s", err), output)
|
|
return
|
|
}
|
|
|
|
SuccessOutput(response.Machine, "Node shared", output)
|
|
},
|
|
}
|
|
|
|
func nodesToPtables(currentNamespace string, machines []*v1.Machine) (pterm.TableData, error) {
|
|
d := pterm.TableData{{"ID", "Name", "NodeKey", "Namespace", "IP address", "Ephemeral", "Last seen", "Online"}}
|
|
|
|
for _, machine := range machines {
|
|
var ephemeral bool
|
|
if machine.PreAuthKey != nil && machine.PreAuthKey.Ephemeral {
|
|
ephemeral = true
|
|
}
|
|
var lastSeen time.Time
|
|
var lastSeenTime string
|
|
if machine.LastSeen != nil {
|
|
lastSeen = machine.LastSeen.AsTime()
|
|
lastSeenTime = lastSeen.Format("2006-01-02 15:04:05")
|
|
}
|
|
nKey, err := wgkey.ParseHex(machine.NodeKey)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
nodeKey := tailcfg.NodeKey(nKey)
|
|
|
|
var online string
|
|
if lastSeen.After(time.Now().Add(-5 * time.Minute)) { // TODO: Find a better way to reliably show if online
|
|
online = pterm.LightGreen("true")
|
|
} else {
|
|
online = pterm.LightRed("false")
|
|
}
|
|
|
|
var namespace string
|
|
if currentNamespace == "" || (currentNamespace == machine.Namespace.Name) {
|
|
namespace = pterm.LightMagenta(machine.Namespace.Name)
|
|
} else {
|
|
// Shared into this namespace
|
|
namespace = pterm.LightYellow(machine.Namespace.Name)
|
|
}
|
|
d = append(
|
|
d,
|
|
[]string{
|
|
strconv.FormatUint(machine.Id, 10),
|
|
machine.Name,
|
|
nodeKey.ShortString(),
|
|
namespace,
|
|
machine.IpAddress,
|
|
strconv.FormatBool(ephemeral),
|
|
lastSeenTime,
|
|
online,
|
|
},
|
|
)
|
|
}
|
|
return d, nil
|
|
}
|