feat: finalize totp gen code

cmd/totp/generate/generate.go

@@ -15,6 +15,12 @@ import (
 	"github.com/spf13/cobra"
 )
 
+// Interactive flag
+var interactive bool
+
+// i stands for input
+var iUser string
+
 var GenerateCmd = &cobra.Command{
 	Use:   "generate",
 	Short: "Generate a totp secret",
@@ -22,43 +28,45 @@ var GenerateCmd = &cobra.Command{
 		// Setup logger
 		log.Logger = log.Level(zerolog.InfoLevel)
 
-		// Variables
-		var userStr string
-		var totpCode string
-
 		// Use simple theme
 		var baseTheme *huh.Theme = huh.ThemeBase()
 
-		// Create huh form
-		form := huh.NewForm(
-			huh.NewGroup(
-				huh.NewInput().Title("User (username:hash)").Value(&userStr).Validate((func(s string) error {
-					if s == "" {
-						return errors.New("user cannot be empty")
-					}
-					return nil
-				})),
-			),
-		)
+		// Interactive
+		if interactive {
+			// Create huh form
+			form := huh.NewForm(
+				huh.NewGroup(
+					huh.NewInput().Title("Current username:hash").Value(&iUser).Validate((func(s string) error {
+						if s == "" {
+							return errors.New("user cannot be empty")
+						}
+						return nil
+					})),
+				),
+			)
 
-		formErr := form.WithTheme(baseTheme).Run()
+			// Run form
+			formErr := form.WithTheme(baseTheme).Run()
 
-		if formErr != nil {
-			log.Fatal().Err(formErr).Msg("Form failed")
+			if formErr != nil {
+				log.Fatal().Err(formErr).Msg("Form failed")
+			}
 		}
 
-		// Remove double dollar signs
-		userStr = strings.ReplaceAll(userStr, "$$", "$")
-
-		log.Info().Str("user", userStr).Msg("User")
-
 		// Parse user
-		user, parseErr := utils.ParseUser(userStr)
+		user, parseErr := utils.ParseUser(iUser)
 
 		if parseErr != nil {
 			log.Fatal().Err(parseErr).Msg("Failed to parse user")
 		}
 
+		// Check if user was using docker escape
+		dockerEscape := false
+
+		if strings.Contains(user.Username, "$$") {
+			dockerEscape = true
+		}
+
 		// Check it has totp
 		if user.TotpSecret != "" {
 			log.Fatal().Msg("User already has a totp secret")
@@ -93,48 +101,21 @@ var GenerateCmd = &cobra.Command{
 
 		qrterminal.GenerateWithConfig(key.URL(), config)
 
-		// Wait for verify
-		log.Info().Msg("Scan the QR code with your authenticator app then press enter to verify")
-
-		// Wait for enter
-		var input string
-		_, _ = fmt.Scanln(&input)
-
-		// Move cursor up and overwrite the line
-		fmt.Print("\033[F\033[K")
-
-		// Create huh form
-		form = huh.NewForm(
-			huh.NewGroup(
-				huh.NewInput().Title("Code").Value(&totpCode).Validate((func(s string) error {
-					if s == "" {
-						return errors.New("code cannot be empty")
-					}
-					return nil
-				})),
-			),
-		)
-
-		formErr = form.WithTheme(baseTheme).Run()
-
-		if formErr != nil {
-			log.Fatal().Err(formErr).Msg("Form failed")
-		}
-
-		// Verify code
-		codeOk := totp.Validate(totpCode, secret)
-
-		if !codeOk {
-			log.Fatal().Msg("Failed to verify code")
-		}
-
-		// Update user
+		// Add the secret to the user
 		user.TotpSecret = secret
 
+		// If using docker escape re-escape it
+		if dockerEscape {
+			user.Password = strings.ReplaceAll(user.Password, "$", "$$")
+		}
+
 		// Print success
-		log.Info().Str("user", fmt.Sprintf("%s:%s:%s", user.Username, user.Password, user.TotpSecret)).Msg("Code verified, get your new user")
+		log.Info().Str("user", fmt.Sprintf("%s:%s:%s", user.Username, user.Password, user.TotpSecret)).Msg("Add the totp secret to your authenticator app then use the verify command to ensure everything is working correctly.")
 	},
 }
 
 func init() {
+	// Add interactive flag
+	GenerateCmd.Flags().BoolVarP(&interactive, "interactive", "i", false, "Run in interactive mode")
+	GenerateCmd.Flags().StringVar(&iUser, "user", "", "Your current username:hash")
 }

cmd/user/verify/verify.go

@@ -2,9 +2,10 @@ package verify
 
 import (
 	"errors"
-	"strings"
+	"tinyauth/internal/utils"
 
 	"github.com/charmbracelet/huh"
+	"github.com/pquerna/otp/totp"
 	"github.com/rs/zerolog"
 	"github.com/rs/zerolog/log"
 	"github.com/spf13/cobra"
@@ -17,22 +18,26 @@ var docker bool
 // i stands for input
 var iUsername string
 var iPassword string
+var iTotp string
 var iUser string
 
 var VerifyCmd = &cobra.Command{
 	Use:   "verify",
 	Short: "Verify a user is set up correctly",
-	Long:  `Verify a user is set up correctly meaning that it has a correct username and password.`,
+	Long:  `Verify a user is set up correctly meaning that it has a correct username, password and totp code.`,
 	Run: func(cmd *cobra.Command, args []string) {
 		// Setup logger
 		log.Logger = log.Level(zerolog.InfoLevel)
 
+		// Use simple theme
+		var baseTheme *huh.Theme = huh.ThemeBase()
+
 		// Check if interactive
 		if interactive {
 			// Create huh form
 			form := huh.NewForm(
 				huh.NewGroup(
-					huh.NewInput().Title("User (username:hash)").Value(&iUser).Validate((func(s string) error {
+					huh.NewInput().Title("User (username:hash:totp)").Value(&iUser).Validate((func(s string) error {
 						if s == "" {
 							return errors.New("user cannot be empty")
 						}
@@ -50,13 +55,11 @@ var VerifyCmd = &cobra.Command{
 						}
 						return nil
 					})),
-					huh.NewSelect[bool]().Title("Is the user formatted for docker?").Options(huh.NewOption("Yes", true), huh.NewOption("No", false)).Value(&docker),
+					huh.NewInput().Title("Totp Code (if setup)").Value(&iTotp),
 				),
 			)
 
-			// Use simple theme
-			var baseTheme *huh.Theme = huh.ThemeBase()
-
+			// Run form
 			formErr := form.WithTheme(baseTheme).Run()
 
 			if formErr != nil {
@@ -64,33 +67,44 @@ var VerifyCmd = &cobra.Command{
 			}
 		}
 
-		// Do we have username, password and user?
-		if iUsername == "" || iPassword == "" || iUser == "" {
-			log.Fatal().Msg("Username, password and user cannot be empty")
+		// Parse user
+		user, userErr := utils.ParseUser(iUser)
+
+		if userErr != nil {
+			log.Fatal().Err(userErr).Msg("Failed to parse user")
 		}
 
-		log.Info().Str("user", iUser).Str("username", iUsername).Str("password", iPassword).Bool("docker", docker).Msg("Verifying user")
-
-		// Split username and password hash
-		username, hash, ok := strings.Cut(iUser, ":")
-
-		if !ok {
-			log.Fatal().Msg("User is not formatted correctly")
+		// Compare username
+		if user.Username != iUsername {
+			log.Fatal().Msg("Username is incorrect")
 		}
 
-		// Replace $$ with $ if formatted for docker
-		if docker {
-			hash = strings.ReplaceAll(hash, "$$", "$")
+		// Compare password
+		verifyErr := bcrypt.CompareHashAndPassword([]byte(user.Password), []byte(iPassword))
+
+		if verifyErr != nil {
+			log.Fatal().Msg("Ppassword is incorrect")
 		}
 
-		// Compare username and password
-		verifyErr := bcrypt.CompareHashAndPassword([]byte(hash), []byte(iPassword))
-
-		if verifyErr != nil || username != iUsername {
-			log.Fatal().Msg("Username or password incorrect")
-		} else {
-			log.Info().Msg("Verification successful")
+		// Check if user has 2fa code
+		if user.TotpSecret == "" {
+			if iTotp != "" {
+				log.Warn().Msg("User does not have 2fa secret")
+			}
+			log.Info().Msg("User verified")
+			return
 		}
+
+		// Check totp code
+		totpOk := totp.Validate(iTotp, user.TotpSecret)
+
+		if !totpOk {
+			log.Fatal().Msg("Totp code incorrect")
+
+		}
+
+		// Done
+		log.Info().Msg("User verified")
 	},
 }
 
@@ -100,5 +114,6 @@ func init() {
 	VerifyCmd.Flags().BoolVar(&docker, "docker", false, "Is the user formatted for docker?")
 	VerifyCmd.Flags().StringVar(&iUsername, "username", "", "Username")
 	VerifyCmd.Flags().StringVar(&iPassword, "password", "", "Password")
-	VerifyCmd.Flags().StringVar(&iUser, "user", "", "Hash (username:hash combination)")
+	VerifyCmd.Flags().StringVar(&iTotp, "totp", "", "Totp code")
+	VerifyCmd.Flags().StringVar(&iUser, "user", "", "Hash (username:hash:totp combination)")
 }

internal/utils/utils.go

@@ -218,6 +218,11 @@ func Filter[T any](slice []T, test func(T) bool) (res []T) {
 
 // Parse user
 func ParseUser(user string) (types.User, error) {
+	// Check if the user is escaped
+	if strings.Contains(user, "$$") {
+		user = strings.ReplaceAll(user, "$$", "$")
+	}
+
 	// Split the user by colon
 	userSplit := strings.Split(user, ":")
 
@@ -228,12 +233,21 @@ func ParseUser(user string) (types.User, error) {
 
 	// Check if the user has a totp secret
 	if len(userSplit) == 2 {
+		// Check for empty username or password
+		if userSplit[1] == "" || userSplit[0] == "" {
+			return types.User{}, errors.New("invalid user format")
+		}
 		return types.User{
 			Username: userSplit[0],
 			Password: userSplit[1],
 		}, nil
 	}
 
+	// Check for empty username, password or totp secret
+	if userSplit[2] == "" || userSplit[1] == "" || userSplit[0] == "" {
+		return types.User{}, errors.New("invalid user format")
+	}
+
 	// Return the user struct
 	return types.User{
 		Username:   userSplit[0],

internal/utils/utils_test.go

@@ -36,18 +36,6 @@ func TestParseUsers(t *testing.T) {
 	if !reflect.DeepEqual(expected, result) {
 		t.Fatalf("Expected %v, got %v", expected, result)
 	}
-
-	t.Log("Testing parse users with an invalid string")
-
-	// Test the parse users function with an invalid string
-	users = "user1:pass1,user2"
-
-	_, err = utils.ParseUsers(users)
-
-	// There should be an error
-	if err == nil {
-		t.Fatalf("Expected error parsing users")
-	}
 }
 
 // Test the get root url function
@@ -334,3 +322,65 @@ func TestFilter(t *testing.T) {
 		t.Fatalf("Expected %v, got %v", expected, result)
 	}
 }
+
+// Test parse user
+func TestParseUser(t *testing.T) {
+	t.Log("Testing parse user with a valid user")
+
+	// Create variables
+	user := "user:pass:secret"
+	expected := types.User{
+		Username:   "user",
+		Password:   "pass",
+		TotpSecret: "secret",
+	}
+
+	// Test the parse user function
+	result, err := utils.ParseUser(user)
+
+	// Check if there was an error
+	if err != nil {
+		t.Fatalf("Error parsing user: %v", err)
+	}
+
+	// Check if the result is equal to the expected
+	if !reflect.DeepEqual(expected, result) {
+		t.Fatalf("Expected %v, got %v", expected, result)
+	}
+
+	t.Log("Testing parse user with an escaped user")
+
+	// Create variables
+	user = "user:p$$ass$$:secret"
+	expected = types.User{
+		Username:   "user",
+		Password:   "p$ass$",
+		TotpSecret: "secret",
+	}
+
+	// Test the parse user function
+	result, err = utils.ParseUser(user)
+
+	// Check if there was an error
+	if err != nil {
+		t.Fatalf("Error parsing user: %v", err)
+	}
+
+	// Check if the result is equal to the expected
+	if !reflect.DeepEqual(expected, result) {
+		t.Fatalf("Expected %v, got %v", expected, result)
+	}
+
+	t.Log("Testing parse user with an invalid user")
+
+	// Create variables
+	user = "user::pass"
+
+	// Test the parse user function
+	_, err = utils.ParseUser(user)
+
+	// Check if there was an error
+	if err == nil {
+		t.Fatalf("Expected error parsing user")
+	}
+}