mcias-junie/cmd/mcias/totp.go

package main

import (
	"database/sql"
	"errors"
	"fmt"
	"log"
	"os"

	"git.wntrmute.dev/kyle/mcias/data"
	_ "github.com/mattn/go-sqlite3"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

const (
	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{
	Use:   "totp",
	Short: "Manage TOTP authentication",
	Long:  `Commands for managing TOTP (Time-based One-Time Password) authentication in the MCIAS system.`,
}

var enableTOTPCmd = &cobra.Command{
	Use:   "enable",
	Short: "Enable TOTP for a user",
	Long: `Enable TOTP (Time-based One-Time Password) authentication for a user in the MCIAS system.
This command requires a username.`,
	Run: func(cmd *cobra.Command, args []string) {
		enableTOTP()
	},
}

var validateTOTPCmd = &cobra.Command{
	Use:   "validate",
	Short: "Validate a TOTP code",
	Long: `Validate a TOTP code for a user in the MCIAS system.
This command requires a username and a TOTP code.`,
	Run: func(cmd *cobra.Command, args []string) {
		validateTOTP()
	},
}

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()
	},
}

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 {
		fmt.Fprintf(os.Stderr, "Error marking username flag as required: %v\n", err)
	}

	validateTOTPCmd.Flags().StringVarP(&totpUsername, "username", "u", "", "Username to validate TOTP code for")
	validateTOTPCmd.Flags().StringVarP(&totpCode, "code", "c", "", "TOTP code to validate")
	if err := validateTOTPCmd.MarkFlagRequired("username"); err != nil {
		fmt.Fprintf(os.Stderr, "Error marking username flag as required: %v\n", err)
	}
	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() {
	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()

	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 errors.Is(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 enabled for user %s\n", totpUsername)
	fmt.Printf("Secret: %s\n", secret)
	fmt.Println("Please save this secret in your authenticator app.")
	fmt.Println("You will need to provide a TOTP code when logging in.")
}

func validateTOTP() {
	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 errors.Is(err, sql.ErrNoRows) {
			logger.Fatalf("User %s does not exist", totpUsername)
		}
		logger.Fatalf("Failed to get user: %v", err)
	}

	// Check if TOTP is enabled
	if !totpSecret.Valid || totpSecret.String == "" {
		logger.Fatalf("TOTP is not enabled for user %s", totpUsername)
	}

	// Create a user object
	user := &data.User{
		ID:         userID,
		Created:    created,
		User:       username,
		Password:   password,
		Salt:       salt,
		TOTPSecret: totpSecret.String,
	}

	// Validate the TOTP code
	valid, err := user.ValidateTOTPCode(totpCode)
	if err != nil {
		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 errors.Is(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)
	}
}