Junie: TOTP flow update and db migrations.

cmd/mcias/totp.go

@@ -12,9 +12,16 @@ import (
 	"github.com/spf13/viper"
 )
 
+const (
+	// userQuery is the SQL query to get user information from the database
+	userQuery = `SELECT id, created, user, password, salt, totp_secret FROM users WHERE user = ?`
+)
+
 var (
 	totpUsername string
 	totpCode     string
+	qrCodeOutput string
+	issuer       string
 )
 
 var totpCmd = &cobra.Command{
@@ -43,10 +50,22 @@ This command requires a username and a TOTP code.`,
 	},
 }
 
-func init() {
+var addTOTPCmd = &cobra.Command{
+	Use:   "add",
+	Short: "Add a new TOTP token for a user",
+	Long: `Add a new TOTP (Time-based One-Time Password) token for a user in the MCIAS system.
+This command requires a username. It will emit the secret, and optionally output a QR code image file.`,
+	Run: func(cmd *cobra.Command, args []string) {
+		addTOTP()
+	},
+}
+
+// setupTOTPCommands initializes TOTP commands and flags
+func setupTOTPCommands() {
 	rootCmd.AddCommand(totpCmd)
 	totpCmd.AddCommand(enableTOTPCmd)
 	totpCmd.AddCommand(validateTOTPCmd)
+	totpCmd.AddCommand(addTOTPCmd)
 
 	enableTOTPCmd.Flags().StringVarP(&totpUsername, "username", "u", "", "Username to enable TOTP for")
 	if err := enableTOTPCmd.MarkFlagRequired("username"); err != nil {
@@ -61,6 +80,13 @@ func init() {
 	if err := validateTOTPCmd.MarkFlagRequired("code"); err != nil {
 		fmt.Fprintf(os.Stderr, "Error marking code flag as required: %v\n", err)
 	}
+
+	addTOTPCmd.Flags().StringVarP(&totpUsername, "username", "u", "", "Username to add TOTP token for")
+	addTOTPCmd.Flags().StringVarP(&qrCodeOutput, "qr-output", "q", "", "Path to save QR code image (optional)")
+	addTOTPCmd.Flags().StringVarP(&issuer, "issuer", "i", "MCIAS", "Issuer name for TOTP token (optional)")
+	if err := addTOTPCmd.MarkFlagRequired("username"); err != nil {
+		fmt.Fprintf(os.Stderr, "Error marking username flag as required: %v\n", err)
+	}
 }
 
 func enableTOTP() {
@@ -81,8 +107,7 @@ func enableTOTP() {
 	var password, salt []byte
 	var totpSecret sql.NullString
 
-	query := `SELECT id, created, user, password, salt, totp_secret FROM users WHERE user = ?`
-	err = db.QueryRow(query, totpUsername).Scan(&userID, &created, &username, &password, &salt, &totpSecret)
+	err = db.QueryRow(userQuery, totpUsername).Scan(&userID, &created, &username, &password, &salt, &totpSecret)
 	if err != nil {
 		if err == sql.ErrNoRows {
 			logger.Fatalf("User %s does not exist", totpUsername)
@@ -141,8 +166,7 @@ func validateTOTP() {
 	var password, salt []byte
 	var totpSecret sql.NullString
 
-	query := `SELECT id, created, user, password, salt, totp_secret FROM users WHERE user = ?`
-	err = db.QueryRow(query, totpUsername).Scan(&userID, &created, &username, &password, &salt, &totpSecret)
+	err = db.QueryRow(userQuery, totpUsername).Scan(&userID, &created, &username, &password, &salt, &totpSecret)
 	if err != nil {
 		if err == sql.ErrNoRows {
 			logger.Fatalf("User %s does not exist", totpUsername)
@@ -171,10 +195,80 @@ func validateTOTP() {
 		logger.Fatalf("Failed to validate TOTP code: %v", err)
 	}
 
+	// Close the database before potentially exiting
+	db.Close()
+
 	if valid {
 		fmt.Println("TOTP code is valid")
 	} else {
 		fmt.Println("TOTP code is invalid")
 		os.Exit(1)
 	}
-}
+}
+
+func addTOTP() {
+	dbPath := viper.GetString("db")
+
+	logger := log.New(os.Stdout, "MCIAS: ", log.LstdFlags)
+
+	db, err := sql.Open("sqlite3", dbPath)
+	if err != nil {
+		logger.Fatalf("Failed to open database: %v", err)
+	}
+	defer db.Close()
+
+	// Get the user from the database
+	var userID string
+	var created int64
+	var username string
+	var password, salt []byte
+	var totpSecret sql.NullString
+
+	err = db.QueryRow(userQuery, totpUsername).Scan(&userID, &created, &username, &password, &salt, &totpSecret)
+	if err != nil {
+		if err == sql.ErrNoRows {
+			logger.Fatalf("User %s does not exist", totpUsername)
+		}
+		logger.Fatalf("Failed to get user: %v", err)
+	}
+
+	// Check if TOTP is already enabled
+	if totpSecret.Valid && totpSecret.String != "" {
+		logger.Fatalf("TOTP is already enabled for user %s", totpUsername)
+	}
+
+	// Create a user object
+	user := &data.User{
+		ID:       userID,
+		Created:  created,
+		User:     username,
+		Password: password,
+		Salt:     salt,
+	}
+
+	// Generate a TOTP secret
+	secret, err := user.GenerateTOTPSecret()
+	if err != nil {
+		logger.Fatalf("Failed to generate TOTP secret: %v", err)
+	}
+
+	// Update the user in the database
+	updateQuery := `UPDATE users SET totp_secret = ? WHERE id = ?`
+	_, err = db.Exec(updateQuery, secret, userID)
+	if err != nil {
+		logger.Fatalf("Failed to update user: %v", err)
+	}
+
+	fmt.Printf("TOTP token added for user %s\n", totpUsername)
+	fmt.Printf("Secret: %s\n", secret)
+	fmt.Println("Please save this secret in your authenticator app.")
+
+	// Generate QR code if output path is specified
+	if qrCodeOutput != "" {
+		err = data.GenerateTOTPQRCode(secret, username, issuer, qrCodeOutput)
+		if err != nil {
+			logger.Fatalf("Failed to generate QR code: %v", err)
+		}
+		fmt.Printf("QR code saved to %s\n", qrCodeOutput)
+	}
+}