From cc09fd0613970beccde0528c01e0a261543ebbb3 Mon Sep 17 00:00:00 2001
From: Dimitrios Karagiannis <dhkarag@gmail.com>
Date: Wed, 22 Jan 2020 12:24:57 +0000
Subject: [PATCH] Rework the megaport_token helper tool

It is no longer interactive and gets its inputs from environment
variables. This should allow us to employ it for automated testing.

Signed-off-by: Dimitrios Karagiannis <dhkarag@gmail.com>
---
 Makefile                    |  3 ++
 README.md                   | 10 +++---
 util/megaport_token/main.go | 72 ++++++++++++++++++++++++-------------
 3 files changed, 56 insertions(+), 29 deletions(-)

diff --git a/Makefile b/Makefile
index f69f36a..4c0cda4 100644
--- a/Makefile
+++ b/Makefile
@@ -68,6 +68,9 @@ providerlint:
 		-c 0 \
 		./...
 
+reset-token:
+	@sh -c 'cd util/megaport_token && go run . --reset'
+
 sweep:
 	@echo "WARNING: This will destroy infrastructure. Use only in development accounts."
 	go test $(TEST) -v -sweep=$(SWEEP) $(SWEEPARGS)
diff --git a/README.md b/README.md
index 0501439..f140fc8 100644
--- a/README.md
+++ b/README.md
@@ -4,12 +4,14 @@
 
 ## Utilities
 
-To grab a token for the megaport api, you can use the helper tool:
+To retrieve a new token for the megaport api:
+```
+$ eval $(make reset-token)
+```
 
+Alternatively, you can use the helper tool directly:
 ```
 $ cd util/megaport_token
 $ go run .
 ```
-
-To revoke a token (and get a new one) you can pass the `--reset` flag to the
-tool.
+To revoke the current token (and get a new one) you can pass the `--reset` flag.
diff --git a/util/megaport_token/main.go b/util/megaport_token/main.go
index 492c151..81707f2 100644
--- a/util/megaport_token/main.go
+++ b/util/megaport_token/main.go
@@ -1,11 +1,12 @@
 package main
 
 import (
-	"bufio"
 	"fmt"
 	"log"
 	"os"
+	"time"
 
+	"github.com/pquerna/otp/totp"
 	"github.com/utilitywarehouse/terraform-provider-megaport/megaport/api"
 )
 
@@ -14,40 +15,61 @@ const (
 )
 
 func main() {
+	var (
+		reset                          = false
+		endpoint                       = api.EndpointStaging
+		username, password, totpSecret string
+	)
 	if len(os.Args) > 2 {
 		log.Fatalln(usage)
 	}
 	if len(os.Args) == 2 {
 		switch os.Args[1] {
 		case "--reset":
-			token := os.Getenv("MEGAPORT_TOKEN")
-			if token == "" {
-				log.Fatal("To reset the token, please export MEGAPORT_TOKEN with your current token")
-			}
-			c := api.NewClient(api.EndpointStaging)
-			c.Token = token
-			if err := c.Logout(); err != nil {
-				log.Fatal(err)
-			}
-			log.Print("Current token has been reset. Please login again to fetch a new one.")
+			reset = true
+			log.Println("The token will be reset to retrieve a new one")
 		default:
 			log.Fatalln(usage)
 		}
 	}
-	scanner := bufio.NewScanner(os.Stdin)
-	var username, password, otp string
-	fmt.Printf("username: ")
-	scanner.Scan()
-	username = scanner.Text()
-	fmt.Printf("password: ")
-	scanner.Scan()
-	password = scanner.Text()
-	fmt.Printf("otp (leave empty if disabled): ")
-	scanner.Scan()
-	otp = scanner.Text()
-	c := api.NewClient(api.EndpointStaging)
+	if v := os.Getenv("MEGAPORT_ENDPOINT"); v != "" {
+		endpoint = v
+	}
+	log.Printf("Endpoint: %s\n", endpoint)
+	if username = os.Getenv("MEGAPORT_USERNAME"); username == "" {
+		log.Fatalln("MEGAPORT_USERNAME is empty")
+	}
+	if password = os.Getenv("MEGAPORT_PASSWORD"); password == "" {
+		log.Fatalln("MEGAPORT_PASSWORD is empty")
+	}
+	totpSecret = os.Getenv("MEGAPORT_TOTP_SECRET")
+	token, err := getToken(endpoint, username, password, totpSecret, reset)
+	if err != nil {
+		log.Fatalln(err)
+	}
+	fmt.Printf("MEGAPORT_TOKEN=%s\n", token)
+}
+
+func getToken(endpoint, username, password, totpSecret string, reset bool) (string, error) {
+	otp := ""
+	if totpSecret != "" {
+		v, err := totp.GenerateCode(totpSecret, time.Now())
+		if err != nil {
+			return "", err
+		}
+		otp = v
+	}
+	c := api.NewClient(endpoint)
+	if reset {
+		if err := c.Login(username, password, otp); err != nil {
+			return "", err
+		}
+		if err := c.Logout(); err != nil {
+			return "", err
+		}
+	}
 	if err := c.Login(username, password, otp); err != nil {
-		log.Fatal(err)
+		return "", err
 	}
-	fmt.Printf("MEGAPORT_TOKEN=%s\n", c.Token)
+	return c.Token, nil
 }