add --follow flag (#9)

README.md

@@ -63,6 +63,12 @@ $ totp config reset
 $ totp --secret seed
 ```
 
+**Continuous code output** can be generated with the `--follow` option.
+
+```
+$ totp --follow mysecretname
+```
+
 **For help** on any of the above, use the `--help` option. Examples are
 
 ```
@@ -104,7 +110,13 @@ $ totp --time 2019-06-01T20:00:00-05:00 --forward 30s --secret seed
 820148
 ```
 
-## Using the Sdio Option
+The `--follow` option is also compatible with the time machine.
+
+```
+$ totp --time 2001-10-31T20:00:00-05:00 --follow --secret seed
+```
+
+## Using the Stdio Option
 
 If storing secrets in the clear isn't ideal for you, `totp` supports streaming the shared secret collection through stdin and stdout with the `--stdio` option. This allows you to roll your own encryption or support other methods of maintaining shared secrets.
 

cmd/root.go

@@ -12,14 +12,19 @@ import (
 )
 
 const (
+	optionBackward = "backward"
 	optionFile     = "file"
+	optionFollow   = "follow"
+	optionForward  = "forward"
 	optionSecret   = "secret"
 	optionStdio    = "stdio"
 	optionTime     = "time"
-	optionBackward = "backward"
-	optionForward  = "forward"
 )
 
+type generateCodesAPI func(time.Duration, time.Duration, time.Duration, func(time.Duration), string, string)
+
+var generateCodesService generateCodesAPI
+
 var rootCmd = &cobra.Command{
 	Use:   "totp",
 	Short: "TOTP Generator",
@@ -78,9 +83,13 @@ var rootCmd = &cobra.Command{
 			return
 		}
 
-		// Default to time of now
-		codeTime := time.Now().Add(-backward)
-		codeTime = codeTime.Add(forward)
+		follow, err := cmd.Flags().GetBool(optionFollow)
+		if err != nil {
+			fmt.Println("Error processing follow option", err)
+			return
+		}
+
+		var codeTime time.Time
 
 		// Override if time was given
 		if len(timeString) > 0 {
@@ -89,9 +98,9 @@ var rootCmd = &cobra.Command{
 				fmt.Println("Error parsing the time option", err)
 				return
 			}
-
-			codeTime = codeTime.Add(-backward)
-			codeTime = codeTime.Add(forward)
+		} else {
+			codeTime = time.Now()
+			// codeOffset is 0
 		}
 
 		secretLen := len(secret)
@@ -128,7 +137,11 @@ var rootCmd = &cobra.Command{
 		}
 
 		// If here then a stored shared secret is wanted
-		generateCode(secretName, secret, codeTime)
+		generateCode(secretName, secret, codeTime.Add(forward-backward))
+
+		if follow {
+			generateCodesService(codeTime.Sub(time.Now())-backward+forward, 0, 30*time.Second, time.Sleep, secretName, secret)
+		}
 	},
 }
 
@@ -160,6 +173,61 @@ func generateCode(name string, secret string, t time.Time) error {
 	return err
 }
 
+func durationToNextInterval(now time.Time) time.Duration {
+	var sleepSeconds int
+
+	s := now.Second()
+	switch {
+	case s == 0, s < 30:
+		sleepSeconds = 30 - s
+	case s >= 30:
+		sleepSeconds = 60 - s
+	}
+
+	return time.Duration(sleepSeconds)*time.Second -
+		time.Duration(now.Nanosecond())*time.Nanosecond
+
+}
+
+func callOnInterval(runtime time.Duration, interval time.Duration, exec func() bool) {
+	stopper := make(chan bool)
+
+	if runtime > 0 {
+		go func() {
+			time.Sleep(runtime)
+			stopper <- true
+		}()
+	}
+
+	if exec != nil && exec() {
+		return
+	}
+
+	ticker := time.NewTicker(interval)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-stopper:
+			return
+		case <-ticker.C:
+			if exec != nil && exec() {
+				return
+			}
+		}
+	}
+}
+
+func generateCodes(timeOffset time.Duration, durationToRun time.Duration, intervalTime time.Duration, sleep func(time.Duration), secretName, secret string) {
+	sleep(durationToNextInterval(time.Now().Add(timeOffset)) + 10*time.Millisecond)
+
+	callOnInterval(durationToRun, intervalTime,
+		func() bool {
+			generateCode(secretName, secret, time.Now().Add(timeOffset))
+			return false
+		})
+}
+
 // Execute adds all child commands to the root command and sets flags appropriately.
 // This is called by main.main(). It only needs to happen once to the rootCmd.
 func Execute() int {
@@ -176,13 +244,16 @@ func Execute() int {
 func init() {
 	var duration time.Duration
 
+	generateCodesService = generateCodes
+
 	rootCmd.PersistentFlags().StringP(optionFile, "f", "", "secret collection file")
 
 	rootCmd.Flags().StringP(optionSecret, "s", "", "TOTP secret value")
 	rootCmd.Flags().BoolP(optionStdio, "", false, "load with stdin")
 	rootCmd.Flags().StringP(optionTime, "", "", "RFC3339 time for TOTP (2019-06-23T20:00:00-05:00)")
 	rootCmd.Flags().DurationP(optionBackward, "", duration, "move time backward (ex. \"30s\")")
 	rootCmd.Flags().DurationP(optionForward, "", duration, "move time forward (ex. \"1m\")")
+	rootCmd.Flags().BoolP(optionFollow, "", false, "continuous output")
 
 	rootCmd.SetUsageTemplate(strings.Replace(rootCmd.UsageTemplate(), "{{.UseLine}}", "{{.UseLine}}\n  {{.CommandPath}} [secret name]", 1))
 }

cmd/root_test.go

@@ -3,6 +3,7 @@ package cmd
 import (
 	"os"
 	"testing"
+	"time"
 
 	"github.com/spf13/pflag"
 )
@@ -35,6 +36,14 @@ func TestRoot(t *testing.T) {
 	// Non-existing entry
 	rootCmd.Run(rootCmd, []string{"invalidsecret"})
 
+	// Test follow condition
+	savedGenerateCodesService := generateCodesService
+	generateCodesService = func(time.Duration, time.Duration, time.Duration, func(time.Duration), string, string) {}
+	rootCmd.Flags().Lookup(optionFollow).Value.Set("true")
+	rootCmd.Run(rootCmd, []string{"name0"})
+	generateCodesService = savedGenerateCodesService
+	rootCmd.Flags().Lookup(optionFollow).Value.Set("false")
+
 	// No collections file
 	os.Remove(collectionFile.filename)
 	rootCmd.Run(rootCmd, []string{secretList[0].name})
@@ -58,6 +67,9 @@ func TestRoot(t *testing.T) {
 	// Stdio option
 	rootCmd.Flags().Set(optionStdio, "true")
 	rootCmd.PersistentPreRun(rootCmd, []string{"secret"})
+	collectionFile.loader = loadCollectionFromDefaultFile
+	collectionFile.useStdio = false
+	rootCmd.Flags().Set(optionStdio, "")
 
 	// Time option
 	rootCmd.Flags().Set(optionTime, "2019-06-01T20:00:00-05:00")
@@ -66,10 +78,12 @@ func TestRoot(t *testing.T) {
 	// Give secret and secret name
 	rootCmd.Flags().Set(optionSecret, "seed")
 	rootCmd.Run(rootCmd, []string{"secretname"})
+	rootCmd.Flags().Set(optionSecret, "")
 
 	// Invalid time option
 	rootCmd.Flags().Set(optionTime, "invalidtime")
 	rootCmd.Run(rootCmd, []string{})
+	rootCmd.Flags().Set(optionTime, "")
 
 	var f *pflag.Flag
 	var savedFlagValue pflag.Value
@@ -117,9 +131,70 @@ func TestRoot(t *testing.T) {
 	rootCmd.Run(rootCmd, []string{})
 	f.Value = savedFlagValue
 
+	// optionFollow error
+	f = rootCmd.Flags().Lookup(optionFollow)
+	savedFlagValue = f.Value
+	f.Value = new(flagValue)
+	rootCmd.Run(rootCmd, []string{})
+	f.Value = savedFlagValue
+
 	Execute()
 	savedArgs := os.Args
 	os.Args = []string{"totp", "--invalidoption"}
 	Execute()
 	os.Args = savedArgs
 }
+
+func TestExecOnInterval(t *testing.T) {
+	execCount := 0
+	preAndExecNormal := func() bool { return false }
+	preAndExecEarlyExit := func() bool { execCount++; return execCount == 2 }
+
+	// Normal execution
+	execCount = 0
+	callOnInterval(2*time.Millisecond, 1*time.Millisecond, preAndExecNormal)
+
+	// Exit at preExec
+	execCount = 1
+	callOnInterval(2*time.Millisecond, 1*time.Millisecond, preAndExecNormal)
+
+	// Exit at top exec
+	execCount = 1
+	callOnInterval(2*time.Millisecond, 1*time.Millisecond, preAndExecEarlyExit)
+
+	// Exit at loop exec
+	execCount = 0
+	callOnInterval(2*time.Millisecond, 1*time.Millisecond, preAndExecEarlyExit)
+
+	// No callbacks
+	callOnInterval(2*time.Millisecond, 1*time.Millisecond, nil)
+}
+
+func TestDurationToNextInterval(t *testing.T) {
+	now, _ := time.Parse(time.RFC3339, "2019-06-23T20:00:01-05:00")
+	expectedResult := time.Duration(29 * time.Second)
+	actualResult := durationToNextInterval(now)
+	if expectedResult != actualResult {
+		t.Errorf("durationToNextInterval(%s) expected %s but returned %s", now, expectedResult, actualResult)
+	}
+
+	now, _ = time.Parse(time.RFC3339, "2019-06-23T20:00:31-05:00")
+	expectedResult = time.Duration(29 * time.Second)
+	actualResult = durationToNextInterval(now)
+	if expectedResult != actualResult {
+		t.Errorf("durationToNextInterval(%s) expected %s but returned %s", now, expectedResult, actualResult)
+	}
+
+	now, _ = time.Parse(time.RFC3339, "2019-06-23T20:00:31.001-05:00")
+	expectedResult = time.Duration(28*time.Second + 999*time.Millisecond)
+	actualResult = durationToNextInterval(now)
+	if expectedResult != actualResult {
+		t.Errorf("durationToNextInterval(%s) expected %s but returned %s", now, expectedResult, actualResult)
+	}
+}
+
+func TestGenerateCodes(t *testing.T) {
+	var d time.Duration
+	generateCodes(d, 2*time.Millisecond, 1*time.Millisecond,
+		func(d time.Duration) {}, "", "seed")
+}