feat: add sha256 and sha512 support (#135)

yitsushi · web-flow · commit 28337f38af94 · 2025-02-26T20:04:39.000+01:00
* feat: add sha256 and sha512 support As the original implementation followed only RFC-4226, it had only SHA1 support. This covers a lot of tokens, but secure tokens are not SHA1, but SHA2 based as described in RFC-6238. Until now, I though only hardware tokens use SHA2 (SHA256, SHA512) and maybe some edge cases, but not common. That was probably true almost 10 years ago when I originally created this tool. Today, even tho I didn't see 256 or 512 tokens yet, I think it's more common, and soon they will be used in a lot of places. Resolves #133 References: - #133 - https://www.ietf.org/rfc/rfc4226.txt - https://www.ietf.org/rfc/rfc6238.txt Signed-off-by: Victoria Nadasdi <victoria@efertone.me>
diff --git a/internal/cmd/add_token.go b/internal/cmd/add_token.go
@@ -17,16 +17,9 @@ func AddTokenCommand() *cli.Command {
 		Usage:     "Add new token.",
 		ArgsUsage: "[namespace] [account]",
 		Flags: []cli.Flag{
-			&cli.UintFlag{
-				Name:  "length",
-				Value: s.DefaultTokenLength,
-				Usage: "Length of the generated token.",
-			},
-			&cli.StringFlag{
-				Name:  "prefix",
-				Value: "",
-				Usage: "Prefix for the token.",
-			},
+			flagLength(),
+			flagPrefix(),
+			flagAlgorithm(),
 		},
 		Action: func(ctx *cli.Context) error {
 			var (
@@ -54,7 +47,13 @@ func AddTokenCommand() *cli.Command {
 				}
 			}
 
-			account = &s.Account{Name: accName, Token: token, Prefix: ctx.String("prefix"), Length: ctx.Uint("length")}
+			account = &s.Account{
+				Name:      accName,
+				Token:     token,
+				Prefix:    ctx.String("prefix"),
+				Length:    ctx.Uint("length"),
+				Algorithm: ctx.String("algorithm"),
+			}
 			namespace.Accounts = append(namespace.Accounts, account)
 
 			err = storage.Save()
diff --git a/internal/cmd/dump.go b/internal/cmd/dump.go
@@ -18,16 +18,8 @@ func DumpCommand() *cli.Command {
 		Usage:     "Dump all available accounts under all namespaces.",
 		ArgsUsage: " ",
 		Flags: []cli.Flag{
-			&cli.BoolFlag{
-				Name:  "yes-please",
-				Value: false,
-				Usage: warningMsg,
-			},
-			&cli.StringFlag{
-				Name:     "output",
-				Usage:    "Output file. (REQUIRED)",
-				Required: true,
-			},
+			flagYesPlease(warningMsg),
+			flagOutput(),
 		},
 		Action: func(ctx *cli.Context) error {
 			if !ctx.Bool("yes-please") {
diff --git a/internal/cmd/error.go b/internal/cmd/error.go
@@ -21,3 +21,16 @@ func (e CommandError) Error() string {
 func resourceNotFoundError(name string) CommandError {
 	return CommandError{Message: name + " does not exist"}
 }
+
+// FlagError is an error during flag parsing.
+type FlagError struct {
+	Message string
+}
+
+func (e FlagError) Error() string {
+	return "flag error: %s" + e.Message
+}
+
+func invalidAlgorithmError(value string) FlagError {
+	return FlagError{Message: "Invalid algorithm: " + value}
+}
diff --git a/internal/cmd/flags.go b/internal/cmd/flags.go
@@ -0,0 +1,85 @@
+package cmd
+
+import (
+	"github.com/urfave/cli/v2"
+	"github.com/yitsushi/totp-cli/internal/storage"
+)
+
+func flagAlgorithm() *cli.StringFlag {
+	return &cli.StringFlag{
+		Name:  "algorithm",
+		Value: "sha1",
+		Usage: "Algorithm to use for HMAC (sha1, sha256, sha512).",
+		Action: func(_ *cli.Context, value string) error {
+			if value != "sha1" && value != "sha256" && value != "sha512" {
+				return invalidAlgorithmError(value)
+			}
+
+			return nil
+		},
+	}
+}
+
+func flagShowRemaining() *cli.BoolFlag {
+	return &cli.BoolFlag{
+		Name:  "show-remaining",
+		Value: false,
+		Usage: "Show how much time left until the code will be invalid.",
+	}
+}
+
+func flagLength() *cli.UintFlag {
+	return &cli.UintFlag{
+		Name:  "length",
+		Value: storage.DefaultTokenLength,
+		Usage: "Length of the generated token.",
+	}
+}
+
+func flagPrefix() *cli.StringFlag {
+	return &cli.StringFlag{
+		Name:  "prefix",
+		Value: "",
+		Usage: "Prefix for the token.",
+	}
+}
+
+func flagYesPlease(usage string) *cli.BoolFlag {
+	return &cli.BoolFlag{
+		Name:  "yes-please",
+		Value: false,
+		Usage: usage,
+	}
+}
+
+func flagOutput() *cli.StringFlag {
+	return &cli.StringFlag{
+		Name:     "output",
+		Usage:    "Output file. (REQUIRED)",
+		Required: true,
+	}
+}
+
+func flagInput() *cli.StringFlag {
+	return &cli.StringFlag{
+		Name:     "input",
+		Usage:    "Input YAML file. (REQUIRED)",
+		Required: true,
+	}
+}
+
+func flagFollow() *cli.BoolFlag {
+	return &cli.BoolFlag{
+		Name:  "follow",
+		Value: false,
+		Usage: "Generate codes continuously.",
+	}
+}
+
+func flagClearPrefix() *cli.BoolFlag {
+	return &cli.BoolFlag{
+		Name:  "clear",
+		Value: false,
+		Usage: "Clear prefix from account.",
+	}
+}
diff --git a/internal/cmd/generate.go b/internal/cmd/generate.go
@@ -14,16 +14,8 @@ func GenerateCommand() *cli.Command {
 		Name:    "generate",
 		Aliases: []string{"g"},
 		Flags: []cli.Flag{
-			&cli.BoolFlag{
-				Name:  "follow",
-				Value: false,
-				Usage: "Generate codes continuously.",
-			},
-			&cli.BoolFlag{
-				Name:  "show-remaining",
-				Value: false,
-				Usage: "Show how much time left until the code will be invalid.",
-			},
+			flagFollow(),
+			flagShowRemaining(),
 		},
 		Usage:     "Generate a specific OTP",
 		ArgsUsage: "<namespace> <account>",
diff --git a/internal/cmd/generatehelper.go b/internal/cmd/generatehelper.go
@@ -5,6 +5,7 @@ import (
 	"time"
 
 	"github.com/yitsushi/totp-cli/internal/security"
+	"github.com/yitsushi/totp-cli/internal/security/algo"
 	s "github.com/yitsushi/totp-cli/internal/storage"
 )
 
@@ -17,7 +18,20 @@ func formatCode(code string, remaining int64, showRemaining bool) string {
 }
 
 func generateCode(account *s.Account) (string, int64) {
-	code, remaining, err := security.GenerateOTPCode(account.Token, time.Now(), account.Length)
+	var algorithm algo.Algorithm
+
+	switch account.Algorithm {
+	case "sha1":
+		algorithm = algo.SHA1{}
+	case "sha256":
+		algorithm = algo.SHA256{}
+	case "sha512":
+		algorithm = algo.SHA512{}
+	default:
+		algorithm = algo.Default{}
+	}
+
+	code, remaining, err := security.GenerateOTPCode(account.Token, time.Now(), account.Length, algorithm)
 	if err != nil {
 		fmt.Printf("Error: %s\n", err.Error())
 
diff --git a/internal/cmd/import.go b/internal/cmd/import.go
@@ -16,11 +16,7 @@ func ImportCommand() *cli.Command {
 		Name:  "import",
 		Usage: "Import tokens from a yaml file.",
 		Flags: []cli.Flag{
-			&cli.StringFlag{
-				Name:     "input",
-				Usage:    "Input YAML file. (REQUIRED)",
-				Required: true,
-			},
+			flagInput(),
 		},
 		Action: func(ctx *cli.Context) (err error) {
 			var file []byte
diff --git a/internal/cmd/instant.go b/internal/cmd/instant.go
@@ -3,10 +3,8 @@ package cmd
 import (
 	"fmt"
 	"os"
-	"time"
 
 	"github.com/urfave/cli/v2"
-	"github.com/yitsushi/totp-cli/internal/security"
 	"github.com/yitsushi/totp-cli/internal/storage"
 	"github.com/yitsushi/totp-cli/internal/terminal"
 )
@@ -18,31 +16,25 @@ func InstantCommand() *cli.Command {
 		Usage:     "Generate an OTP from TOTP_TOKEN or stdin without the Storage backend.",
 		ArgsUsage: " ",
 		Flags: []cli.Flag{
-			&cli.UintFlag{
-				Name:  "length",
-				Value: storage.DefaultTokenLength,
-				Usage: "Length of the generated token.",
-			},
-			&cli.BoolFlag{
-				Name:  "show-remaining",
-				Value: false,
-				Usage: "Show how much time left until the code will be invalid.",
-			},
+			flagLength(),
+			flagShowRemaining(),
+			flagAlgorithm(),
 		},
 		Action: func(ctx *cli.Context) error {
-			token := os.Getenv("TOTP_TOKEN")
-			if token == "" {
-				term := terminal.New(os.Stdin, os.Stdout, os.Stderr)
-				token, _ = term.Read("")
+			account := storage.Account{
+				Name:      "instant",
+				Token:     os.Getenv("TOTP_TOKEN"),
+				Length:    ctx.Uint("length"),
+				Algorithm: ctx.String("algorithm"),
 			}
 
-			length := ctx.Uint("length")
-
-			code, remaining, err := security.GenerateOTPCode(token, time.Now(), length)
-			if err != nil {
-				return err
+			if account.Token == "" {
+				term := terminal.New(os.Stdin, os.Stdout, os.Stderr)
+				account.Token, _ = term.Read("")
 			}
 
+			code, remaining := generateCode(&account)
+
 			fmt.Println(formatCode(code, remaining, ctx.Bool("show-remaining")))
 
 			return nil
diff --git a/internal/cmd/set_prefix.go b/internal/cmd/set_prefix.go
@@ -16,11 +16,7 @@ func SetPrefixCommand() *cli.Command {
 		Usage:     "Set prefix for a token.",
 		ArgsUsage: "[namespace] [account] [prefix]",
 		Flags: []cli.Flag{
-			&cli.BoolFlag{
-				Name:  "clear",
-				Value: false,
-				Usage: "Clear prefix from account.",
-			},
+			flagClearPrefix(),
 		},
 		Action: func(ctx *cli.Context) (err error) {
 			var (
diff --git a/internal/security/algo/default.go b/internal/security/algo/default.go
@@ -0,0 +1,4 @@
+package algo
+
+// Default is the default algorithm to use.
+type Default = SHA1
diff --git a/internal/security/algo/main.go b/internal/security/algo/main.go
@@ -0,0 +1,12 @@
+package algo
+
+import "hash"
+
+// HasherFn is a function that returns a new hash.Hash.
+type HasherFn func() hash.Hash
+
+// Algorithm is an interface that defines the methods that an algorithm must implement.
+type Algorithm interface {
+	// Hasher returns a function that returns a new hash.Hash.
+	Hasher() HasherFn
+}
diff --git a/internal/security/algo/sha1.go b/internal/security/algo/sha1.go
@@ -0,0 +1,13 @@
+package algo
+
+import "crypto/sha1" //nolint:gosec // RFC-4226 defined SHA-1 as algorithm.
+
+// SHA1 is a struct that implements the Algorithm interface for SHA1.
+type SHA1 struct{}
+
+// Hasher returns a function that returns a new hash.Hash.
+func (s SHA1) Hasher() HasherFn {
+	return sha1.New
+}
+
+var _ Algorithm = SHA1{}
diff --git a/internal/security/algo/sha256.go b/internal/security/algo/sha256.go
@@ -0,0 +1,13 @@
+package algo
+
+import "crypto/sha256"
+
+// SHA256 is a struct that implements the Algorithm interface for SHA256.
+type SHA256 struct{}
+
+// Hasher returns a function that returns a new hash.Hash.
+func (s SHA256) Hasher() HasherFn {
+	return sha256.New
+}
+
+var _ Algorithm = SHA256{}
diff --git a/internal/security/algo/sha512.go b/internal/security/algo/sha512.go
@@ -0,0 +1,13 @@
+package algo
+
+import "crypto/sha512"
+
+// SHA512 is a struct that implements the Algorithm interface for SHA512.
+type SHA512 struct{}
+
+// Hasher returns a function that returns a new hash.Hash.
+func (s SHA512) Hasher() HasherFn {
+	return sha512.New
+}
+
+var _ Algorithm = SHA512{}
diff --git a/internal/security/otp.go b/internal/security/otp.go
@@ -2,13 +2,14 @@ package security
 
 import (
 	"crypto/hmac"
-	"crypto/sha1" //nolint:gosec // This is an implementation of an RFC that used SHA-1
 	"encoding/base32"
 	"encoding/binary"
 	"fmt"
 	"math"
 	"strings"
 	"time"
+
+	"github.com/yitsushi/totp-cli/internal/security/algo"
 )
 
 const (
@@ -24,7 +25,7 @@ const (
 )
 
 // GenerateOTPCode generates a 6 digit TOTP from the secret Token.
-func GenerateOTPCode(token string, when time.Time, length uint) (string, int64, error) {
+func GenerateOTPCode(token string, when time.Time, length uint, algorithm algo.Algorithm) (string, int64, error) {
 	timer := uint64(math.Floor(float64(when.Unix()) / float64(timeSplitInSeconds)))
 	remainingTime := timeSplitInSeconds - when.Unix()%timeSplitInSeconds
 
@@ -46,7 +47,7 @@ func GenerateOTPCode(token string, when time.Time, length uint) (string, int64,
 	}
 
 	buf := make([]byte, sumByteLength)
-	mac := hmac.New(sha1.New, secretBytes)
+	mac := hmac.New(algorithm.Hasher(), secretBytes)
 
 	binary.BigEndian.PutUint64(buf, timer)
 	_, _ = mac.Write(buf)
diff --git a/internal/security/otp_test.go b/internal/security/otp_test.go
diff --git a/internal/storage/account.go b/internal/storage/account.go
diff --git a/main.go b/main.go
