From d860270733e5441b47d773c2d111b5d9d5b52f3f Mon Sep 17 00:00:00 2001
From: Kristoffer Dalby <kradalby@kradalby.no>
Date: Mon, 30 May 2022 16:10:39 +0200
Subject: [PATCH] Use Prometheus duration parser (support days and weeks)

---
 cmd/headscale/cli/api_key.go     | 23 ++++++++++++++++++-----
 cmd/headscale/cli/preauthkeys.go | 23 ++++++++++++++++++-----
 2 files changed, 36 insertions(+), 10 deletions(-)

diff --git a/cmd/headscale/cli/api_key.go b/cmd/headscale/cli/api_key.go
index aa056c5..d97cefa 100644
--- a/cmd/headscale/cli/api_key.go
+++ b/cmd/headscale/cli/api_key.go
@@ -7,6 +7,7 @@ import (
 
 	"github.com/juanfont/headscale"
 	v1 "github.com/juanfont/headscale/gen/go/headscale/v1"
+	"github.com/prometheus/common/model"
 	"github.com/pterm/pterm"
 	"github.com/rs/zerolog/log"
 	"github.com/spf13/cobra"
@@ -15,7 +16,7 @@ import (
 
 const (
 	// 90 days.
-	DefaultAPIKeyExpiry = 90 * 24 * time.Hour
+	DefaultAPIKeyExpiry = "90d"
 )
 
 func init() {
@@ -23,7 +24,7 @@ func init() {
 	apiKeysCmd.AddCommand(listAPIKeys)
 
 	createAPIKeyCmd.Flags().
-		DurationP("expiration", "e", DefaultAPIKeyExpiry, "Human-readable expiration of the key (e.g. 30m, 24h)")
+		StringP("expiration", "e", DefaultAPIKeyExpiry, "Human-readable expiration of the key (e.g. 30m, 24h)")
 
 	apiKeysCmd.AddCommand(createAPIKeyCmd)
 
@@ -118,10 +119,22 @@ If you loose a key, create a new one and revoke (expire) the old one.`,
 
 		request := &v1.CreateApiKeyRequest{}
 
-		duration, _ := cmd.Flags().GetDuration("expiration")
-		expiration := time.Now().UTC().Add(duration)
+		durationStr, _ := cmd.Flags().GetString("expiration")
 
-		log.Trace().Dur("expiration", duration).Msg("expiration has been set")
+		duration, err := model.ParseDuration(durationStr)
+		if err != nil {
+			ErrorOutput(
+				err,
+				fmt.Sprintf("Could not parse duration: %s\n", err),
+				output,
+			)
+
+			return
+		}
+
+		expiration := time.Now().UTC().Add(time.Duration(duration))
+
+		log.Trace().Dur("expiration", time.Duration(duration)).Msg("expiration has been set")
 
 		request.Expiration = timestamppb.New(expiration)
 
diff --git a/cmd/headscale/cli/preauthkeys.go b/cmd/headscale/cli/preauthkeys.go
index 7efb72f..ffa1a81 100644
--- a/cmd/headscale/cli/preauthkeys.go
+++ b/cmd/headscale/cli/preauthkeys.go
@@ -6,6 +6,7 @@ import (
 	"time"
 
 	v1 "github.com/juanfont/headscale/gen/go/headscale/v1"
+	"github.com/prometheus/common/model"
 	"github.com/pterm/pterm"
 	"github.com/rs/zerolog/log"
 	"github.com/spf13/cobra"
@@ -13,7 +14,7 @@ import (
 )
 
 const (
-	DefaultPreAuthKeyExpiry = 1 * time.Hour
+	DefaultPreAuthKeyExpiry = "1h"
 )
 
 func init() {
@@ -31,7 +32,7 @@ func init() {
 	createPreAuthKeyCmd.PersistentFlags().
 		Bool("ephemeral", false, "Preauthkey for ephemeral nodes")
 	createPreAuthKeyCmd.Flags().
-		DurationP("expiration", "e", DefaultPreAuthKeyExpiry, "Human-readable expiration of the key (e.g. 30m, 24h)")
+		StringP("expiration", "e", DefaultPreAuthKeyExpiry, "Human-readable expiration of the key (e.g. 30m, 24h)")
 }
 
 var preauthkeysCmd = &cobra.Command{
@@ -148,10 +149,22 @@ var createPreAuthKeyCmd = &cobra.Command{
 			Ephemeral: ephemeral,
 		}
 
-		duration, _ := cmd.Flags().GetDuration("expiration")
-		expiration := time.Now().UTC().Add(duration)
+		durationStr, _ := cmd.Flags().GetString("expiration")
 
-		log.Trace().Dur("expiration", duration).Msg("expiration has been set")
+		duration, err := model.ParseDuration(durationStr)
+		if err != nil {
+			ErrorOutput(
+				err,
+				fmt.Sprintf("Could not parse duration: %s\n", err),
+				output,
+			)
+
+			return
+		}
+
+		expiration := time.Now().UTC().Add(time.Duration(duration))
+
+		log.Trace().Dur("expiration", time.Duration(duration)).Msg("expiration has been set")
 
 		request.Expiration = timestamppb.New(expiration)