reliability fixes and disconnect-all command

management/cmd/admin.go

@@ -3,8 +3,10 @@ package cmd
 import (
 	"context"
 	"fmt"
+	"path"
 	"path/filepath"
 
+	"github.com/dexidp/dex/storage"
 	log "github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
 
@@ -12,6 +14,8 @@ import (
 	admincmd "github.com/netbirdio/netbird/management/cmd/admin"
 	tokencmd "github.com/netbirdio/netbird/management/cmd/token"
 	nbconfig "github.com/netbirdio/netbird/management/internals/server/config"
+	"github.com/netbirdio/netbird/management/server/activity"
+	activitystore "github.com/netbirdio/netbird/management/server/activity/store"
 	"github.com/netbirdio/netbird/management/server/store"
 	"github.com/netbirdio/netbird/util"
 )
@@ -20,70 +24,154 @@ var adminDatadir string
 
 // newAdminCommands creates the admin command tree with management-specific resource openers.
 func newAdminCommands() *cobra.Command {
-	cmd := admincmd.NewCommands(withAdminResources)
+	cmd := admincmd.NewCommands(admincmd.Openers{
+		Resources: withAdminResources,
+		Store:     withAdminStoreOnly,
+		IDP:       withAdminIDPOnly,
+	})
 	cmd.PersistentFlags().StringVar(&adminDatadir, "datadir", "", "Override the data directory from config (used for store.db and the default idp.db)")
-	cmd.AddCommand(tokencmd.NewCommands(withAdminTokenStore))
+	return cmd
+}
+
+func newLegacyTokenCommand() *cobra.Command {
+	cmd := tokencmd.NewCommands(tokencmd.StoreOpener(withAdminStoreOnly))
+	cmd.Deprecated = "use 'admin token' instead"
+	cmd.PersistentFlags().StringVar(&nbconfig.MgmtConfigPath, "config", defaultMgmtConfig, "Netbird config file location")
 	return cmd
 }
 
 // withAdminResources initializes logging, loads config, opens the management store
 // and embedded IdP storage, and calls fn.
 func withAdminResources(cmd *cobra.Command, fn func(ctx context.Context, resources admincmd.Resources) error) error {
-	return withAdminStore(cmd, func(ctx context.Context, managementStore store.Store, config *nbconfig.Config) error {
-		idpStorage, err := admincmd.OpenEmbeddedIDPStorage(config.EmbeddedIdP)
+	return withAdminConfig(cmd, true, func(ctx context.Context, config *nbconfig.Config, datadir string) error {
+		managementStore, err := openAdminStore(ctx, config, datadir)
 		if err != nil {
 			return err
 		}
-		defer func() {
-			if err := idpStorage.Close(); err != nil {
-				log.Debugf("close embedded IdP storage: %v", err)
-			}
-		}()
+		defer admincmd.CloseStore(ctx, managementStore)
 
-		return fn(ctx, admincmd.Resources{Store: managementStore, IDPStorage: idpStorage})
+		idpStorage, idpStorageFile, err := admincmd.OpenIDPStorage(config)
+		if err != nil {
+			return err
+		}
+		defer admincmd.CloseIDPStorage(idpStorage)
+
+		eventStore, esErr := openAdminEventStore(ctx, config, datadir)
+		if esErr != nil {
+			_, _ = fmt.Fprintf(cmd.ErrOrStderr(), "Warning: audit events will not be recorded: %v\n", esErr)
+		}
+		if eventStore != nil {
+			defer func() {
+				if err := eventStore.Close(ctx); err != nil {
+					log.Debugf("close activity event store: %v", err)
+				}
+			}()
+		}
+
+		return fn(ctx, admincmd.Resources{Store: managementStore, IDPStorage: idpStorage, IDPStorageFile: idpStorageFile, EventStore: eventStore})
 	})
 }
 
-// withAdminTokenStore opens only the management store for admin token commands.
-func withAdminTokenStore(cmd *cobra.Command, fn func(ctx context.Context, s store.Store) error) error {
-	return withAdminStore(cmd, func(ctx context.Context, managementStore store.Store, _ *nbconfig.Config) error {
+// withAdminStoreOnly opens only the management store for admin subcommands that do not
+// need embedded IdP storage.
+func withAdminStoreOnly(cmd *cobra.Command, fn func(ctx context.Context, s store.Store) error) error {
+	return withAdminConfig(cmd, false, func(ctx context.Context, config *nbconfig.Config, datadir string) error {
+		managementStore, err := openAdminStore(ctx, config, datadir)
+		if err != nil {
+			return err
+		}
+		defer admincmd.CloseStore(ctx, managementStore)
+
 		return fn(ctx, managementStore)
 	})
 }
 
-func withAdminStore(cmd *cobra.Command, fn func(ctx context.Context, s store.Store, config *nbconfig.Config) error) error {
+func withAdminIDPOnly(cmd *cobra.Command, fn func(ctx context.Context, idpStorage storage.Storage, storageFile string) error) error {
+	return withAdminConfig(cmd, true, func(ctx context.Context, config *nbconfig.Config, _ string) error {
+		idpStorage, idpStorageFile, err := admincmd.OpenIDPStorage(config)
+		if err != nil {
+			return err
+		}
+		defer admincmd.CloseIDPStorage(idpStorage)
+
+		return fn(ctx, idpStorage, idpStorageFile)
+	})
+}
+
+func withAdminConfig(cmd *cobra.Command, applyIDPDefaults bool, fn func(ctx context.Context, config *nbconfig.Config, datadir string) error) error {
 	if err := util.InitLog("error", "console"); err != nil {
 		return fmt.Errorf("init log: %w", err)
 	}
 
 	ctx := context.WithValue(cmd.Context(), hook.ExecutionContextKey, hook.SystemSource) //nolint:staticcheck
 
-	config, err := LoadMgmtConfig(ctx, nbconfig.MgmtConfigPath)
+	config, datadir, err := loadAdminMgmtConfig(ctx, applyIDPDefaults)
 	if err != nil {
 		return fmt.Errorf("load config: %w", err)
 	}
 
-	datadir := config.Datadir
-	if adminDatadir != "" {
-		oldDatadir := datadir
-		datadir = adminDatadir
-		if config.EmbeddedIdP != nil && config.EmbeddedIdP.Storage.Type == "sqlite3" {
-			defaultIDPFile := filepath.Join(oldDatadir, "idp.db")
-			if config.EmbeddedIdP.Storage.Config.File == "" || config.EmbeddedIdP.Storage.Config.File == defaultIDPFile {
-				config.EmbeddedIdP.Storage.Config.File = filepath.Join(datadir, "idp.db")
-			}
+	return fn(ctx, config, datadir)
+}
+
+func loadAdminMgmtConfig(ctx context.Context, applyIDPDefaults bool) (*nbconfig.Config, string, error) {
+	config := &nbconfig.Config{}
+	if _, err := util.ReadJsonWithEnvSub(nbconfig.MgmtConfigPath, config); err != nil {
+		return nil, "", err
+	}
+
+	if applyIDPDefaults {
+		if err := ApplyEmbeddedIdPConfig(ctx, config); err != nil {
+			return nil, "", err
 		}
 	}
 
+	datadir := config.Datadir
+	applyAdminDatadirOverride(config, &datadir)
+	return config, datadir, nil
+}
+
+func applyAdminDatadirOverride(config *nbconfig.Config, datadir *string) {
+	if adminDatadir == "" {
+		return
+	}
+
+	oldDatadir := *datadir
+	*datadir = adminDatadir
+	if config.EmbeddedIdP != nil && config.EmbeddedIdP.Storage.Type == "sqlite3" && isDefaultIDPStorageFile(config.EmbeddedIdP.Storage.Config.File, oldDatadir) {
+		config.EmbeddedIdP.Storage.Config.File = filepath.Join(*datadir, "idp.db")
+	}
+}
+
+func isDefaultIDPStorageFile(file, datadir string) bool {
+	if file == "" {
+		return true
+	}
+	defaultFile := filepath.Join(datadir, "idp.db")
+	legacyDefaultFile := path.Join(datadir, "idp.db")
+	legacySlashDefaultFile := path.Join(filepath.ToSlash(datadir), "idp.db")
+	return filepath.Clean(file) == filepath.Clean(defaultFile) ||
+		file == legacyDefaultFile ||
+		filepath.ToSlash(file) == legacySlashDefaultFile
+}
+
+func openAdminStore(ctx context.Context, config *nbconfig.Config, datadir string) (store.Store, error) {
 	managementStore, err := store.NewStore(ctx, config.StoreConfig.Engine, datadir, nil, true)
 	if err != nil {
-		return fmt.Errorf("create store: %w", err)
+		return nil, fmt.Errorf("create store: %w", err)
 	}
-	defer func() {
-		if err := managementStore.Close(ctx); err != nil {
-			log.Debugf("close store: %v", err)
-		}
-	}()
+	return managementStore, nil
+}
 
-	return fn(ctx, managementStore, config)
+func openAdminEventStore(ctx context.Context, config *nbconfig.Config, datadir string) (activity.Store, error) {
+	if config.DataStoreEncryptionKey == "" {
+		return nil, fmt.Errorf("data store encryption key is not configured")
+	}
+	eventStore, err := activitystore.NewSqlStore(ctx, datadir, config.DataStoreEncryptionKey)
+	if err != nil {
+		return nil, fmt.Errorf("open activity event store: %w", err)
+	}
+	if eventStore == nil {
+		return nil, fmt.Errorf("open activity event store: returned nil store")
+	}
+	return eventStore, nil
 }

management/cmd/admin/admin.go

@@ -11,34 +11,49 @@ import (
 	"log/slog"
 	"os"
 	"strings"
+	"time"
 
 	"github.com/dexidp/dex/storage"
+	log "github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
 	"golang.org/x/crypto/bcrypt"
 
+	"github.com/netbirdio/netbird/formatter/hook"
 	nbdex "github.com/netbirdio/netbird/idp/dex"
+	"github.com/netbirdio/netbird/management/cmd/proxy"
+	"github.com/netbirdio/netbird/management/cmd/token"
+	nbconfig "github.com/netbirdio/netbird/management/internals/server/config"
 	"github.com/netbirdio/netbird/management/server"
+	"github.com/netbirdio/netbird/management/server/activity"
 	"github.com/netbirdio/netbird/management/server/idp"
 	"github.com/netbirdio/netbird/management/server/store"
 	"github.com/netbirdio/netbird/management/server/types"
 )
 
-const (
-	localConnectorID           = "local"
-	dashboardClientID          = "netbird-dashboard"
-	cliClientID                = "netbird-cli"
-	defaultTOTPAuthenticatorID = "default-totp"
-)
-
 // Resources contains the storages required by the admin commands.
 type Resources struct {
-	Store      store.Store
-	IDPStorage storage.Storage
+	Store          store.Store
+	IDPStorage     storage.Storage
+	IDPStorageFile string
+	EventStore     activity.Store
 }
 
 // Opener initializes command resources from the command context and calls fn.
 type Opener func(cmd *cobra.Command, fn func(ctx context.Context, resources Resources) error) error
 
+// StoreOpener initializes only the management store from the command context and calls fn.
+type StoreOpener func(cmd *cobra.Command, fn func(ctx context.Context, s store.Store) error) error
+
+// IDPOpener initializes only the embedded IdP storage from the command context and calls fn.
+type IDPOpener func(cmd *cobra.Command, fn func(ctx context.Context, idpStorage storage.Storage, storageFile string) error) error
+
+// Openers contains the resource openers needed by the admin command tree.
+type Openers struct {
+	Resources Opener
+	Store     StoreOpener
+	IDP       IDPOpener
+}
+
 type userSelector struct {
 	email  string
 	userID string
@@ -59,8 +74,8 @@ func (s userSelector) validate() error {
 	return nil
 }
 
-// NewCommands creates the admin command tree with the given resource opener.
-func NewCommands(opener Opener) *cobra.Command {
+// NewCommands creates the admin command tree with the given resource openers.
+func NewCommands(openers Openers) *cobra.Command {
 	adminCmd := &cobra.Command{
 		Use:   "admin",
 		Short: "Self-hosted administrator helpers",
@@ -81,12 +96,15 @@ func NewCommands(opener Opener) *cobra.Command {
 		Short:   "Change a local user's password",
 		Args:    cobra.NoArgs,
 		RunE: func(cmd *cobra.Command, _ []string) error {
+			if err := passwordSelector.validate(); err != nil {
+				return err
+			}
 			newPassword, err := resolvePasswordInput(cmd, password, passwordFile)
 			if err != nil {
 				return err
 			}
-			return opener(cmd, func(ctx context.Context, resources Resources) error {
-				return runChangePassword(ctx, resources.IDPStorage, cmd.OutOrStdout(), passwordSelector, newPassword)
+			return openers.IDP(cmd, func(ctx context.Context, idpStorage storage.Storage, storageFile string) error {
+				return runChangePassword(ctx, idpStorage, cmd.OutOrStdout(), passwordSelector, newPassword, storageFile)
 			})
 		},
 	}
@@ -100,8 +118,11 @@ func NewCommands(opener Opener) *cobra.Command {
 		Short: "Reset a local user's MFA enrollment",
 		Args:  cobra.NoArgs,
 		RunE: func(cmd *cobra.Command, _ []string) error {
-			return opener(cmd, func(ctx context.Context, resources Resources) error {
-				return runResetMFA(ctx, resources.IDPStorage, cmd.OutOrStdout(), resetSelector)
+			if err := resetSelector.validate(); err != nil {
+				return err
+			}
+			return openers.IDP(cmd, func(ctx context.Context, idpStorage storage.Storage, storageFile string) error {
+				return runResetMFA(ctx, idpStorage, cmd.OutOrStdout(), resetSelector, storageFile)
 			})
 		},
 	}
@@ -119,7 +140,7 @@ func NewCommands(opener Opener) *cobra.Command {
 		Short: "Enable MFA for local embedded IdP users",
 		Args:  cobra.NoArgs,
 		RunE: func(cmd *cobra.Command, _ []string) error {
-			return opener(cmd, func(ctx context.Context, resources Resources) error {
+			return openers.Resources(cmd, func(ctx context.Context, resources Resources) error {
 				return runSetMFAEnabled(ctx, resources, cmd.OutOrStdout(), true)
 			})
 		},
@@ -130,7 +151,7 @@ func NewCommands(opener Opener) *cobra.Command {
 		Short: "Disable MFA for local embedded IdP users",
 		Args:  cobra.NoArgs,
 		RunE: func(cmd *cobra.Command, _ []string) error {
-			return opener(cmd, func(ctx context.Context, resources Resources) error {
+			return openers.Resources(cmd, func(ctx context.Context, resources Resources) error {
 				return runSetMFAEnabled(ctx, resources, cmd.OutOrStdout(), false)
 			})
 		},
@@ -141,7 +162,7 @@ func NewCommands(opener Opener) *cobra.Command {
 		Short: "Show local MFA status",
 		Args:  cobra.NoArgs,
 		RunE: func(cmd *cobra.Command, _ []string) error {
-			return opener(cmd, func(ctx context.Context, resources Resources) error {
+			return openers.Resources(cmd, func(ctx context.Context, resources Resources) error {
 				return runMFAStatus(ctx, resources, cmd.OutOrStdout())
 			})
 		},
@@ -149,6 +170,10 @@ func NewCommands(opener Opener) *cobra.Command {
 
 	mfaCmd.AddCommand(enableCmd, disableCmd, statusCmd)
 	adminCmd.AddCommand(userCmd, mfaCmd)
+	if openers.Store != nil {
+		adminCmd.AddCommand(tokencmd.NewCommands(tokencmd.StoreOpener(openers.Store)))
+		adminCmd.AddCommand(proxycmd.NewCommands(proxycmd.StoreOpener(openers.Store)))
+	}
 	return adminCmd
 }
 
@@ -171,6 +196,45 @@ func OpenEmbeddedIDPStorage(cfg *idp.EmbeddedIdPConfig) (storage.Storage, error)
 	return st, nil
 }
 
+// CloseStore closes the management store and logs cleanup errors at debug level.
+func CloseStore(ctx context.Context, s store.Store) {
+	if s == nil {
+		return
+	}
+	if err := s.Close(ctx); err != nil {
+		log.Debugf("close store: %v", err)
+	}
+}
+
+// OpenIDPStorage opens embedded IdP storage and returns its sqlite file path when applicable.
+func OpenIDPStorage(config *nbconfig.Config) (storage.Storage, string, error) {
+	if config == nil {
+		return nil, "", fmt.Errorf("management config is required")
+	}
+	idpStorage, err := OpenEmbeddedIDPStorage(config.EmbeddedIdP)
+	if err != nil {
+		return nil, "", err
+	}
+	return idpStorage, embeddedIDPStorageFile(config), nil
+}
+
+func embeddedIDPStorageFile(config *nbconfig.Config) string {
+	if config.EmbeddedIdP == nil || config.EmbeddedIdP.Storage.Type != "sqlite3" {
+		return ""
+	}
+	return config.EmbeddedIdP.Storage.Config.File
+}
+
+// CloseIDPStorage closes embedded IdP storage and logs cleanup errors at debug level.
+func CloseIDPStorage(s storage.Storage) {
+	if s == nil {
+		return
+	}
+	if err := s.Close(); err != nil {
+		log.Debugf("close embedded IdP storage: %v", err)
+	}
+}
+
 func addUserSelectorFlags(cmd *cobra.Command, selector *userSelector) {
 	cmd.Flags().StringVar(&selector.email, "email", "", "User email")
 	cmd.Flags().StringVar(&selector.userID, "user-id", "", "User ID")
@@ -197,7 +261,7 @@ func resolvePasswordInput(cmd *cobra.Command, password, passwordFile string) (st
 	return strings.TrimRight(string(data), "\r\n"), nil
 }
 
-func runChangePassword(ctx context.Context, idpStorage storage.Storage, w io.Writer, selector userSelector, password string) error {
+func runChangePassword(ctx context.Context, idpStorage storage.Storage, w io.Writer, selector userSelector, password string, idpStorageFile string) error {
 	if idpStorage == nil {
 		return fmt.Errorf("embedded IdP storage is required")
 	}
@@ -212,7 +276,7 @@ func runChangePassword(ctx context.Context, idpStorage storage.Storage, w io.Wri
 		return fmt.Errorf("invalid password: %w", err)
 	}
 
-	user, err := findLocalUser(ctx, idpStorage, selector)
+	user, err := findLocalUser(ctx, idpStorage, selector, idpStorageFile)
 	if err != nil {
 		return err
 	}
@@ -237,7 +301,7 @@ func runChangePassword(ctx context.Context, idpStorage storage.Storage, w io.Wri
 	return nil
 }
 
-func runResetMFA(ctx context.Context, idpStorage storage.Storage, w io.Writer, selector userSelector) error {
+func runResetMFA(ctx context.Context, idpStorage storage.Storage, w io.Writer, selector userSelector, idpStorageFile string) error {
 	if idpStorage == nil {
 		return fmt.Errorf("embedded IdP storage is required")
 	}
@@ -246,13 +310,13 @@ func runResetMFA(ctx context.Context, idpStorage storage.Storage, w io.Writer, s
 		return err
 	}
 
-	user, err := findLocalUser(ctx, idpStorage, selector)
+	user, err := findLocalUser(ctx, idpStorage, selector, idpStorageFile)
 	if err != nil {
 		return err
 	}
 
 	reset := false
-	err = idpStorage.UpdateUserIdentity(ctx, user.UserID, localConnectorID, func(old storage.UserIdentity) (storage.UserIdentity, error) {
+	err = idpStorage.UpdateUserIdentity(ctx, user.UserID, idp.LocalConnectorID, func(old storage.UserIdentity) (storage.UserIdentity, error) {
 		reset = reset || len(old.MFASecrets) > 0 || len(old.WebAuthnCredentials) > 0
 		old.MFASecrets = map[string]*storage.MFASecret{}
 		old.WebAuthnCredentials = map[string][]storage.WebAuthnCredential{}
@@ -289,24 +353,30 @@ func runSetMFAEnabled(ctx context.Context, resources Resources, w io.Writer, ena
 		return fmt.Errorf("embedded IdP storage is required")
 	}
 
-	accounts := resources.Store.GetAllAccounts(ctx)
-	if len(accounts) != 1 {
-		return fmt.Errorf("expected exactly one account, got %d; local MFA is supported only in single-account embedded IdP deployments", len(accounts))
+	accountID, settings, err := getSingleAccountSettings(ctx, resources.Store)
+	if err != nil {
+		return err
 	}
 
-	settings := &types.Settings{}
-	if accounts[0].Settings != nil {
-		settings = accounts[0].Settings.Copy()
-	}
-	settings.LocalMfaEnabled = enabled
-	if err := resources.Store.SaveAccountSettings(ctx, accounts[0].Id, settings); err != nil {
-		return fmt.Errorf("save local MFA account setting: %w", err)
-	}
+	oldEnabled := settings.LocalMfaEnabled
+	newSettings := settings.Copy()
+	newSettings.LocalMfaEnabled = enabled
 
 	if err := setIDPClientsMFA(ctx, resources.IDPStorage, enabled); err != nil {
 		return err
 	}
 
+	if err := resources.Store.SaveAccountSettings(ctx, accountID, newSettings); err != nil {
+		if rollbackErr := setIDPClientsMFA(ctx, resources.IDPStorage, oldEnabled); rollbackErr != nil {
+			return fmt.Errorf("save local MFA account setting: %w (also failed to roll back embedded IdP MFA state: %v)", err, rollbackErr)
+		}
+		return fmt.Errorf("save local MFA account setting: %w", err)
+	}
+
+	if err := storeMFAActivity(ctx, resources.EventStore, accountID, enabled); err != nil {
+		_, _ = fmt.Fprintf(w, "Warning: failed to record audit event: %v\n", err)
+	}
+
 	state := "disabled"
 	if enabled {
 		state = "enabled"
@@ -323,13 +393,13 @@ func runMFAStatus(ctx context.Context, resources Resources, w io.Writer) error {
 		return fmt.Errorf("embedded IdP storage is required")
 	}
 
-	accounts := resources.Store.GetAllAccounts(ctx)
-	accountStatus := "unknown"
-	if len(accounts) == 1 && accounts[0].Settings != nil {
-		accountStatus = "disabled"
-		if accounts[0].Settings.LocalMfaEnabled {
-			accountStatus = "enabled"
-		}
+	_, settings, err := getSingleAccountSettings(ctx, resources.Store)
+	if err != nil {
+		return err
+	}
+	accountStatus := "disabled"
+	if settings.LocalMfaEnabled {
+		accountStatus = "enabled"
 	}
 
 	clientStatus, err := idpClientsMFAStatus(ctx, resources.IDPStorage)
@@ -342,7 +412,52 @@ func runMFAStatus(ctx context.Context, resources Resources, w io.Writer) error {
 	return nil
 }
 
-func findLocalUser(ctx context.Context, idpStorage storage.Storage, selector userSelector) (storage.Password, error) {
+func getSingleAccountSettings(ctx context.Context, s store.Store) (string, *types.Settings, error) {
+	count, err := s.GetAccountsCounter(ctx)
+	if err != nil {
+		return "", nil, fmt.Errorf("count accounts: %w", err)
+	}
+	if count != 1 {
+		return "", nil, fmt.Errorf("expected exactly one account, got %d; local MFA is supported only in single-account embedded IdP deployments", count)
+	}
+
+	accountID, err := s.GetAnyAccountID(ctx)
+	if err != nil {
+		return "", nil, fmt.Errorf("get account ID: %w", err)
+	}
+
+	settings, err := s.GetAccountSettings(ctx, store.LockingStrengthNone, accountID)
+	if err != nil {
+		return "", nil, fmt.Errorf("get account settings: %w", err)
+	}
+	if settings == nil {
+		settings = &types.Settings{}
+	}
+	return accountID, settings, nil
+}
+
+func storeMFAActivity(ctx context.Context, eventStore activity.Store, accountID string, enabled bool) error {
+	if eventStore == nil {
+		return nil
+	}
+	event := activity.AccountLocalMfaDisabled
+	if enabled {
+		event = activity.AccountLocalMfaEnabled
+	}
+	_, err := eventStore.Save(ctx, &activity.Event{
+		Timestamp:   time.Now().UTC(),
+		Activity:    event,
+		InitiatorID: string(hook.SystemSource),
+		TargetID:    accountID,
+		AccountID:   accountID,
+	})
+	if err != nil {
+		return fmt.Errorf("save local MFA audit event: %w", err)
+	}
+	return nil
+}
+
+func findLocalUser(ctx context.Context, idpStorage storage.Storage, selector userSelector, idpStorageFile string) (storage.Password, error) {
 	selector = selector.normalized()
 	if err := selector.validate(); err != nil {
 		return storage.Password{}, err
@@ -351,6 +466,11 @@ func findLocalUser(ctx context.Context, idpStorage storage.Storage, selector use
 	if selector.email != "" {
 		user, err := idpStorage.GetPassword(ctx, selector.email)
 		if errors.Is(err, storage.ErrNotFound) {
+			if empty, listErr := localUsersEmpty(ctx, idpStorage); listErr != nil {
+				return storage.Password{}, listErr
+			} else if empty {
+				return storage.Password{}, noLocalUsersError(idpStorageFile)
+			}
 			return storage.Password{}, fmt.Errorf("local user with email %q not found", selector.email)
 		}
 		if err != nil {
@@ -374,11 +494,31 @@ func findLocalUser(ctx context.Context, idpStorage storage.Storage, selector use
 		}
 	}
 
+	if len(users) == 0 {
+		return storage.Password{}, noLocalUsersError(idpStorageFile)
+	}
+
 	return storage.Password{}, fmt.Errorf("local user with ID %q not found", selector.userID)
 }
 
+func localUsersEmpty(ctx context.Context, idpStorage storage.Storage) (bool, error) {
+	users, err := idpStorage.ListPasswords(ctx)
+	if err != nil {
+		return false, fmt.Errorf("list local users: %w", err)
+	}
+	return len(users) == 0, nil
+}
+
+func noLocalUsersError(idpStorageFile string) error {
+	location := ""
+	if idpStorageFile != "" {
+		location = fmt.Sprintf(" (%s)", idpStorageFile)
+	}
+	return fmt.Errorf("no local users exist in the embedded IdP storage%s; the management server may never have started with this config, or --datadir points at the wrong location", location)
+}
+
 func deleteLocalAuthSession(ctx context.Context, idpStorage storage.Storage, userID string) error {
-	err := idpStorage.DeleteAuthSession(ctx, userID, localConnectorID)
+	err := idpStorage.DeleteAuthSession(ctx, userID, idp.LocalConnectorID)
 	if err == nil || errors.Is(err, storage.ErrNotFound) {
 		return nil
 	}
@@ -388,25 +528,21 @@ func deleteLocalAuthSession(ctx context.Context, idpStorage storage.Storage, use
 func setIDPClientsMFA(ctx context.Context, idpStorage storage.Storage, enabled bool) error {
 	var mfaChain []string
 	if enabled {
-		mfaChain = []string{defaultTOTPAuthenticatorID}
+		mfaChain = []string{idp.DefaultTOTPAuthenticatorID}
 	}
 
-	for _, clientID := range []string{cliClientID, dashboardClientID} {
-		if err := idpStorage.UpdateClient(ctx, clientID, func(old storage.Client) (storage.Client, error) {
-			old.MFAChain = mfaChain
-			return old, nil
-		}); err != nil {
-			if errors.Is(err, storage.ErrNotFound) {
-				return fmt.Errorf("embedded IdP client %q not found; start the management server once before toggling MFA", clientID)
-			}
-			return fmt.Errorf("update MFA chain on embedded IdP client %q: %w", clientID, err)
+	clientIDs := []string{idp.StaticClientCLI, idp.StaticClientDashboard}
+	if err := nbdex.SetClientsMFAChain(ctx, idpStorage, clientIDs, mfaChain); err != nil {
+		if errors.Is(err, storage.ErrNotFound) {
+			return fmt.Errorf("embedded IdP client not found; start the management server once before toggling MFA: %w", err)
 		}
+		return fmt.Errorf("update MFA chain on embedded IdP clients: %w", err)
 	}
 	return nil
 }
 
 func idpClientsMFAStatus(ctx context.Context, idpStorage storage.Storage) (string, error) {
-	clientIDs := []string{cliClientID, dashboardClientID}
+	clientIDs := []string{idp.StaticClientCLI, idp.StaticClientDashboard}
 	enabledCount := 0
 	for _, clientID := range clientIDs {
 		client, err := idpStorage.GetClient(ctx, clientID)
@@ -416,7 +552,7 @@ func idpClientsMFAStatus(ctx context.Context, idpStorage storage.Storage) (strin
 		if err != nil {
 			return "unknown", fmt.Errorf("get embedded IdP client %q: %w", clientID, err)
 		}
-		if hasAuthenticator(client.MFAChain, defaultTOTPAuthenticatorID) {
+		if hasAuthenticator(client.MFAChain, idp.DefaultTOTPAuthenticatorID) {
 			enabledCount++
 		}
 	}

management/cmd/admin/admin_test.go

@@ -16,6 +16,9 @@ import (
 	"golang.org/x/crypto/bcrypt"
 
 	nbdex "github.com/netbirdio/netbird/idp/dex"
+	"github.com/netbirdio/netbird/management/server/idp"
+	mgmtstore "github.com/netbirdio/netbird/management/server/store"
+	"github.com/netbirdio/netbird/management/server/types"
 )
 
 func newTestIDPStorage(t *testing.T) storage.Storage {
@@ -33,10 +36,10 @@ func newTestIDPStorage(t *testing.T) storage.Storage {
 	}))
 	require.NoError(t, st.CreateUserIdentity(context.Background(), storage.UserIdentity{
 		UserID:      "user-1",
-		ConnectorID: localConnectorID,
+		ConnectorID: idp.LocalConnectorID,
 		MFASecrets: map[string]*storage.MFASecret{
-			defaultTOTPAuthenticatorID: {
-				AuthenticatorID: defaultTOTPAuthenticatorID,
+			idp.DefaultTOTPAuthenticatorID: {
+				AuthenticatorID: idp.DefaultTOTPAuthenticatorID,
 				Type:            "TOTP",
 				Secret:          "otpauth://totp/NetBird:user@example.com?secret=ABC",
 				Confirmed:       true,
@@ -49,11 +52,11 @@ func newTestIDPStorage(t *testing.T) storage.Storage {
 	}))
 	require.NoError(t, st.CreateAuthSession(context.Background(), storage.AuthSession{
 		UserID:      "user-1",
-		ConnectorID: localConnectorID,
+		ConnectorID: idp.LocalConnectorID,
 		Nonce:       "nonce",
 	}))
-	require.NoError(t, st.CreateClient(context.Background(), storage.Client{ID: cliClientID, Name: "CLI"}))
-	require.NoError(t, st.CreateClient(context.Background(), storage.Client{ID: dashboardClientID, Name: "Dashboard"}))
+	require.NoError(t, st.CreateClient(context.Background(), storage.Client{ID: idp.StaticClientCLI, Name: "CLI"}))
+	require.NoError(t, st.CreateClient(context.Background(), storage.Client{ID: idp.StaticClientDashboard, Name: "Dashboard"}))
 
 	return st
 }
@@ -63,7 +66,7 @@ func TestRunChangePassword(t *testing.T) {
 	st := newTestIDPStorage(t)
 	var out bytes.Buffer
 
-	err := runChangePassword(ctx, st, &out, userSelector{email: "user@example.com"}, "NewPass1!")
+	err := runChangePassword(ctx, st, &out, userSelector{email: "user@example.com"}, "NewPass1!", "")
 	require.NoError(t, err)
 	require.Contains(t, out.String(), "Password updated")
 
@@ -71,13 +74,13 @@ func TestRunChangePassword(t *testing.T) {
 	require.NoError(t, err)
 	require.NoError(t, bcrypt.CompareHashAndPassword(user.Hash, []byte("NewPass1!")))
 
-	_, err = st.GetAuthSession(ctx, "user-1", localConnectorID)
+	_, err = st.GetAuthSession(ctx, "user-1", idp.LocalConnectorID)
 	require.ErrorIs(t, err, storage.ErrNotFound)
 }
 
 func TestRunChangePasswordValidatesPassword(t *testing.T) {
 	st := newTestIDPStorage(t)
-	err := runChangePassword(context.Background(), st, io.Discard, userSelector{email: "user@example.com"}, "short")
+	err := runChangePassword(context.Background(), st, io.Discard, userSelector{email: "user@example.com"}, "short", "")
 	require.Error(t, err)
 	require.Contains(t, err.Error(), "invalid password")
 }
@@ -87,31 +90,31 @@ func TestRunResetMFA(t *testing.T) {
 	st := newTestIDPStorage(t)
 	var out bytes.Buffer
 
-	encodedUserID := nbdex.EncodeDexUserID("user-1", localConnectorID)
-	err := runResetMFA(ctx, st, &out, userSelector{userID: encodedUserID})
+	encodedUserID := nbdex.EncodeDexUserID("user-1", idp.LocalConnectorID)
+	err := runResetMFA(ctx, st, &out, userSelector{userID: encodedUserID}, "")
 	require.NoError(t, err)
 	require.Contains(t, out.String(), "MFA reset")
 
-	identity, err := st.GetUserIdentity(ctx, "user-1", localConnectorID)
+	identity, err := st.GetUserIdentity(ctx, "user-1", idp.LocalConnectorID)
 	require.NoError(t, err)
 	require.Empty(t, identity.MFASecrets)
 	require.Empty(t, identity.WebAuthnCredentials)
 
-	_, err = st.GetAuthSession(ctx, "user-1", localConnectorID)
+	_, err = st.GetAuthSession(ctx, "user-1", idp.LocalConnectorID)
 	require.ErrorIs(t, err, storage.ErrNotFound)
 }
 
 func TestRunResetMFAWithoutEnrollment(t *testing.T) {
 	ctx := context.Background()
 	st := newTestIDPStorage(t)
-	require.NoError(t, st.UpdateUserIdentity(ctx, "user-1", localConnectorID, func(old storage.UserIdentity) (storage.UserIdentity, error) {
+	require.NoError(t, st.UpdateUserIdentity(ctx, "user-1", idp.LocalConnectorID, func(old storage.UserIdentity) (storage.UserIdentity, error) {
 		old.MFASecrets = nil
 		old.WebAuthnCredentials = nil
 		return old, nil
 	}))
 
 	var out bytes.Buffer
-	err := runResetMFA(ctx, st, &out, userSelector{email: "user@example.com"})
+	err := runResetMFA(ctx, st, &out, userSelector{email: "user@example.com"}, "")
 	require.NoError(t, err)
 	require.Contains(t, out.String(), "No MFA enrollment found")
 }
@@ -131,6 +134,67 @@ func TestSetIDPClientsMFA(t *testing.T) {
 	require.Equal(t, "disabled", status)
 }
 
+func newTestManagementStore(t *testing.T, localMFAEnabled bool) mgmtstore.Store {
+	t.Helper()
+	ctx := context.Background()
+	st, err := mgmtstore.NewStore(ctx, types.SqliteStoreEngine, t.TempDir(), nil, false)
+	require.NoError(t, err)
+	t.Cleanup(func() { require.NoError(t, st.Close(ctx)) })
+	require.NoError(t, st.SaveAccount(ctx, &types.Account{
+		Id:       "account-1",
+		Settings: &types.Settings{LocalMfaEnabled: localMFAEnabled},
+	}))
+	return st
+}
+
+func TestRunSetMFAEnabledDoesNotSaveWhenIDPUpdateFails(t *testing.T) {
+	ctx := context.Background()
+	managementStore := newTestManagementStore(t, false)
+	idpStorage := memory.New(slog.New(slog.NewTextHandler(io.Discard, nil)))
+
+	err := runSetMFAEnabled(ctx, Resources{Store: managementStore, IDPStorage: idpStorage}, io.Discard, true)
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "embedded IdP client")
+
+	settings, err := managementStore.GetAccountSettings(ctx, mgmtstore.LockingStrengthNone, "account-1")
+	require.NoError(t, err)
+	require.False(t, settings.LocalMfaEnabled)
+}
+
+func TestRunSetMFAEnabledUpdatesSettingsAfterIDP(t *testing.T) {
+	ctx := context.Background()
+	managementStore := newTestManagementStore(t, false)
+	idpStorage := newTestIDPStorage(t)
+
+	err := runSetMFAEnabled(ctx, Resources{Store: managementStore, IDPStorage: idpStorage}, io.Discard, true)
+	require.NoError(t, err)
+
+	settings, err := managementStore.GetAccountSettings(ctx, mgmtstore.LockingStrengthNone, "account-1")
+	require.NoError(t, err)
+	require.True(t, settings.LocalMfaEnabled)
+	clientStatus, err := idpClientsMFAStatus(ctx, idpStorage)
+	require.NoError(t, err)
+	require.Equal(t, "enabled", clientStatus)
+}
+
+func TestRunSetMFAEnabledSucceedsWithNilEventStore(t *testing.T) {
+	ctx := context.Background()
+	managementStore := newTestManagementStore(t, false)
+	idpStorage := newTestIDPStorage(t)
+	var out bytes.Buffer
+	var err error
+
+	require.NotPanics(t, func() {
+		err = runSetMFAEnabled(ctx, Resources{Store: managementStore, IDPStorage: idpStorage, EventStore: nil}, &out, true)
+	})
+	require.NoError(t, err)
+	require.Contains(t, out.String(), "Local MFA enabled")
+
+	settings, err := managementStore.GetAccountSettings(ctx, mgmtstore.LockingStrengthNone, "account-1")
+	require.NoError(t, err)
+	require.True(t, settings.LocalMfaEnabled)
+}
+
 func TestUserSelectorValidate(t *testing.T) {
 	require.NoError(t, userSelector{email: " user@example.com "}.validate())
 	require.NoError(t, userSelector{userID: "user-1"}.validate())
@@ -140,11 +204,37 @@ func TestUserSelectorValidate(t *testing.T) {
 
 func TestFindLocalUserNotFound(t *testing.T) {
 	st := newTestIDPStorage(t)
-	_, err := findLocalUser(context.Background(), st, userSelector{email: "missing@example.com"})
+	_, err := findLocalUser(context.Background(), st, userSelector{email: "missing@example.com"}, "")
 	require.Error(t, err)
 	require.True(t, strings.Contains(err.Error(), "not found"))
 }
 
+func TestFindLocalUserZeroUsersIncludesStoragePath(t *testing.T) {
+	st := memory.New(slog.New(slog.NewTextHandler(io.Discard, nil)))
+	_, err := findLocalUser(context.Background(), st, userSelector{email: "missing@example.com"}, "/var/lib/netbird/idp.db")
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "no local users exist")
+	require.Contains(t, err.Error(), "/var/lib/netbird/idp.db")
+}
+
+func TestUserCommandValidatesSelectorBeforeOpeningStorage(t *testing.T) {
+	opened := false
+	cmd := NewCommands(Openers{
+		IDP: func(cmd *cobra.Command, fn func(ctx context.Context, idpStorage storage.Storage, storageFile string) error) error {
+			opened = true
+			return nil
+		},
+	})
+	cmd.SetArgs([]string{"user", "change-password", "--password", "NewPass1!"})
+	cmd.SetOut(io.Discard)
+	cmd.SetErr(io.Discard)
+
+	err := cmd.Execute()
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "provide exactly one")
+	require.False(t, opened)
+}
+
 func TestResolvePasswordInputFromStdin(t *testing.T) {
 	cmd := &cobra.Command{}
 	cmd.SetIn(strings.NewReader("NewPass1!\n"))

management/cmd/admin_config_test.go

@@ -0,0 +1,80 @@
+package cmd
+
+import (
+	"context"
+	"path"
+	"path/filepath"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	nbconfig "github.com/netbirdio/netbird/management/internals/server/config"
+	"github.com/netbirdio/netbird/management/server/idp"
+)
+
+func TestApplyAdminDatadirOverrideRelocatesDefaultIDPStorage(t *testing.T) {
+	oldDatadir := filepath.Join(t.TempDir(), "old")
+	newDatadir := filepath.Join(t.TempDir(), "new")
+
+	for _, defaultFile := range []string{
+		"",
+		filepath.Join(oldDatadir, "idp.db"),
+		path.Join(oldDatadir, "idp.db"),
+	} {
+		t.Run(defaultFile, func(t *testing.T) {
+			cfg := &nbconfig.Config{
+				EmbeddedIdP: &idp.EmbeddedIdPConfig{
+					Enabled: true,
+					Storage: idp.EmbeddedStorageConfig{
+						Type: "sqlite3",
+						Config: idp.EmbeddedStorageTypeConfig{
+							File: defaultFile,
+						},
+					},
+				},
+			}
+			datadir := oldDatadir
+			oldAdminDatadir := adminDatadir
+			adminDatadir = newDatadir
+			t.Cleanup(func() { adminDatadir = oldAdminDatadir })
+
+			applyAdminDatadirOverride(cfg, &datadir)
+
+			require.Equal(t, newDatadir, datadir)
+			require.Equal(t, filepath.Join(newDatadir, "idp.db"), cfg.EmbeddedIdP.Storage.Config.File)
+		})
+	}
+}
+
+func TestOpenAdminEventStoreMissingEncryptionKeyReturnsNilInterface(t *testing.T) {
+	eventStore, err := openAdminEventStore(context.Background(), &nbconfig.Config{}, t.TempDir())
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "encryption key")
+	require.Nil(t, eventStore)
+}
+
+func TestApplyAdminDatadirOverrideKeepsExplicitIDPStorage(t *testing.T) {
+	oldDatadir := filepath.Join(t.TempDir(), "old")
+	newDatadir := filepath.Join(t.TempDir(), "new")
+	explicitFile := filepath.Join(t.TempDir(), "custom-idp.db")
+	cfg := &nbconfig.Config{
+		EmbeddedIdP: &idp.EmbeddedIdPConfig{
+			Enabled: true,
+			Storage: idp.EmbeddedStorageConfig{
+				Type: "sqlite3",
+				Config: idp.EmbeddedStorageTypeConfig{
+					File: explicitFile,
+				},
+			},
+		},
+	}
+	datadir := oldDatadir
+	oldAdminDatadir := adminDatadir
+	adminDatadir = newDatadir
+	t.Cleanup(func() { adminDatadir = oldAdminDatadir })
+
+	applyAdminDatadirOverride(cfg, &datadir)
+
+	require.Equal(t, newDatadir, datadir)
+	require.Equal(t, explicitFile, cfg.EmbeddedIdP.Storage.Config.File)
+}

management/cmd/management.go

@@ -13,6 +13,7 @@ import (
 	"os"
 	"os/signal"
 	"path"
+	"path/filepath"
 	"strings"
 	"syscall"
 
@@ -209,7 +210,7 @@ func ApplyEmbeddedIdPConfig(ctx context.Context, cfg *nbconfig.Config) error {
 		cfg.EmbeddedIdP.Storage.Type = "sqlite3"
 	}
 	if cfg.EmbeddedIdP.Storage.Config.File == "" && cfg.Datadir != "" {
-		cfg.EmbeddedIdP.Storage.Config.File = path.Join(cfg.Datadir, "idp.db")
+		cfg.EmbeddedIdP.Storage.Config.File = filepath.Join(cfg.Datadir, "idp.db")
 	}
 
 	issuer := cfg.EmbeddedIdP.Issuer

management/cmd/proxy/proxy.go

@@ -0,0 +1,141 @@
+// Package proxycmd provides reusable cobra commands for managing reverse proxy instances.
+// Both the management and combined binaries use these commands, each providing
+// their own StoreOpener to handle config loading and store initialization.
+package proxycmd
+
+import (
+	"bufio"
+	"context"
+	"fmt"
+	"io"
+	"strings"
+	"text/tabwriter"
+
+	"github.com/spf13/cobra"
+
+	rpproxy "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/proxy"
+	"github.com/netbirdio/netbird/management/server/store"
+)
+
+// StoreOpener initializes a store from the command context and calls fn.
+type StoreOpener func(cmd *cobra.Command, fn func(ctx context.Context, s store.Store) error) error
+
+const disconnectAllConfirmation = "disconnect all proxies"
+
+// NewCommands creates the proxy command tree with the given store opener.
+// Returns the parent "proxy" command with the disconnect-all subcommand.
+func NewCommands(opener StoreOpener) *cobra.Command {
+	var dryRun bool
+	var force bool
+
+	proxyCmd := &cobra.Command{
+		Use:   "proxy",
+		Short: "Manage reverse proxy instances",
+		Long:  "Commands for inspecting and repairing the reverse proxy instances registered with the management server.",
+	}
+
+	disconnectAllCmd := &cobra.Command{
+		Use:   "disconnect-all",
+		Short: "Force-mark all reverse proxy instances as disconnected",
+		Long: "Lists all reverse proxy instances and force-marks them as disconnected, regardless of their session state. " +
+			"Use this to repair stale connection state, e.g. after an unclean management server shutdown. " +
+			"By default, it asks for manual confirmation before changing state. Use --dry-run to preview without changing state, or --force to skip confirmation. " +
+			"Run during a maintenance window; affected live proxies may stay hidden until their next heartbeat or reconnect/re-register.",
+		Args: cobra.NoArgs,
+		RunE: func(cmd *cobra.Command, _ []string) error {
+			return opener(cmd, func(ctx context.Context, s store.Store) error {
+				return runDisconnectAll(ctx, s, cmd.OutOrStdout(), cmd.InOrStdin(), dryRun, force)
+			})
+		},
+	}
+	disconnectAllCmd.Flags().BoolVar(&dryRun, "dry-run", false, "List reverse proxy instances that would be disconnected without changing state")
+	disconnectAllCmd.Flags().BoolVar(&force, "force", false, "Skip the confirmation prompt and apply the repair")
+
+	proxyCmd.AddCommand(disconnectAllCmd)
+	return proxyCmd
+}
+
+func runDisconnectAll(ctx context.Context, s store.Store, out io.Writer, in io.Reader, dryRun, force bool) error {
+	proxies, err := s.GetAllProxies(ctx)
+	if err != nil {
+		return fmt.Errorf("list proxies: %w", err)
+	}
+
+	if len(proxies) == 0 {
+		_, _ = fmt.Fprintln(out, "No reverse proxy instances found.")
+		return nil
+	}
+
+	toDisconnect := 0
+	w := tabwriter.NewWriter(out, 0, 0, 2, ' ', 0)
+	_, _ = fmt.Fprintln(w, "ID\tCLUSTER\tIP\tACCOUNT\tSTATUS\tLAST SEEN")
+	_, _ = fmt.Fprintln(w, "--\t-------\t--\t-------\t------\t---------")
+
+	for _, p := range proxies {
+		if p.Status != rpproxy.StatusDisconnected {
+			toDisconnect++
+		}
+
+		account := "-"
+		if p.AccountID != nil {
+			account = *p.AccountID
+		}
+
+		_, _ = fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\t%s\n",
+			p.ID,
+			p.ClusterAddress,
+			p.IPAddress,
+			account,
+			p.Status,
+			p.LastSeen.Format("2006-01-02 15:04:05"),
+		)
+	}
+	if err := w.Flush(); err != nil {
+		return fmt.Errorf("write proxy list: %w", err)
+	}
+
+	if dryRun {
+		_, _ = fmt.Fprintf(out, "\nDry run: would force-mark %d of %d reverse proxy instance(s) as disconnected.\n", toDisconnect, len(proxies))
+		return nil
+	}
+
+	if !force {
+		confirmed, err := confirmDisconnectAll(out, in)
+		if err != nil {
+			return err
+		}
+		if !confirmed {
+			_, _ = fmt.Fprintln(out, "Aborted. No reverse proxy instances were changed.")
+			return nil
+		}
+	}
+
+	disconnected, err := s.DisconnectAllProxies(ctx)
+	if err != nil {
+		return fmt.Errorf("disconnect proxies: %w", err)
+	}
+
+	_, _ = fmt.Fprintf(out, "\nForce-marked %d of %d reverse proxy instance(s) as disconnected.\n", disconnected, len(proxies))
+	return nil
+}
+
+func confirmDisconnectAll(out io.Writer, in io.Reader) (bool, error) {
+	if in == nil {
+		in = strings.NewReader("")
+	}
+
+	_, _ = fmt.Fprintln(out, "\nWARNING: This command changes stored reverse proxy state for every non-disconnected instance.")
+	_, _ = fmt.Fprintln(out, "Run it during a maintenance window; affected live proxies may stay hidden until "+
+		"their next heartbeat or reconnect/re-register.")
+	_, _ = fmt.Fprintf(out, "Type %q to continue: ", disconnectAllConfirmation)
+
+	scanner := bufio.NewScanner(in)
+	if !scanner.Scan() {
+		if err := scanner.Err(); err != nil {
+			return false, fmt.Errorf("read confirmation: %w", err)
+		}
+		return false, nil
+	}
+
+	return strings.EqualFold(strings.TrimSpace(scanner.Text()), disconnectAllConfirmation), nil
+}

management/cmd/proxy/proxy_test.go

@@ -0,0 +1,180 @@
+package proxycmd
+
+import (
+	"bytes"
+	"context"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/spf13/cobra"
+	"github.com/stretchr/testify/require"
+
+	rpproxy "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/proxy"
+	"github.com/netbirdio/netbird/management/server/store"
+)
+
+func newTestStore(t *testing.T) store.Store {
+	t.Helper()
+
+	s, cleanup, err := store.NewTestStoreFromSQL(context.Background(), "", t.TempDir())
+	require.NoError(t, err)
+	t.Cleanup(cleanup)
+
+	return s
+}
+
+func seedProxies(t *testing.T, ctx context.Context, s store.Store) {
+	t.Helper()
+
+	accountID := "account-1"
+	alreadyDisconnectedAt := time.Now().Add(-time.Hour)
+	seed := []*rpproxy.Proxy{
+		{
+			ID:             "proxy-1",
+			SessionID:      "session-1",
+			ClusterAddress: "cluster-a.example.com",
+			IPAddress:      "10.0.0.1",
+			LastSeen:       time.Now(),
+			Status:         rpproxy.StatusConnected,
+		},
+		{
+			ID:             "proxy-2",
+			SessionID:      "session-2",
+			ClusterAddress: "cluster-b.example.com",
+			IPAddress:      "10.0.0.2",
+			AccountID:      &accountID,
+			LastSeen:       time.Now(),
+			Status:         rpproxy.StatusConnected,
+		},
+		{
+			ID:             "proxy-3",
+			SessionID:      "session-3",
+			ClusterAddress: "cluster-a.example.com",
+			IPAddress:      "10.0.0.3",
+			LastSeen:       time.Now().Add(-time.Hour),
+			Status:         rpproxy.StatusDisconnected,
+			DisconnectedAt: &alreadyDisconnectedAt,
+		},
+	}
+	for _, p := range seed {
+		require.NoError(t, s.SaveProxy(ctx, p))
+	}
+}
+
+func proxiesByID(t *testing.T, ctx context.Context, s store.Store) map[string]*rpproxy.Proxy {
+	t.Helper()
+
+	proxies, err := s.GetAllProxies(ctx)
+	require.NoError(t, err)
+	require.Len(t, proxies, 3)
+
+	byID := make(map[string]*rpproxy.Proxy, len(proxies))
+	for _, p := range proxies {
+		byID[p.ID] = p
+	}
+	return byID
+}
+
+func TestRunDisconnectAllWithConfirmation(t *testing.T) {
+	ctx := context.Background()
+	s := newTestStore(t)
+	seedProxies(t, ctx, s)
+
+	var out bytes.Buffer
+	require.NoError(t, runDisconnectAll(ctx, s, &out, strings.NewReader(disconnectAllConfirmation+"\n"), false, false))
+
+	output := out.String()
+	require.Contains(t, output, "proxy-1")
+	require.Contains(t, output, "proxy-2")
+	require.Contains(t, output, "proxy-3")
+	require.Contains(t, output, "cluster-a.example.com")
+	require.Contains(t, output, "account-1")
+	require.Contains(t, output, "Type \"disconnect all proxies\" to continue")
+	require.Contains(t, output, "Force-marked 2 of 3 reverse proxy instance(s) as disconnected.")
+
+	for _, p := range proxiesByID(t, ctx, s) {
+		require.Equal(t, rpproxy.StatusDisconnected, p.Status, "proxy %s should be disconnected", p.ID)
+		require.NotNil(t, p.DisconnectedAt, "proxy %s should have a disconnected timestamp", p.ID)
+	}
+}
+
+func TestRunDisconnectAllForceSkipsConfirmation(t *testing.T) {
+	ctx := context.Background()
+	s := newTestStore(t)
+	seedProxies(t, ctx, s)
+
+	var out bytes.Buffer
+	require.NoError(t, runDisconnectAll(ctx, s, &out, strings.NewReader(""), false, true))
+
+	output := out.String()
+	require.NotContains(t, output, "Type \"disconnect all proxies\" to continue")
+	require.Contains(t, output, "Force-marked 2 of 3 reverse proxy instance(s) as disconnected.")
+}
+
+func TestRunDisconnectAllAbortLeavesProxiesUnchanged(t *testing.T) {
+	ctx := context.Background()
+	s := newTestStore(t)
+	seedProxies(t, ctx, s)
+
+	var out bytes.Buffer
+	require.NoError(t, runDisconnectAll(ctx, s, &out, strings.NewReader("no\n"), false, false))
+
+	output := out.String()
+	require.Contains(t, output, "Type \"disconnect all proxies\" to continue")
+	require.Contains(t, output, "Aborted. No reverse proxy instances were changed.")
+
+	byID := proxiesByID(t, ctx, s)
+	require.Equal(t, rpproxy.StatusConnected, byID["proxy-1"].Status)
+	require.Equal(t, rpproxy.StatusConnected, byID["proxy-2"].Status)
+	require.Equal(t, rpproxy.StatusDisconnected, byID["proxy-3"].Status)
+}
+
+func TestRunDisconnectAllDryRunLeavesProxiesUnchanged(t *testing.T) {
+	ctx := context.Background()
+	s := newTestStore(t)
+	seedProxies(t, ctx, s)
+
+	var out bytes.Buffer
+	require.NoError(t, runDisconnectAll(ctx, s, &out, strings.NewReader(""), true, false))
+
+	output := out.String()
+	require.Contains(t, output, "Dry run: would force-mark 2 of 3 reverse proxy instance(s) as disconnected.")
+	require.NotContains(t, output, "Type \"disconnect all proxies\" to continue")
+
+	byID := proxiesByID(t, ctx, s)
+	require.Equal(t, rpproxy.StatusConnected, byID["proxy-1"].Status)
+	require.Equal(t, rpproxy.StatusConnected, byID["proxy-2"].Status)
+	require.Equal(t, rpproxy.StatusDisconnected, byID["proxy-3"].Status)
+}
+
+func TestNewCommandsDisconnectAllDryRun(t *testing.T) {
+	ctx := context.Background()
+	s := newTestStore(t)
+	seedProxies(t, ctx, s)
+
+	opened := false
+	cmd := NewCommands(func(cmd *cobra.Command, fn func(ctx context.Context, s store.Store) error) error {
+		opened = true
+		return fn(cmd.Context(), s)
+	})
+
+	var out bytes.Buffer
+	cmd.SetOut(&out)
+	cmd.SetErr(&out)
+	cmd.SetIn(strings.NewReader(""))
+	cmd.SetArgs([]string{"disconnect-all", "--dry-run"})
+
+	require.NoError(t, cmd.ExecuteContext(ctx))
+	require.True(t, opened)
+	require.Contains(t, out.String(), "Dry run: would force-mark 2 of 3 reverse proxy instance(s) as disconnected.")
+}
+
+func TestRunDisconnectAllEmpty(t *testing.T) {
+	ctx := context.Background()
+	s := newTestStore(t)
+
+	var out bytes.Buffer
+	require.NoError(t, runDisconnectAll(ctx, s, &out, strings.NewReader(""), false, false))
+	require.Contains(t, out.String(), "No reverse proxy instances found.")
+}

management/cmd/root.go

@@ -86,4 +86,5 @@ func init() {
 	ac := newAdminCommands()
 	ac.PersistentFlags().StringVar(&nbconfig.MgmtConfigPath, "config", defaultMgmtConfig, "Netbird config file location")
 	rootCmd.AddCommand(ac)
+	rootCmd.AddCommand(newLegacyTokenCommand())
 }