From 9f5f4adddbd2b58549fbd989aa453013d424a6d7 Mon Sep 17 00:00:00 2001 From: Stavros Date: Thu, 6 Mar 2025 16:41:57 +0200 Subject: [PATCH] feat: finalize totp gen code --- cmd/totp/generate/generate.go | 103 ++++++++++++++-------------------- cmd/user/verify/verify.go | 71 ++++++++++++++--------- internal/utils/utils.go | 14 +++++ internal/utils/utils_test.go | 74 ++++++++++++++++++++---- 4 files changed, 161 insertions(+), 101 deletions(-) diff --git a/cmd/totp/generate/generate.go b/cmd/totp/generate/generate.go index fc382a9..cc82da8 100644 --- a/cmd/totp/generate/generate.go +++ b/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") } diff --git a/cmd/user/verify/verify.go b/cmd/user/verify/verify.go index ba64fb4..138bf91 100644 --- a/cmd/user/verify/verify.go +++ b/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)") } diff --git a/internal/utils/utils.go b/internal/utils/utils.go index 98dda21..1e68aee 100644 --- a/internal/utils/utils.go +++ b/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], diff --git a/internal/utils/utils_test.go b/internal/utils/utils_test.go index 36a1c1b..b3774ce 100644 --- a/internal/utils/utils_test.go +++ b/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") + } +}