Junie: TOTP flow update and db migrations.

api/auth.go

@@ -22,6 +22,12 @@ type TokenResponse struct {
 	Expires int64  `json:"expires"`
 }
 
+type TOTPVerifyRequest struct {
+	Version  string `json:"version"`
+	Username string `json:"username"`
+	TOTPCode string `json:"totp_code"`
+}
+
 type ErrorResponse struct {
 	Error string `json:"error"`
 }
@@ -57,17 +63,35 @@ func (s *Server) handlePasswordLogin(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	// Check password and TOTP if enabled
-	if !user.Check(&req.Login) {
-		// If TOTP is enabled but no code was provided, return a special error
-		if user.HasTOTP() && req.Login.TOTPCode == "" {
-			s.sendError(w, "TOTP code required", http.StatusUnauthorized)
-			return
-		}
+	// Check password only first
+	if !user.CheckPassword(&req.Login) {
 		s.sendError(w, "Invalid username or password", http.StatusUnauthorized)
 		return
 	}
 
+	// If TOTP is enabled and a code was provided, verify it
+	if user.HasTOTP() {
+		if req.Login.TOTPCode == "" {
+			// TOTP is enabled but no code was provided
+			// Return a special response indicating TOTP is required
+			w.Header().Set("Content-Type", "application/json")
+			w.WriteHeader(http.StatusUnauthorized)
+			if err := json.NewEncoder(w).Encode(ErrorResponse{
+				Error: "TOTP code required",
+			}); err != nil {
+				s.Logger.Printf("Error encoding response: %v", err)
+			}
+			return
+		}
+
+		// Validate the TOTP code
+		valid, validErr := user.ValidateTOTPCode(req.Login.TOTPCode)
+		if validErr != nil || !valid {
+			s.sendError(w, "Invalid TOTP code", http.StatusUnauthorized)
+			return
+		}
+	}
+
 	token, expires, err := s.createToken(user.ID)
 	if err != nil {
 		s.Logger.Printf("Token creation error: %v", err)
@@ -228,6 +252,60 @@ func (s *Server) renewToken(username, token string) (int64, error) {
 	return expires, nil
 }
 
+func (s *Server) handleTOTPVerify(w http.ResponseWriter, r *http.Request) {
+	var req TOTPVerifyRequest
+	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
+		s.sendError(w, "Invalid request format", http.StatusBadRequest)
+		return
+	}
+
+	if req.Version != "v1" || req.Username == "" || req.TOTPCode == "" {
+		s.sendError(w, "Invalid TOTP verification request", http.StatusBadRequest)
+		return
+	}
+
+	user, err := s.getUserByUsername(req.Username)
+	if err != nil {
+		if errors.Is(err, sql.ErrNoRows) {
+			s.sendError(w, "User not found", http.StatusUnauthorized)
+		} else {
+			s.Logger.Printf("Database error: %v", err)
+			s.sendError(w, "Internal server error", http.StatusInternalServerError)
+		}
+		return
+	}
+
+	// Check if TOTP is enabled for the user
+	if !user.HasTOTP() {
+		s.sendError(w, "TOTP not enabled for user", http.StatusBadRequest)
+		return
+	}
+
+	// Validate the TOTP code
+	valid, validErr := user.ValidateTOTPCode(req.TOTPCode)
+	if validErr != nil || !valid {
+		s.sendError(w, "Invalid TOTP code", http.StatusUnauthorized)
+		return
+	}
+
+	// TOTP code is valid, create a token
+	token, expires, err := s.createToken(user.ID)
+	if err != nil {
+		s.Logger.Printf("Token creation error: %v", err)
+		s.sendError(w, "Internal server error", http.StatusInternalServerError)
+		return
+	}
+
+	w.Header().Set("Content-Type", "application/json")
+	w.WriteHeader(http.StatusOK)
+	if err := json.NewEncoder(w).Encode(TokenResponse{
+		Token:   token,
+		Expires: expires,
+	}); err != nil {
+		s.Logger.Printf("Error encoding response: %v", err)
+	}
+}
+
 func (s *Server) handleDatabaseCredentials(w http.ResponseWriter, r *http.Request) {
 	// Extract authorization header
 	authHeader := r.Header.Get("Authorization")

api/server.go

@@ -32,6 +32,7 @@ func NewServer(db *sql.DB, logger *log.Logger) *Server {
 func (s *Server) registerRoutes() {
 	s.Router.HandleFunc("POST /v1/login/password", s.handlePasswordLogin)
 	s.Router.HandleFunc("POST /v1/login/token", s.handleTokenLogin)
+	s.Router.HandleFunc("POST /v1/login/totp", s.handleTOTPVerify)
 	s.Router.HandleFunc("GET /v1/database/credentials", s.handleDatabaseCredentials)
 }