Add per-user TOTP 2FA, client-level user assignment, self-service portal

- TOTP (RFC 6238, stdlib-only) enrollment in profile, login step-up, admin emergency reset. - Admins can grant a user visibility into individual clients (User.ClientIDs) in addition to whole-server access (User.ServerIDs). - New "My Access" page: non-admin users see only their assigned clients (view/QR/download only, no management), reachable from the main nav. - GetUser/GetUsers now redact TOTPSecret before returning JSON. No Go toolchain was available while writing this - not yet build-verified. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01PvrfUytqd74H6WcQkRzFM4
2026-07-25 00:42:05 +02:00
parent c29edfdcc3
commit 34bc8f76f9
12 changed files with 1213 additions and 36 deletions
@@ -24,6 +24,7 @@ import (
 	"golang.zx2c4.com/wireguard/wgctrl"
 	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"

+	"github.com/ngoduykhanh/wireguard-ui/auth"
 	"github.com/ngoduykhanh/wireguard-ui/backup"
 	"github.com/ngoduykhanh/wireguard-ui/emailer"
 	"github.com/ngoduykhanh/wireguard-ui/firewall"
@@ -135,43 +136,30 @@ func Login(db store.IStore) echo.HandlerFunc {
 		}

 		if userCorrect && passwordCorrect {
-			ageMax := 0
-			if rememberMe {
-				ageMax = 86400 * 7
+			if dbuser.TOTPEnabled {
+				// Password check passed, but this account requires a second
+				// factor. Stash a short-lived, NOT-yet-authenticated pending
+				// state instead of finalizing the session.
+				cookiePath := util.GetCookiePath()
+
+				sess, _ := session.Get("session", c)
+				sess.Options = &sessions.Options{
+					Path:     cookiePath,
+					MaxAge:   300, // 5 minutes to enter the TOTP code
+					HttpOnly: true,
+					SameSite: http.SameSiteLaxMode,
+				}
+				sess.Values["pending_totp_user"] = dbuser.Username
+				sess.Values["pending_remember_me"] = rememberMe
+				sess.Save(c.Request(), c.Response())
+
+				return c.JSON(http.StatusOK, map[string]interface{}{"success": true, "totp_required": true})
 			}

-			cookiePath := util.GetCookiePath()
-
-			sess, _ := session.Get("session", c)
-			sess.Options = &sessions.Options{
-				Path:     cookiePath,
-				MaxAge:   ageMax,
-				HttpOnly: true,
-				SameSite: http.SameSiteLaxMode,
+			if err := finalizeLoginSession(c, dbuser, rememberMe); err != nil {
+				return c.JSON(http.StatusInternalServerError, jsonHTTPResponse{false, err.Error()})
 			}

-			// set session_token
-			tokenUID := xid.New().String()
-			now := time.Now().UTC().Unix()
-			sess.Values["username"] = dbuser.Username
-			sess.Values["user_hash"] = util.GetDBUserCRC32(dbuser)
-			sess.Values["admin"] = dbuser.Admin
-			sess.Values["session_token"] = tokenUID
-			sess.Values["max_age"] = ageMax
-			sess.Values["created_at"] = now
-			sess.Values["updated_at"] = now
-			sess.Save(c.Request(), c.Response())
-
-			// set session_token in cookie
-			cookie := new(http.Cookie)
-			cookie.Name = "session_token"
-			cookie.Path = cookiePath
-			cookie.Value = tokenUID
-			cookie.MaxAge = ageMax
-			cookie.HttpOnly = true
-			cookie.SameSite = http.SameSiteLaxMode
-			c.SetCookie(cookie)
-
 			log.Infof("Logged in successfully user %s (%s)", username, c.Request().RemoteAddr)
 			return c.JSON(http.StatusOK, jsonHTTPResponse{true, "Logged in successfully"})
 		}
@@ -181,6 +169,98 @@ func Login(db store.IStore) echo.HandlerFunc {
 	}
 }

+// finalizeLoginSession fully authenticates dbuser by writing the real
+// session values and the session_token cookie. It is shared between the
+// no-TOTP path of Login and the second-factor confirmation in
+// VerifyTOTPLogin so the two code paths cannot drift apart.
+func finalizeLoginSession(c echo.Context, dbuser model.User, rememberMe bool) error {
+	ageMax := 0
+	if rememberMe {
+		ageMax = 86400 * 7
+	}
+
+	cookiePath := util.GetCookiePath()
+
+	sess, _ := session.Get("session", c)
+	sess.Options = &sessions.Options{
+		Path:     cookiePath,
+		MaxAge:   ageMax,
+		HttpOnly: true,
+		SameSite: http.SameSiteLaxMode,
+	}
+
+	// set session_token
+	tokenUID := xid.New().String()
+	now := time.Now().UTC().Unix()
+	sess.Values["username"] = dbuser.Username
+	sess.Values["user_hash"] = util.GetDBUserCRC32(dbuser)
+	sess.Values["admin"] = dbuser.Admin
+	sess.Values["session_token"] = tokenUID
+	sess.Values["max_age"] = ageMax
+	sess.Values["created_at"] = now
+	sess.Values["updated_at"] = now
+	if err := sess.Save(c.Request(), c.Response()); err != nil {
+		return err
+	}
+
+	// set session_token in cookie
+	cookie := new(http.Cookie)
+	cookie.Name = "session_token"
+	cookie.Path = cookiePath
+	cookie.Value = tokenUID
+	cookie.MaxAge = ageMax
+	cookie.HttpOnly = true
+	cookie.SameSite = http.SameSiteLaxMode
+	c.SetCookie(cookie)
+
+	return nil
+}
+
+// VerifyTOTPLogin completes a login that was put into the "pending TOTP"
+// state by Login, after the user supplies a valid 6-digit code.
+func VerifyTOTPLogin(db store.IStore) echo.HandlerFunc {
+	return func(c echo.Context) error {
+		sess, _ := session.Get("session", c)
+
+		pendingUsername, ok := sess.Values["pending_totp_user"].(string)
+		if !ok || pendingUsername == "" {
+			return c.JSON(http.StatusUnauthorized, jsonHTTPResponse{false, "No pending login"})
+		}
+
+		rememberMe, _ := sess.Values["pending_remember_me"].(bool)
+
+		data := make(map[string]interface{})
+		if err := json.NewDecoder(c.Request().Body).Decode(&data); err != nil {
+			return c.JSON(http.StatusBadRequest, jsonHTTPResponse{false, "Bad post data"})
+		}
+
+		code, _ := data["code"].(string)
+
+		dbuser, err := db.GetUserByName(pendingUsername)
+		if err != nil {
+			log.Warnf("Invalid pending TOTP login. Cannot query user %s from DB (%s)", pendingUsername, c.Request().RemoteAddr)
+			return c.JSON(http.StatusUnauthorized, jsonHTTPResponse{false, "No pending login"})
+		}
+
+		if !dbuser.TOTPEnabled || !auth.Validate(dbuser.TOTPSecret, code, time.Now()) {
+			log.Warnf("Invalid TOTP code for user %s (%s)", pendingUsername, c.Request().RemoteAddr)
+			return c.JSON(http.StatusUnauthorized, jsonHTTPResponse{false, "Invalid code"})
+		}
+
+		// clear the pending state before finalizing the real session
+		delete(sess.Values, "pending_totp_user")
+		delete(sess.Values, "pending_remember_me")
+		sess.Save(c.Request(), c.Response())
+
+		if err := finalizeLoginSession(c, dbuser, rememberMe); err != nil {
+			return c.JSON(http.StatusInternalServerError, jsonHTTPResponse{false, err.Error()})
+		}
+
+		log.Infof("Logged in successfully user %s (%s)", pendingUsername, c.Request().RemoteAddr)
+		return c.JSON(http.StatusOK, jsonHTTPResponse{true, "Logged in successfully"})
+	}
+}
+
 // GetUsers handler return a JSON list of all users
 func GetUsers(db store.IStore) echo.HandlerFunc {
 	return func(c echo.Context) error {
@@ -191,6 +271,11 @@ func GetUsers(db store.IStore) echo.HandlerFunc {
 			})
 		}

+		// never expose live TOTP secrets to API callers
+		for i := range usersList {
+			usersList[i].TOTPSecret = ""
+		}
+
 		return c.JSON(http.StatusOK, usersList)
 	}
 }
@@ -213,6 +298,10 @@ func GetUser(db store.IStore) echo.HandlerFunc {
 			return c.JSON(http.StatusNotFound, jsonHTTPResponse{false, "User not found"})
 		}

+		// Never expose the raw TOTP secret through this endpoint. Enrollment
+		// has its own dedicated, self-only route (/profile/totp/enroll).
+		userData.TOTPSecret = ""
+
 		return c.JSON(http.StatusOK, userData)
 	}
 }
@@ -309,7 +398,7 @@ func UpdateUser(db store.IStore) echo.HandlerFunc {
 			user.Admin = admin
 		}

-		// only an admin may change which servers a user can access
+		// only an admin may change which servers/clients a user can access
 		if isAdmin(c) {
 			if rawIDs, ok := data["server_ids"].([]interface{}); ok {
 				serverIDs := make([]string, 0, len(rawIDs))
@@ -320,6 +409,16 @@ func UpdateUser(db store.IStore) echo.HandlerFunc {
 				}
 				user.ServerIDs = serverIDs
 			}
+
+			if rawIDs, ok := data["client_ids"].([]interface{}); ok {
+				clientIDs := make([]string, 0, len(rawIDs))
+				for _, v := range rawIDs {
+					if id, ok := v.(string); ok && util.ValidateRecordID(id) {
+						clientIDs = append(clientIDs, id)
+					}
+				}
+				user.ClientIDs = clientIDs
+			}
 		}

 		if err := db.DeleteUser(previousUsername); err != nil {
@@ -338,6 +437,128 @@ func UpdateUser(db store.IStore) echo.HandlerFunc {
 	}
 }

+// totpIssuer is the "issuer" name embedded in the otpauth:// provisioning
+// URI, shown by authenticator apps next to the account name.
+const totpIssuer = "WireGuard-UI-Multi"
+
+// EnrollTOTP starts (or resumes) self-service TOTP enrollment for the
+// currently logged-in user. It generates a secret on first call (leaving
+// TOTPEnabled false until confirmed via ConfirmTOTP) and always returns a
+// QR code / provisioning secret for the currently stored (unconfirmed or
+// confirmed) secret.
+func EnrollTOTP(db store.IStore) echo.HandlerFunc {
+	return func(c echo.Context) error {
+		username := currentUser(c)
+
+		user, err := db.GetUserByName(username)
+		if err != nil {
+			return c.JSON(http.StatusNotFound, jsonHTTPResponse{false, "User not found"})
+		}
+
+		if user.TOTPSecret == "" {
+			secret, err := auth.GenerateSecret()
+			if err != nil {
+				return c.JSON(http.StatusInternalServerError, jsonHTTPResponse{false, "Cannot generate TOTP secret: " + err.Error()})
+			}
+			user.TOTPSecret = secret
+			user.TOTPEnabled = false
+			if err := db.SaveUser(user); err != nil {
+				return c.JSON(http.StatusInternalServerError, jsonHTTPResponse{false, err.Error()})
+			}
+		}
+
+		uri := auth.ProvisioningURI(user.TOTPSecret, user.Username, totpIssuer)
+
+		png, err := qrcode.Encode(uri, qrcode.Medium, 256)
+		if err != nil {
+			return c.JSON(http.StatusInternalServerError, jsonHTTPResponse{false, "qr gen: " + err.Error()})
+		}
+		qrDataURI := "data:image/png;base64," + base64.StdEncoding.EncodeToString(png)
+
+		return c.JSON(http.StatusOK, map[string]interface{}{
+			"secret":  user.TOTPSecret,
+			"qrcode":  qrDataURI,
+			"enabled": user.TOTPEnabled,
+		})
+	}
+}
+
+// ConfirmTOTP verifies a code against the currently pending (unconfirmed)
+// TOTP secret for the logged-in user, and enables TOTP on success.
+func ConfirmTOTP(db store.IStore) echo.HandlerFunc {
+	return func(c echo.Context) error {
+		username := currentUser(c)
+
+		data := make(map[string]interface{})
+		if err := json.NewDecoder(c.Request().Body).Decode(&data); err != nil {
+			return c.JSON(http.StatusBadRequest, jsonHTTPResponse{false, "Bad post data"})
+		}
+		code, _ := data["code"].(string)
+
+		user, err := db.GetUserByName(username)
+		if err != nil {
+			return c.JSON(http.StatusNotFound, jsonHTTPResponse{false, "User not found"})
+		}
+
+		if user.TOTPSecret == "" {
+			return c.JSON(http.StatusBadRequest, jsonHTTPResponse{false, "TOTP enrollment has not been started"})
+		}
+
+		if !auth.Validate(user.TOTPSecret, code, time.Now()) {
+			log.Warnf("Invalid TOTP confirmation code for user %s (%s)", username, c.Request().RemoteAddr)
+			return c.JSON(http.StatusUnauthorized, jsonHTTPResponse{false, "Invalid code"})
+		}
+
+		user.TOTPEnabled = true
+		if err := db.SaveUser(user); err != nil {
+			return c.JSON(http.StatusInternalServerError, jsonHTTPResponse{false, err.Error()})
+		}
+
+		log.Infof("User %s enabled two-factor authentication", username)
+		return c.JSON(http.StatusOK, jsonHTTPResponse{true, "Two-factor authentication enabled"})
+	}
+}
+
+// DisableTOTP disables TOTP for the current user, or - if the caller is an
+// admin and explicitly names a different user - resets that other user's
+// TOTP as an emergency lockout recovery mechanism.
+func DisableTOTP(db store.IStore) echo.HandlerFunc {
+	return func(c echo.Context) error {
+		username := currentUser(c)
+
+		data := make(map[string]interface{})
+		// Body is optional for the self-service case.
+		_ = json.NewDecoder(c.Request().Body).Decode(&data)
+
+		targetUsername := username
+		if requested, ok := data["username"].(string); ok && requested != "" {
+			if isAdmin(c) && requested != currentUser(c) {
+				targetUsername = requested
+			} else if requested != currentUser(c) {
+				return c.JSON(http.StatusForbidden, jsonHTTPResponse{false, "Manager cannot access other user data"})
+			}
+		}
+
+		if !usernameRegexp.MatchString(targetUsername) {
+			return c.JSON(http.StatusBadRequest, jsonHTTPResponse{false, "Please provide a valid username"})
+		}
+
+		user, err := db.GetUserByName(targetUsername)
+		if err != nil {
+			return c.JSON(http.StatusNotFound, jsonHTTPResponse{false, "User not found"})
+		}
+
+		user.TOTPSecret = ""
+		user.TOTPEnabled = false
+		if err := db.SaveUser(user); err != nil {
+			return c.JSON(http.StatusInternalServerError, jsonHTTPResponse{false, err.Error()})
+		}
+
+		log.Infof("Two-factor authentication disabled for user %s (by %s)", targetUsername, username)
+		return c.JSON(http.StatusOK, jsonHTTPResponse{true, "Two-factor authentication disabled"})
+	}
+}
+
 // CreateUser to create new user
 func CreateUser(db store.IStore) echo.HandlerFunc {
 	return func(c echo.Context) error {
@@ -382,6 +603,14 @@ func CreateUser(db store.IStore) echo.HandlerFunc {
 			}
 		}

+		if rawIDs, ok := data["client_ids"].([]interface{}); ok {
+			for _, v := range rawIDs {
+				if id, ok := v.(string); ok && util.ValidateRecordID(id) {
+					user.ClientIDs = append(user.ClientIDs, id)
+				}
+			}
+		}
+
 		if err := db.SaveUser(user); err != nil {
 			return c.JSON(http.StatusInternalServerError, jsonHTTPResponse{false, err.Error()})
 		}