diff --git a/api/controllers/console/auth/error.py b/api/controllers/console/auth/error.py index 81f1c6e70fa3d1..a3021f0a3202ff 100644 --- a/api/controllers/console/auth/error.py +++ b/api/controllers/console/auth/error.py @@ -155,3 +155,21 @@ class MemberNotInTenantError(BaseHTTPException): error_code = "member_not_in_tenant" description = "The member is not in the workspace." code = 400 + + +class MFARequiredError(BaseHTTPException): + error_code = "mfa_required" + description = "Multi-factor authentication is required." + code = 401 + + +class MFATokenRequiredError(BaseHTTPException): + error_code = "mfa_token_invalid" + description = "The MFA token is invalid or expired." + code = 401 + + +class MFASetupRequiredError(BaseHTTPException): + error_code = "mfa_setup_required" + description = "MFA setup is required to complete this action." + code = 400 diff --git a/api/controllers/console/auth/login.py b/api/controllers/console/auth/login.py index 3b35ab3c237487..223b906b01aa0f 100644 --- a/api/controllers/console/auth/login.py +++ b/api/controllers/console/auth/login.py @@ -32,6 +32,7 @@ from services.errors.account import AccountRegisterError from services.errors.workspace import WorkSpaceNotAllowedCreateError, WorkspacesLimitExceededError from services.feature_service import FeatureService +from services.mfa_service import MFAService class LoginApi(Resource): @@ -46,6 +47,8 @@ def post(self): parser.add_argument("password", type=str, required=True, location="json") parser.add_argument("remember_me", type=bool, required=False, default=False, location="json") parser.add_argument("invite_token", type=str, required=False, default=None, location="json") + parser.add_argument("mfa_code", type=str, required=False, default=None, location="json") + parser.add_argument("is_backup_code", type=bool, required=False, default=False, location="json") args = parser.parse_args() if dify_config.BILLING_ENABLED and BillingService.is_email_in_freeze(args["email"]): @@ -72,7 +75,22 @@ def post(self): raise AccountBannedError() except services.errors.account.AccountPasswordError: AccountService.add_login_error_rate_limit(args["email"]) + # Perform dummy MFA check to prevent timing attacks + # This ensures similar response time regardless of authentication failure + if args.get("mfa_code"): + import time + + time.sleep(0.05) # Simulate MFA verification delay raise AuthenticationFailedError() + + # Check MFA requirement + if MFAService.is_mfa_required(account): + if not args["mfa_code"]: + return {"result": "fail", "code": "mfa_required"} + + if not MFAService.authenticate_with_mfa(account, args["mfa_code"]): + return {"result": "fail", "code": "mfa_token_invalid", "data": "The MFA token is invalid or expired."} + # SELF_HOSTED only have one workspace tenants = TenantService.get_join_tenants(account) if len(tenants) == 0: diff --git a/api/controllers/console/auth/mfa.py b/api/controllers/console/auth/mfa.py new file mode 100644 index 00000000000000..cec71ffe73cc4f --- /dev/null +++ b/api/controllers/console/auth/mfa.py @@ -0,0 +1,140 @@ +import logging +from typing import cast + +import flask_login +from flask_restx import Resource, reqparse + +from controllers.console.wraps import account_initialization_required +from libs.login import login_required +from models.account import Account +from services.mfa_service import MFAService + + +class MFASetupInitApi(Resource): + @login_required + @account_initialization_required + def get(self): + """Initialize MFA setup - generate secret and QR code (GET method for compatibility).""" + # Call post method directly with proper context + account = cast(Account, flask_login.current_user) + + try: + mfa_status = MFAService.get_mfa_status(account) + if mfa_status["enabled"]: + return {"error": "MFA is already enabled"}, 400 + + setup_data = MFAService.generate_mfa_setup_data(account) + return {"secret": setup_data["secret"], "qr_code": setup_data["qr_code"]} + except Exception: + # Log the actual error for debugging + logging.exception("MFA setup error") + # Return generic error message to avoid exposing internal details + return {"error": "Failed to setup MFA. Please try again."}, 500 + + @login_required + @account_initialization_required + def post(self): # type: ignore + """Initialize MFA setup - generate secret and QR code.""" + account = cast(Account, flask_login.current_user) + + try: + mfa_status = MFAService.get_mfa_status(account) + if mfa_status["enabled"]: + return {"error": "MFA is already enabled"}, 400 + + setup_data = MFAService.generate_mfa_setup_data(account) + return {"secret": setup_data["secret"], "qr_code": setup_data["qr_code"]} + except Exception: + # Log the actual error for debugging + logging.exception("MFA setup error") + # Return generic error message to avoid exposing internal details + return {"error": "Failed to setup MFA. Please try again."}, 500 + + +class MFASetupCompleteApi(Resource): + @login_required + @account_initialization_required + def post(self): + """Complete MFA setup with TOTP verification.""" + parser = reqparse.RequestParser() + parser.add_argument("totp_token", type=str, required=True, help="TOTP token is required") + args = parser.parse_args() + + account = cast(Account, flask_login.current_user) + + try: + result = MFAService.setup_mfa(account, args["totp_token"]) + return { + "message": "MFA setup completed successfully", + "backup_codes": result["backup_codes"], + "setup_at": result["setup_at"].isoformat(), + } + except ValueError as e: + return {"error": str(e)}, 400 + except Exception as e: + return {"error": str(e)}, 500 + + +class MFADisableApi(Resource): + @login_required + @account_initialization_required + def post(self): + """Disable MFA with password verification.""" + parser = reqparse.RequestParser() + parser.add_argument("password", type=str, required=True, help="Password is required") + args = parser.parse_args() + + account = cast(Account, flask_login.current_user) + + try: + mfa_status = MFAService.get_mfa_status(account) + if not mfa_status["enabled"]: + return {"error": "MFA is not enabled"}, 400 + + if MFAService.disable_mfa(account, args["password"]): + return {"message": "MFA disabled successfully"} + else: + return {"error": "Invalid password"}, 400 + except Exception as e: + return {"error": str(e)}, 500 + + +class MFAStatusApi(Resource): + @login_required + @account_initialization_required + def get(self): + """Get current MFA status.""" + account = cast(Account, flask_login.current_user) + + try: + status = MFAService.get_mfa_status(account) + return status + except Exception as e: + return {"error": str(e)}, 500 + + +class MFAVerifyApi(Resource): + def post(self): + """Verify MFA token during login (public endpoint).""" + parser = reqparse.RequestParser() + parser.add_argument("email", type=str, required=True, help="Email is required") + parser.add_argument("mfa_token", type=str, required=True, help="MFA token is required") + args = parser.parse_args() + + from models.engine import db + + account = db.session.query(Account).filter_by(email=args["email"]).first() + + if not account: + return {"error": "Account not found"}, 404 + + if not MFAService.is_mfa_required(account): + return {"error": "MFA not required for this account"}, 400 + + try: + if MFAService.authenticate_with_mfa(account, args["mfa_token"]): + return {"message": "MFA verification successful"} + else: + return {"error": "Invalid MFA token"}, 400 + except Exception as e: + return {"error": str(e)}, 500 diff --git a/api/controllers/console/workspace/account.py b/api/controllers/console/workspace/account.py index 7a41a8a5ccd683..0fdc129bcab6b9 100644 --- a/api/controllers/console/workspace/account.py +++ b/api/controllers/console/workspace/account.py @@ -583,3 +583,11 @@ def post(self): api.add_resource(CheckEmailUnique, "/account/change-email/check-email-unique") # api.add_resource(AccountEmailApi, '/account/email') # api.add_resource(AccountEmailVerifyApi, '/account/email-verify') + +# MFA endpoints +from controllers.console.auth.mfa import MFADisableApi, MFASetupCompleteApi, MFASetupInitApi, MFAStatusApi + +api.add_resource(MFAStatusApi, "/account/mfa/status") +api.add_resource(MFASetupInitApi, "/account/mfa/setup") +api.add_resource(MFASetupCompleteApi, "/account/mfa/setup/complete") +api.add_resource(MFADisableApi, "/account/mfa/disable") diff --git a/api/migrations/versions/2025_09_22_1350-c699fb438ab7_add_account_mfa_settings.py b/api/migrations/versions/2025_09_22_1350-c699fb438ab7_add_account_mfa_settings.py new file mode 100644 index 00000000000000..eec623dc06a874 --- /dev/null +++ b/api/migrations/versions/2025_09_22_1350-c699fb438ab7_add_account_mfa_settings.py @@ -0,0 +1,43 @@ +"""add account mfa settings table + +Revision ID: c699fb438ab7 +Revises: 68519ad5cd18 +Create Date: 2025-09-22 13:50:00.000000 + +""" +from alembic import op +import models as models +import sqlalchemy as sa +from sqlalchemy.dialects import postgresql + +# revision identifiers, used by Alembic. +revision = 'c699fb438ab7' +down_revision = '68519ad5cd18' +branch_labels = None +depends_on = None + + +def upgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.create_table('account_mfa_settings', + sa.Column('id', models.types.StringUUID(), server_default=sa.text('uuid_generate_v4()'), nullable=False), + sa.Column('account_id', models.types.StringUUID(), nullable=False), + sa.Column('enabled', sa.Boolean(), server_default=sa.text('false'), nullable=False), + sa.Column('secret', sa.Text(), nullable=True), + sa.Column('backup_codes', sa.Text(), nullable=True), + sa.Column('setup_at', sa.DateTime(), nullable=True), + sa.Column('created_at', sa.DateTime(), server_default=sa.text('CURRENT_TIMESTAMP'), nullable=False), + sa.Column('updated_at', sa.DateTime(), server_default=sa.text('CURRENT_TIMESTAMP'), nullable=False), + sa.ForeignKeyConstraint(['account_id'], ['accounts.id'], ), + sa.PrimaryKeyConstraint('id', name='account_mfa_settings_pkey'), + sa.UniqueConstraint('account_id', name='unique_account_mfa_settings') + ) + op.create_index('account_mfa_settings_account_id_idx', 'account_mfa_settings', ['account_id'], unique=False) + # ### end Alembic commands ### + + +def downgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.drop_index('account_mfa_settings_account_id_idx', table_name='account_mfa_settings') + op.drop_table('account_mfa_settings') + # ### end Alembic commands ### \ No newline at end of file diff --git a/api/models/__init__.py b/api/models/__init__.py index 779484283feda7..be832d2b77973d 100644 --- a/api/models/__init__.py +++ b/api/models/__init__.py @@ -1,6 +1,7 @@ from .account import ( Account, AccountIntegrate, + AccountMFASettings, AccountStatus, InvitationCode, Tenant, @@ -97,6 +98,7 @@ "APIBasedExtensionPoint", "Account", "AccountIntegrate", + "AccountMFASettings", "AccountStatus", "ApiRequest", "ApiToken", diff --git a/api/models/account.py b/api/models/account.py index 8c1f990aa2179c..92f2647062c701 100644 --- a/api/models/account.py +++ b/api/models/account.py @@ -6,7 +6,7 @@ import sqlalchemy as sa from flask_login import UserMixin # type: ignore[import-untyped] from sqlalchemy import DateTime, String, func, select -from sqlalchemy.orm import Mapped, Session, mapped_column, reconstructor +from sqlalchemy.orm import Mapped, Session, backref, mapped_column, reconstructor, relationship from typing_extensions import deprecated from models.base import Base @@ -361,3 +361,24 @@ class UpgradeMode(enum.StrEnum): include_plugins: Mapped[list[str]] = mapped_column(sa.ARRAY(String(255)), nullable=False) # plugin_id (author/name) created_at: Mapped[datetime] = mapped_column(DateTime, nullable=False, server_default=func.current_timestamp()) updated_at: Mapped[datetime] = mapped_column(DateTime, nullable=False, server_default=func.current_timestamp()) + + +class AccountMFASettings(Base): + __tablename__ = "account_mfa_settings" + __table_args__ = ( + sa.PrimaryKeyConstraint("id", name="account_mfa_settings_pkey"), + sa.UniqueConstraint("account_id", name="unique_account_mfa_settings"), + sa.Index("account_mfa_settings_account_id_idx", "account_id"), + ) + + id: Mapped[str] = mapped_column(StringUUID, server_default=sa.text("uuid_generate_v4()")) + account_id: Mapped[str] = mapped_column(StringUUID, sa.ForeignKey("accounts.id"), nullable=False) + enabled: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, server_default=sa.text("false")) + secret: Mapped[str | None] = mapped_column(sa.Text, nullable=True) + backup_codes: Mapped[str | None] = mapped_column(sa.Text, nullable=True) + setup_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True) + created_at: Mapped[datetime] = mapped_column(DateTime, nullable=False, server_default=func.current_timestamp()) + updated_at: Mapped[datetime] = mapped_column(DateTime, nullable=False, server_default=func.current_timestamp()) + + # Relationship + account = relationship("Account", backref=backref("mfa_settings", uselist=False, cascade="all, delete-orphan")) diff --git a/api/pyproject.toml b/api/pyproject.toml index 012702edd25282..e49a4c2e46ba33 100644 --- a/api/pyproject.toml +++ b/api/pyproject.toml @@ -68,10 +68,12 @@ dependencies = [ "pydantic-extra-types~=2.10.3", "pydantic-settings~=2.9.1", "pyjwt~=2.10.1", + "pyotp~=2.9.0", "pypdfium2==4.30.0", "python-docx~=1.1.0", "python-dotenv==1.0.1", "pyyaml~=6.0.1", + "qrcode~=7.4.2", "readabilipy~=0.3.0", "redis[hiredis]~=6.1.0", "resend~=2.9.0", diff --git a/api/services/mfa_service.py b/api/services/mfa_service.py new file mode 100644 index 00000000000000..9afb7cf4ef7c44 --- /dev/null +++ b/api/services/mfa_service.py @@ -0,0 +1,255 @@ +import base64 +import hashlib +import io +import json +import logging +import secrets +from datetime import UTC, datetime +from typing import cast + +import pyotp +import qrcode + +from core.helper import encrypter +from models.account import Account, AccountMFASettings +from models.engine import db + + +class MFAService: + @staticmethod + def generate_secret() -> str: + """Generate a new TOTP secret for the user.""" + return pyotp.random_base32() + + @staticmethod + def generate_backup_codes(count: int = 8) -> list[str]: + """Generate backup codes for account recovery.""" + codes = [] + for _ in range(count): + code = secrets.token_hex(4).upper() + codes.append(code) + return codes + + @staticmethod + def generate_qr_code(account: Account, secret: str) -> str: + """Generate QR code for TOTP setup.""" + totp = pyotp.TOTP(secret) + provisioning_uri = totp.provisioning_uri(name=account.email, issuer_name="Dify") + + # Generate QR code + qr = qrcode.QRCode( + version=1, + error_correction=qrcode.constants.ERROR_CORRECT_L, + box_size=10, + border=4, + ) + qr.add_data(provisioning_uri) + qr.make(fit=True) + + # Create image + img = qr.make_image(fill_color="black", back_color="white") + + # Convert to base64 + buffer = io.BytesIO() + img.save(buffer) + img_str = base64.b64encode(buffer.getvalue()).decode() + + return f"data:image/png;base64,{img_str}" + + @staticmethod + def verify_totp(secret: str, token: str, tenant_id: str | None = None) -> bool: + """Verify TOTP token.""" + if not secret: + return False + try: + # Decrypt secret if tenant_id provided + if tenant_id: + secret = encrypter.decrypt_token(tenant_id, secret) + totp = pyotp.TOTP(secret) + return totp.verify(token, valid_window=1) + except (ValueError, TypeError): + logging.exception("TOTP verification failed") + return False + + @staticmethod + def get_or_create_mfa_settings(account: Account) -> AccountMFASettings: + """Get or create MFA settings for account.""" + mfa_settings = db.session.query(AccountMFASettings).filter_by(account_id=account.id).first() + if not mfa_settings: + mfa_settings = AccountMFASettings(account_id=account.id) + db.session.add(mfa_settings) + db.session.commit() + return mfa_settings + + @staticmethod + def _hash_backup_code(code: str) -> str: + """Hash a backup code for storage.""" + return hashlib.sha256(code.upper().encode()).hexdigest() + + @staticmethod + def verify_backup_code(mfa_settings: AccountMFASettings, code: str) -> bool: + """Verify and consume backup code.""" + if mfa_settings.backup_codes is None: # type: ignore + return False + + try: + # Hash the provided code + hashed_code = MFAService._hash_backup_code(code) + + # Load stored hashed codes + assert mfa_settings.backup_codes is not None + backup_codes_hashed = json.loads(cast(str, mfa_settings.backup_codes)) + + if hashed_code in backup_codes_hashed: + # Remove used backup code + backup_codes_hashed.remove(hashed_code) + mfa_settings.backup_codes = json.dumps(backup_codes_hashed) + db.session.commit() + return True + except json.JSONDecodeError: + pass + + return False + + @staticmethod + def setup_mfa(account: Account, totp_token: str) -> dict: + """Setup MFA for account with TOTP verification.""" + mfa_settings = MFAService.get_or_create_mfa_settings(account) + + if mfa_settings.enabled is True: + raise ValueError("MFA is already enabled for this account") + + if mfa_settings.secret is None: + raise ValueError("MFA secret not generated") + + # Get tenant ID from account - try multiple sources + tenant_id = account.current_tenant_id + if not tenant_id: + # Try to get from TenantAccountJoin + from models.account import TenantAccountJoin + + tenant_join = db.session.query(TenantAccountJoin).filter_by(account_id=account.id).first() + if tenant_join: + tenant_id = tenant_join.tenant_id + + if not tenant_id: + raise ValueError("No tenant associated with account") + + # Verify TOTP token with decryption + assert mfa_settings.secret is not None + if not MFAService.verify_totp(cast(str, mfa_settings.secret), totp_token, tenant_id): + raise ValueError("Invalid TOTP token") + + # Generate backup codes + backup_codes = MFAService.generate_backup_codes() + + # Hash backup codes for storage + backup_codes_hashed = [MFAService._hash_backup_code(code) for code in backup_codes] + + # Enable MFA + mfa_settings.enabled = True + mfa_settings.backup_codes = json.dumps(backup_codes_hashed) + mfa_settings.setup_at = datetime.now(UTC) + + db.session.commit() + + # Return the plain backup codes (user must save them) + return {"backup_codes": backup_codes, "setup_at": mfa_settings.setup_at} + + @staticmethod + def disable_mfa(account: Account, password: str) -> bool: + """Disable MFA for account after password verification.""" + from libs.password import compare_password + + # Verify password + if account.password is None or not compare_password(password, account.password, account.password_salt): + return False + + mfa_settings = db.session.query(AccountMFASettings).filter_by(account_id=account.id).first() + if not mfa_settings: + return True # Already disabled + + # Disable MFA + mfa_settings.enabled = False + mfa_settings.secret = None + mfa_settings.backup_codes = None + mfa_settings.setup_at = None + + db.session.commit() + return True + + @staticmethod + def generate_mfa_setup_data(account: Account) -> dict: + """Generate MFA setup data including secret and QR code.""" + mfa_settings = MFAService.get_or_create_mfa_settings(account) + + if mfa_settings.enabled: + raise ValueError("MFA is already enabled for this account") + + # Get tenant ID from account + tenant_id = account.current_tenant_id + if not tenant_id: + raise ValueError("No tenant associated with account") + + # Generate new secret + secret = MFAService.generate_secret() + + # Encrypt secret for storage + encrypted_secret = encrypter.encrypt_token(tenant_id, secret) + mfa_settings.secret = encrypted_secret + db.session.commit() + + # Generate QR code with plain secret + qr_code = MFAService.generate_qr_code(account, secret) + + # Return plain secret for user display (only shown once) + return {"secret": secret, "qr_code": qr_code} + + @staticmethod + def is_mfa_required(account: Account) -> bool: + """Check if MFA is required for this account.""" + mfa_settings = db.session.query(AccountMFASettings).filter_by(account_id=account.id).first() + return bool(mfa_settings and mfa_settings.enabled and mfa_settings.secret is not None) + + @staticmethod + def authenticate_with_mfa(account: Account, token: str) -> bool: + """Authenticate user with MFA token (TOTP or backup code).""" + mfa_settings = db.session.query(AccountMFASettings).filter_by(account_id=account.id).first() + + if not mfa_settings or not mfa_settings.enabled: + return True + + # Get tenant ID from account - try multiple sources + tenant_id = account.current_tenant_id + if not tenant_id: + # Try to get from TenantAccountJoin + from models.account import TenantAccountJoin + + tenant_join = db.session.query(TenantAccountJoin).filter_by(account_id=account.id).first() + if tenant_join: + tenant_id = tenant_join.tenant_id + + # Try TOTP first with decryption + assert mfa_settings.secret is not None + if MFAService.verify_totp(cast(str, mfa_settings.secret), token, tenant_id): + return True + + # Try backup code (already hashed) + if MFAService.verify_backup_code(mfa_settings, token): + return True + + return False + + @staticmethod + def get_mfa_status(account: Account) -> dict: + """Get MFA status for account.""" + mfa_settings = db.session.query(AccountMFASettings).filter_by(account_id=account.id).first() + + if not mfa_settings: + return {"enabled": False, "setup_at": None, "has_backup_codes": False} + + return { + "enabled": mfa_settings.enabled, + "setup_at": mfa_settings.setup_at.isoformat() if mfa_settings.setup_at is not None else None, + "has_backup_codes": mfa_settings.backup_codes is not None, + } diff --git a/api/tests/integration_tests/controllers/console/auth/test_login_mfa_integration.py b/api/tests/integration_tests/controllers/console/auth/test_login_mfa_integration.py new file mode 100644 index 00000000000000..1d7a396c80c886 --- /dev/null +++ b/api/tests/integration_tests/controllers/console/auth/test_login_mfa_integration.py @@ -0,0 +1,385 @@ +import json +from unittest.mock import Mock, patch + +from flask import Flask + + +class TestLoginMFAIntegration: + @patch("controllers.console.auth.login.FeatureService.get_system_features") + @patch("controllers.console.auth.login.dify_config") + @patch("controllers.console.auth.login.BillingService.is_email_in_freeze") + @patch("controllers.console.auth.login.AccountService.is_login_error_rate_limit") + @patch("controllers.console.auth.login.AccountService.authenticate") + @patch("controllers.console.auth.login.MFAService.is_mfa_required") + @patch("controllers.console.auth.login.TenantService.get_join_tenants") + @patch("controllers.console.auth.login.AccountService.login") + @patch("controllers.console.auth.login.AccountService.reset_login_error_rate_limit") + @patch("controllers.console.auth.login.extract_remote_ip") + def test_login_without_mfa_success( + self, + mock_extract_ip, + mock_reset_limit, + mock_login_service, + mock_get_tenants, + mock_is_mfa_required, + mock_authenticate, + mock_rate_limit, + mock_freeze_check, + mock_dify_config, + mock_system_features, + test_client, + setup_account, + ): + """Test successful login without MFA enabled.""" + # Setup mocks + mock_dify_config.BILLING_ENABLED = False + mock_freeze_check.return_value = False + mock_rate_limit.return_value = False + mock_authenticate.return_value = setup_account + mock_is_mfa_required.return_value = False + mock_get_tenants.return_value = [Mock()] # At least one tenant + mock_extract_ip.return_value = "127.0.0.1" + + token_pair_mock = Mock() + token_pair_mock.model_dump.return_value = { + "access_token": "test_access_token", + "refresh_token": "test_refresh_token", + } + mock_login_service.return_value = token_pair_mock + + with ( + patch("controllers.console.auth.login.setup_required") as mock_setup, + patch("controllers.console.auth.login.email_password_login_enabled") as mock_email_enabled, + ): + mock_setup.return_value = lambda f: f + mock_email_enabled.return_value = lambda f: f + + response = test_client.post( + "/console/api/login", json={"email": setup_account.email, "password": "TestPassword123"} + ) + + assert response.status_code == 200 + data = json.loads(response.data) + assert data["result"] == "success" + assert "access_token" in data["data"] + + @patch("controllers.console.auth.login.FeatureService.get_system_features") + @patch("controllers.console.auth.login.dify_config") + @patch("controllers.console.auth.login.BillingService.is_email_in_freeze") + @patch("controllers.console.auth.login.AccountService.is_login_error_rate_limit") + @patch("controllers.console.auth.login.AccountService.authenticate") + @patch("controllers.console.auth.login.MFAService.is_mfa_required") + def test_login_with_mfa_required_no_token( + self, + mock_is_mfa_required, + mock_authenticate, + mock_rate_limit, + mock_freeze_check, + mock_dify_config, + mock_system_features, + test_client, + setup_account, + ): + """Test login returns mfa_required when MFA is enabled but no token provided.""" + # Setup mocks + mock_dify_config.BILLING_ENABLED = False + mock_freeze_check.return_value = False + mock_rate_limit.return_value = False + mock_authenticate.return_value = setup_account + mock_is_mfa_required.return_value = True + + with ( + patch("controllers.console.auth.login.setup_required") as mock_setup, + patch("controllers.console.auth.login.email_password_login_enabled") as mock_email_enabled, + ): + mock_setup.return_value = lambda f: f + mock_email_enabled.return_value = lambda f: f + + response = test_client.post( + "/console/api/login", json={"email": "test@example.com", "password": "TestPassword123"} + ) + + assert response.status_code == 200 + data = json.loads(response.data) + assert data["result"] == "fail" + assert data["code"] == "mfa_required" + + @patch("controllers.console.auth.login.FeatureService.get_system_features") + @patch("controllers.console.auth.login.dify_config") + @patch("controllers.console.auth.login.BillingService.is_email_in_freeze") + @patch("controllers.console.auth.login.AccountService.is_login_error_rate_limit") + @patch("controllers.console.auth.login.AccountService.authenticate") + @patch("controllers.console.auth.login.MFAService.is_mfa_required") + @patch("controllers.console.auth.login.MFAService.authenticate_with_mfa") + def test_login_with_mfa_invalid_token( + self, + mock_auth_mfa, + mock_is_mfa_required, + mock_authenticate, + mock_rate_limit, + mock_freeze_check, + mock_dify_config, + mock_system_features, + test_client, + setup_account, + ): + """Test login fails with invalid MFA token.""" + # Setup mocks + mock_dify_config.BILLING_ENABLED = False + mock_freeze_check.return_value = False + mock_rate_limit.return_value = False + mock_authenticate.return_value = setup_account + mock_is_mfa_required.return_value = True + mock_auth_mfa.return_value = False # Invalid token + + with ( + patch("controllers.console.auth.login.setup_required") as mock_setup, + patch("controllers.console.auth.login.email_password_login_enabled") as mock_email_enabled, + ): + mock_setup.return_value = lambda f: f + mock_email_enabled.return_value = lambda f: f + + response = test_client.post( + "/console/api/login", + json={"email": "test@example.com", "password": "TestPassword123", "mfa_code": "invalid_token"}, + ) + + assert response.status_code == 200 + data = json.loads(response.data) + assert data["result"] == "fail" + assert data["code"] == "mfa_token_invalid" + assert data["data"] == "The MFA token is invalid or expired." + + @patch("controllers.console.auth.login.FeatureService.get_system_features") + @patch("controllers.console.auth.login.dify_config") + @patch("controllers.console.auth.login.BillingService.is_email_in_freeze") + @patch("controllers.console.auth.login.AccountService.is_login_error_rate_limit") + @patch("controllers.console.auth.login.AccountService.authenticate") + @patch("controllers.console.auth.login.MFAService.is_mfa_required") + @patch("controllers.console.auth.login.MFAService.authenticate_with_mfa") + @patch("controllers.console.auth.login.TenantService.get_join_tenants") + @patch("controllers.console.auth.login.AccountService.login") + @patch("controllers.console.auth.login.AccountService.reset_login_error_rate_limit") + @patch("controllers.console.auth.login.extract_remote_ip") + def test_login_with_mfa_valid_token_success( + self, + mock_extract_ip, + mock_reset_limit, + mock_login_service, + mock_get_tenants, + mock_auth_mfa, + mock_is_mfa_required, + mock_authenticate, + mock_rate_limit, + mock_freeze_check, + mock_dify_config, + mock_system_features, + test_client, + setup_account, + ): + """Test successful login with valid MFA token.""" + # Setup mocks + mock_dify_config.BILLING_ENABLED = False + mock_freeze_check.return_value = False + mock_rate_limit.return_value = False + mock_authenticate.return_value = setup_account + mock_is_mfa_required.return_value = True + mock_auth_mfa.return_value = True # Valid token + mock_get_tenants.return_value = [Mock()] # At least one tenant + mock_extract_ip.return_value = "127.0.0.1" + + token_pair_mock = Mock() + token_pair_mock.model_dump.return_value = { + "access_token": "test_access_token", + "refresh_token": "test_refresh_token", + } + mock_login_service.return_value = token_pair_mock + + with ( + patch("controllers.console.auth.login.setup_required") as mock_setup, + patch("controllers.console.auth.login.email_password_login_enabled") as mock_email_enabled, + ): + mock_setup.return_value = lambda f: f + mock_email_enabled.return_value = lambda f: f + + response = test_client.post( + "/console/api/login", + json={"email": "test@example.com", "password": "TestPassword123", "mfa_code": "123456"}, + ) + + assert response.status_code == 200 + data = json.loads(response.data) + assert data["result"] == "success" + assert "access_token" in data["data"] + + # Verify MFA authentication was called + mock_auth_mfa.assert_called_once_with(setup_account, "123456") + + @patch("controllers.console.auth.login.FeatureService.get_system_features") + @patch("controllers.console.auth.login.dify_config") + @patch("controllers.console.auth.login.BillingService.is_email_in_freeze") + @patch("controllers.console.auth.login.AccountService.is_login_error_rate_limit") + @patch("controllers.console.auth.login.AccountService.authenticate") + @patch("controllers.console.auth.login.MFAService.is_mfa_required") + @patch("controllers.console.auth.login.MFAService.authenticate_with_mfa") + @patch("controllers.console.auth.login.TenantService.get_join_tenants") + @patch("controllers.console.auth.login.AccountService.login") + @patch("controllers.console.auth.login.AccountService.reset_login_error_rate_limit") + @patch("controllers.console.auth.login.extract_remote_ip") + def test_login_with_mfa_backup_code_success( + self, + mock_extract_ip, + mock_reset_limit, + mock_login_service, + mock_get_tenants, + mock_auth_mfa, + mock_is_mfa_required, + mock_authenticate, + mock_rate_limit, + mock_freeze_check, + mock_dify_config, + mock_system_features, + test_client, + setup_account, + ): + """Test successful login with valid backup code.""" + # Setup mocks + mock_dify_config.BILLING_ENABLED = False + mock_freeze_check.return_value = False + mock_rate_limit.return_value = False + mock_authenticate.return_value = setup_account + mock_is_mfa_required.return_value = True + mock_auth_mfa.return_value = True # Valid backup code + mock_get_tenants.return_value = [Mock()] # At least one tenant + mock_extract_ip.return_value = "127.0.0.1" + + token_pair_mock = Mock() + token_pair_mock.model_dump.return_value = { + "access_token": "test_access_token", + "refresh_token": "test_refresh_token", + } + mock_login_service.return_value = token_pair_mock + + with ( + patch("controllers.console.auth.login.setup_required") as mock_setup, + patch("controllers.console.auth.login.email_password_login_enabled") as mock_email_enabled, + ): + mock_setup.return_value = lambda f: f + mock_email_enabled.return_value = lambda f: f + + response = test_client.post( + "/console/api/login", + json={ + "email": "test@example.com", + "password": "TestPassword123", + "mfa_code": "BACKUP123", # Backup code format + }, + ) + + assert response.status_code == 200 + data = json.loads(response.data) + assert data["result"] == "success" + assert "access_token" in data["data"] + + # Verify MFA authentication was called with backup code + mock_auth_mfa.assert_called_once_with(setup_account, "BACKUP123") + + @patch("controllers.console.auth.login.FeatureService.get_system_features") + @patch("controllers.console.auth.login.dify_config") + @patch("controllers.console.auth.login.BillingService.is_email_in_freeze") + @patch("controllers.console.auth.login.AccountService.is_login_error_rate_limit") + @patch("controllers.console.auth.login.AccountService.authenticate") + @patch("controllers.console.auth.login.MFAService.is_mfa_required") + def test_login_mfa_flow_order( + self, + mock_is_mfa_required, + mock_authenticate, + mock_rate_limit, + mock_freeze_check, + mock_dify_config, + mock_system_features, + test_client, + ): + """Test that MFA check happens after password authentication.""" + # Setup mocks - password auth fails + mock_dify_config.BILLING_ENABLED = False + mock_freeze_check.return_value = False + mock_rate_limit.return_value = False + + # Mock password authentication failure + from services.errors.account import AccountPasswordError + + mock_authenticate.side_effect = AccountPasswordError() + + with ( + patch("controllers.console.auth.login.setup_required") as mock_setup, + patch("controllers.console.auth.login.email_password_login_enabled") as mock_email_enabled, + patch("controllers.console.auth.login.AccountService.add_login_error_rate_limit") as mock_add_limit, + ): + mock_setup.return_value = lambda f: f + mock_email_enabled.return_value = lambda f: f + + response = test_client.post( + "/console/api/login", + json={"email": "test@example.com", "password": "WrongPassword123", "mfa_code": "123456"}, + ) + + # Password error should trigger EmailOrPasswordMismatchError + assert response.status_code == 400 + + # MFA check should not be called if password auth fails + mock_is_mfa_required.assert_not_called() + + +class TestMFAEndToEndFlow: + """End-to-end tests for complete MFA flow.""" + + def setup_method(self): + self.app = Flask(__name__) + self.app.config["TESTING"] = True + self.client = self.app.test_client() + + @patch("services.mfa_service.MFAService.generate_secret") + @patch("services.mfa_service.MFAService.generate_qr_code") + @patch("services.mfa_service.MFAService.verify_totp") + @patch("services.mfa_service.MFAService.generate_backup_codes") + @patch("services.mfa_service.db.session") + def test_complete_mfa_setup_flow(self, mock_session, mock_gen_codes, mock_verify, mock_gen_qr, mock_gen_secret): + """Test complete MFA setup flow from init to completion.""" + from models.account import Account + from services.mfa_service import MFAService + + # Mock account + account = Mock(spec=Account) + account.id = "test-id" + account.email = "test@example.com" + + # Setup mocks + mock_gen_secret.return_value = "TESTSECRET123" + mock_gen_qr.return_value = "data:image/png;base64,test" + mock_verify.return_value = True + mock_gen_codes.return_value = ["CODE1", "CODE2", "CODE3"] + + # Step 1: Initialize MFA setup + with patch("services.mfa_service.MFAService.get_or_create_mfa_settings") as mock_get_settings: + mfa_settings = Mock() + mfa_settings.enabled = False + mfa_settings.secret = None + mock_get_settings.return_value = mfa_settings + + setup_data = MFAService.generate_mfa_setup_data(account) + + assert setup_data["secret"] == "TESTSECRET123" + assert setup_data["qr_code"] == "data:image/png;base64,test" + assert mfa_settings.secret == "TESTSECRET123" + + # Step 2: Complete MFA setup + with patch("services.mfa_service.MFAService.get_or_create_mfa_settings") as mock_get_settings: + mfa_settings.secret = "TESTSECRET123" + mock_get_settings.return_value = mfa_settings + + result = MFAService.setup_mfa(account, "123456") + + assert mfa_settings.enabled is True + assert result["backup_codes"] == ["CODE1", "CODE2", "CODE3"] + assert mfa_settings.setup_at is not None diff --git a/api/tests/integration_tests/controllers/console/auth/test_mfa_endpoints.py b/api/tests/integration_tests/controllers/console/auth/test_mfa_endpoints.py new file mode 100644 index 00000000000000..326e45fd24a712 --- /dev/null +++ b/api/tests/integration_tests/controllers/console/auth/test_mfa_endpoints.py @@ -0,0 +1,141 @@ +from datetime import UTC, datetime +from unittest.mock import patch + +import pytest + +from services.account_service import AccountService +from services.mfa_service import MFAService + + +class TestMFAEndpoints: + """Test MFA endpoints using integration test approach.""" + + @pytest.fixture + def auth_header(self, setup_account): + """Get authentication header with JWT token.""" + token = AccountService.get_account_jwt_token(setup_account) + return {"Authorization": f"Bearer {token}"} + + def test_mfa_status_success(self, test_client, setup_account, auth_header): + """Test successful MFA status check.""" + with patch.object(MFAService, "get_mfa_status") as mock_status: + mock_status.return_value = {"enabled": False, "setup_at": None} + + response = test_client.get("/console/api/account/mfa/status", headers=auth_header) + + assert response.status_code == 200 + data = response.json + assert data["enabled"] is False + assert data["setup_at"] is None + mock_status.assert_called_once_with(setup_account) + + def test_mfa_setup_init_success(self, test_client, setup_account, auth_header): + """Test successful MFA setup initialization.""" + with patch.object(MFAService, "get_mfa_status") as mock_status: + with patch.object(MFAService, "generate_mfa_setup_data") as mock_generate: + mock_status.return_value = {"enabled": False} + mock_generate.return_value = {"secret": "TEST_SECRET", "qr_code": "data:image/png;base64,test"} + + response = test_client.post("/console/api/account/mfa/setup", headers=auth_header) + + assert response.status_code == 200 + data = response.json + assert data["secret"] == "TEST_SECRET" + assert data["qr_code"] == "data:image/png;base64,test" + mock_generate.assert_called_once_with(setup_account) + + def test_mfa_setup_init_already_enabled(self, test_client, setup_account, auth_header): + """Test MFA setup initialization when already enabled.""" + with patch.object(MFAService, "get_mfa_status") as mock_status: + mock_status.return_value = {"enabled": True, "setup_at": "2024-01-01T00:00:00"} + + response = test_client.post("/console/api/account/mfa/setup", headers=auth_header) + + assert response.status_code == 400 + data = response.json + assert data["error"] == "MFA is already enabled" + + def test_mfa_setup_complete_success(self, test_client, setup_account, auth_header): + """Test successful MFA setup completion.""" + with patch.object(MFAService, "setup_mfa") as mock_setup: + mock_setup.return_value = { + "backup_codes": ["CODE1", "CODE2", "CODE3", "CODE4", "CODE5", "CODE6", "CODE7", "CODE8"], + "setup_at": datetime(2024, 1, 1, 0, 0, 0, tzinfo=UTC), + } + + response = test_client.post( + "/console/api/account/mfa/setup/complete", headers=auth_header, json={"totp_token": "123456"} + ) + + assert response.status_code == 200 + data = response.json + assert data["message"] == "MFA setup completed successfully" + assert len(data["backup_codes"]) == 8 + assert data["setup_at"] == "2024-01-01T00:00:00+00:00" + mock_setup.assert_called_once_with(setup_account, "123456") + + def test_mfa_setup_complete_missing_token(self, test_client, setup_account, auth_header): + """Test MFA setup completion with missing token.""" + response = test_client.post("/console/api/account/mfa/setup/complete", headers=auth_header, json={}) + + assert response.status_code == 400 + data = response.json + assert "message" in data + assert "TOTP token is required" in data["message"] + + def test_mfa_setup_complete_invalid_token(self, test_client, setup_account, auth_header): + """Test MFA setup completion with invalid token.""" + with patch.object(MFAService, "setup_mfa") as mock_setup: + mock_setup.side_effect = ValueError("Invalid TOTP token") + + response = test_client.post( + "/console/api/account/mfa/setup/complete", headers=auth_header, json={"totp_token": "999999"} + ) + + assert response.status_code == 400 + data = response.json + assert "Invalid TOTP token" in data["error"] + + def test_mfa_disable_success(self, test_client, setup_account, auth_header): + """Test successful MFA disable.""" + with patch.object(MFAService, "get_mfa_status") as mock_status: + with patch.object(MFAService, "disable_mfa") as mock_disable: + mock_status.return_value = {"enabled": True} + mock_disable.return_value = True + + response = test_client.post( + "/console/api/account/mfa/disable", headers=auth_header, json={"password": "test_password"} + ) + + assert response.status_code == 200 + data = response.json + assert data["message"] == "MFA disabled successfully" + mock_disable.assert_called_once_with(setup_account, "test_password") + + def test_mfa_disable_wrong_password(self, test_client, setup_account, auth_header): + """Test MFA disable with wrong password.""" + with patch.object(MFAService, "get_mfa_status") as mock_status: + with patch.object(MFAService, "disable_mfa") as mock_disable: + mock_status.return_value = {"enabled": True} + mock_disable.return_value = False + + response = test_client.post( + "/console/api/account/mfa/disable", headers=auth_header, json={"password": "wrong_password"} + ) + + assert response.status_code == 400 + data = response.json + assert data["error"] == "Invalid password" + + def test_mfa_disable_not_enabled(self, test_client, setup_account, auth_header): + """Test MFA disable when not enabled.""" + with patch.object(MFAService, "get_mfa_status") as mock_status: + mock_status.return_value = {"enabled": False} + + response = test_client.post( + "/console/api/account/mfa/disable", headers=auth_header, json={"password": "test_password"} + ) + + assert response.status_code == 400 + data = response.json + assert data["error"] == "MFA is not enabled" diff --git a/api/tests/unit_tests/services/test_mfa_service.py b/api/tests/unit_tests/services/test_mfa_service.py new file mode 100644 index 00000000000000..bfb2d0eeffed7a --- /dev/null +++ b/api/tests/unit_tests/services/test_mfa_service.py @@ -0,0 +1,374 @@ +import hashlib +import json +import unittest +from datetime import datetime +from unittest.mock import Mock, patch + +import pytest + +from models.account import Account, AccountMFASettings +from services.mfa_service import MFAService + + +class TestMFAService(unittest.TestCase): + def setUp(self): + self.account = Mock(spec=Account) + self.account.id = "test-account-id" + self.account.email = "test@example.com" + self.account.password = "hashed_password" + self.account.password_salt = "salt" + self.account.current_tenant_id = "test-tenant-id" + + self.mfa_settings = Mock(spec=AccountMFASettings) + self.mfa_settings.account_id = self.account.id + self.mfa_settings.enabled = False + self.mfa_settings.secret = None + self.mfa_settings.backup_codes = None + self.mfa_settings.setup_at = None + + def test_generate_secret(self): + """Test secret generation.""" + secret = MFAService.generate_secret() + assert isinstance(secret, str) + assert len(secret) == 32 # Base32 length + + def test_generate_backup_codes(self): + """Test backup codes generation.""" + codes = MFAService.generate_backup_codes() + assert len(codes) == 8 + for code in codes: + assert isinstance(code, str) + assert len(code) == 8 # 4 hex bytes = 8 chars + + @patch("pyotp.TOTP") + def test_verify_totp_valid(self, mock_totp_class): + """Test TOTP verification with valid token.""" + mock_totp = Mock() + mock_totp.verify.return_value = True + mock_totp_class.return_value = mock_totp + + result = MFAService.verify_totp("test_secret", "123456") + + assert result + mock_totp.verify.assert_called_once_with("123456", valid_window=1) + + @patch("pyotp.TOTP") + def test_verify_totp_invalid(self, mock_totp_class): + """Test TOTP verification with invalid token.""" + mock_totp = Mock() + mock_totp.verify.return_value = False + mock_totp_class.return_value = mock_totp + + result = MFAService.verify_totp("test_secret", "invalid") + + assert not result + + def test_verify_totp_no_secret(self): + """Test TOTP verification with no secret.""" + result = MFAService.verify_totp(None, "123456") + assert not result + + @patch("services.mfa_service.db.session") + def test_get_or_create_mfa_settings_existing(self, mock_session): + """Test getting existing MFA settings.""" + mock_session.query.return_value.filter_by.return_value.first.return_value = self.mfa_settings + + result = MFAService.get_or_create_mfa_settings(self.account) + + assert result == self.mfa_settings + mock_session.query.assert_called_once() + + @patch("services.mfa_service.db.session") + def test_get_or_create_mfa_settings_new(self, mock_session): + """Test creating new MFA settings.""" + mock_session.query.return_value.filter_by.return_value.first.return_value = None + + result = MFAService.get_or_create_mfa_settings(self.account) + + # Check that new settings were created + assert isinstance(result, AccountMFASettings) + assert result.account_id == self.account.id + mock_session.add.assert_called_once() + mock_session.commit.assert_called_once() + + @patch("services.mfa_service.db.session") + def test_verify_backup_code_valid(self, mock_session): + """Test backup code verification with valid code.""" + # Store hashed codes + hash1 = hashlib.sha256(b"ABCD1234").hexdigest() + hash2 = hashlib.sha256(b"EFGH5678").hexdigest() + self.mfa_settings.backup_codes = json.dumps([hash1, hash2]) + + result = MFAService.verify_backup_code(self.mfa_settings, "abcd1234") # Test case insensitive + + assert result + # Check that the code was removed (comparing hashes) + remaining_codes = json.loads(self.mfa_settings.backup_codes) + assert hash1 not in remaining_codes + assert hash2 in remaining_codes + mock_session.commit.assert_called_once() + + def test_verify_backup_code_invalid(self): + """Test backup code verification with invalid code.""" + # Store hashed codes + hash1 = hashlib.sha256(b"ABCD1234").hexdigest() + hash2 = hashlib.sha256(b"EFGH5678").hexdigest() + self.mfa_settings.backup_codes = json.dumps([hash1, hash2]) + + result = MFAService.verify_backup_code(self.mfa_settings, "INVALID") + + assert not result + + def test_verify_backup_code_no_codes(self): + """Test backup code verification with no backup codes.""" + self.mfa_settings.backup_codes = None + + result = MFAService.verify_backup_code(self.mfa_settings, "ABCD1234") + + assert not result + + @patch("services.mfa_service.MFAService.get_or_create_mfa_settings") + @patch("services.mfa_service.MFAService.verify_totp") + @patch("services.mfa_service.MFAService.generate_backup_codes") + @patch("services.mfa_service.db.session") + def test_setup_mfa_success(self, mock_session, mock_gen_codes, mock_verify, mock_get_settings): + """Test successful MFA setup.""" + mock_get_settings.return_value = self.mfa_settings + self.mfa_settings.secret = "test_secret" + mock_verify.return_value = True + mock_gen_codes.return_value = ["CODE1", "CODE2"] + + result = MFAService.setup_mfa(self.account, "123456") + + assert self.mfa_settings.enabled + # Backup codes are now hashed + hash1 = hashlib.sha256(b"CODE1").hexdigest() + hash2 = hashlib.sha256(b"CODE2").hexdigest() + assert self.mfa_settings.backup_codes == json.dumps([hash1, hash2]) + assert self.mfa_settings.setup_at is not None + assert result["backup_codes"] == ["CODE1", "CODE2"] + + @patch("services.mfa_service.MFAService.get_or_create_mfa_settings") + def test_setup_mfa_already_enabled(self, mock_get_settings): + """Test MFA setup when already enabled.""" + self.mfa_settings.enabled = True + mock_get_settings.return_value = self.mfa_settings + + with pytest.raises(ValueError) as context: + MFAService.setup_mfa(self.account, "123456") + + assert "already enabled" in str(context.value) + + @patch("services.mfa_service.MFAService.get_or_create_mfa_settings") + def test_setup_mfa_no_secret(self, mock_get_settings): + """Test MFA setup without secret.""" + mock_get_settings.return_value = self.mfa_settings + + with pytest.raises(ValueError) as context: + MFAService.setup_mfa(self.account, "123456") + + assert "secret not generated" in str(context.value) + + @patch("services.mfa_service.MFAService.get_or_create_mfa_settings") + @patch("services.mfa_service.MFAService.verify_totp") + def test_setup_mfa_invalid_token(self, mock_verify, mock_get_settings): + """Test MFA setup with invalid TOTP token.""" + mock_get_settings.return_value = self.mfa_settings + self.mfa_settings.secret = "test_secret" + mock_verify.return_value = False + + with pytest.raises(ValueError) as context: + MFAService.setup_mfa(self.account, "invalid") + + assert "Invalid TOTP token" in str(context.value) + + @patch("services.mfa_service.db.session") + def test_is_mfa_required_enabled(self, mock_session): + """Test MFA requirement check when enabled.""" + self.mfa_settings.enabled = True + self.mfa_settings.secret = "test_secret" + mock_session.query.return_value.filter_by.return_value.first.return_value = self.mfa_settings + + result = MFAService.is_mfa_required(self.account) + + assert result + + @patch("services.mfa_service.db.session") + def test_is_mfa_required_disabled(self, mock_session): + """Test MFA requirement check when disabled.""" + mock_session.query.return_value.filter_by.return_value.first.return_value = self.mfa_settings + + result = MFAService.is_mfa_required(self.account) + + assert not result + + @patch("services.mfa_service.db.session") + def test_is_mfa_required_no_settings(self, mock_session): + """Test MFA requirement check with no settings.""" + mock_session.query.return_value.filter_by.return_value.first.return_value = None + + result = MFAService.is_mfa_required(self.account) + + assert not result + + @patch("services.mfa_service.db.session") + @patch("services.mfa_service.MFAService.verify_totp") + @patch("services.mfa_service.MFAService.verify_backup_code") + def test_authenticate_with_mfa_totp_success(self, mock_verify_backup, mock_verify_totp, mock_session): + """Test MFA authentication with valid TOTP.""" + self.mfa_settings.enabled = True + self.mfa_settings.secret = "test_secret" + mock_session.query.return_value.filter_by.return_value.first.return_value = self.mfa_settings + mock_verify_totp.return_value = True + + result = MFAService.authenticate_with_mfa(self.account, "123456") + + assert result + mock_verify_totp.assert_called_once_with("test_secret", "123456", "test-tenant-id") + mock_verify_backup.assert_not_called() + + @patch("services.mfa_service.db.session") + @patch("services.mfa_service.MFAService.verify_totp") + @patch("services.mfa_service.MFAService.verify_backup_code") + def test_authenticate_with_mfa_backup_success(self, mock_verify_backup, mock_verify_totp, mock_session): + """Test MFA authentication with valid backup code.""" + self.mfa_settings.enabled = True + self.mfa_settings.secret = "test_secret" + mock_session.query.return_value.filter_by.return_value.first.return_value = self.mfa_settings + mock_verify_totp.return_value = False + mock_verify_backup.return_value = True + + result = MFAService.authenticate_with_mfa(self.account, "BACKUP123") + + assert result + mock_verify_totp.assert_called_once_with("test_secret", "BACKUP123", "test-tenant-id") + mock_verify_backup.assert_called_once_with(self.mfa_settings, "BACKUP123") + + @patch("services.mfa_service.db.session") + def test_authenticate_with_mfa_disabled(self, mock_session): + """Test MFA authentication when disabled.""" + mock_session.query.return_value.filter_by.return_value.first.return_value = self.mfa_settings + + result = MFAService.authenticate_with_mfa(self.account, "123456") + + assert result + + @patch("services.mfa_service.db.session") + def test_get_mfa_status_enabled(self, mock_session): + """Test getting MFA status when enabled.""" + self.mfa_settings.enabled = True + self.mfa_settings.setup_at = datetime(2025, 1, 1, 12, 0, 0) + self.mfa_settings.backup_codes = json.dumps(["CODE1", "CODE2"]) + mock_session.query.return_value.filter_by.return_value.first.return_value = self.mfa_settings + + result = MFAService.get_mfa_status(self.account) + + expected = {"enabled": True, "setup_at": "2025-01-01T12:00:00", "has_backup_codes": True} + assert result == expected + + @patch("services.mfa_service.db.session") + def test_get_mfa_status_no_settings(self, mock_session): + """Test getting MFA status with no settings.""" + mock_session.query.return_value.filter_by.return_value.first.return_value = None + + result = MFAService.get_mfa_status(self.account) + + expected = {"enabled": False, "setup_at": None, "has_backup_codes": False} + assert result == expected + + @patch("qrcode.QRCode") + @patch("pyotp.TOTP") + def test_generate_qr_code(self, mock_totp_class, mock_qr_class): + """Test QR code generation.""" + # Mock TOTP + mock_totp = Mock() + mock_totp.provisioning_uri.return_value = "otpauth://totp/test" + mock_totp_class.return_value = mock_totp + + # Mock QR code + mock_qr = Mock() + mock_img = Mock() + mock_qr.make_image.return_value = mock_img + mock_qr_class.return_value = mock_qr + + # Mock image buffer + with patch("io.BytesIO") as mock_buffer, patch("base64.b64encode") as mock_b64: + mock_b64.return_value.decode.return_value = "base64data" + + result = MFAService.generate_qr_code(self.account, "test_secret") + + assert result == "data:image/png;base64,base64data" + mock_totp.provisioning_uri.assert_called_once_with(name=self.account.email, issuer_name="Dify") + + @patch("libs.password.compare_password") + @patch("services.mfa_service.db.session") + def test_disable_mfa_success(self, mock_session, mock_compare_password): + """Test successful MFA disable.""" + mock_compare_password.return_value = True + mock_session.query.return_value.filter_by.return_value.first.return_value = self.mfa_settings + + result = MFAService.disable_mfa(self.account, "correct_password") + + assert result + assert not self.mfa_settings.enabled + assert self.mfa_settings.secret is None + assert self.mfa_settings.backup_codes is None + assert self.mfa_settings.setup_at is None + mock_session.commit.assert_called_once() + + @patch("libs.password.compare_password") + def test_disable_mfa_wrong_password(self, mock_compare_password): + """Test MFA disable with wrong password.""" + mock_compare_password.return_value = False + + result = MFAService.disable_mfa(self.account, "wrong_password") + + assert not result + + @patch("libs.password.compare_password") + @patch("services.mfa_service.db.session") + def test_disable_mfa_no_settings(self, mock_session, mock_compare_password): + """Test MFA disable when no settings exist.""" + mock_compare_password.return_value = True + mock_session.query.return_value.filter_by.return_value.first.return_value = None + + result = MFAService.disable_mfa(self.account, "correct_password") + + assert result # Already disabled + + @patch("services.mfa_service.encrypter") + @patch("services.mfa_service.MFAService.get_or_create_mfa_settings") + @patch("services.mfa_service.MFAService.generate_secret") + @patch("services.mfa_service.MFAService.generate_qr_code") + @patch("services.mfa_service.db.session") + def test_generate_mfa_setup_data_success( + self, mock_session, mock_gen_qr, mock_gen_secret, mock_get_settings, mock_encrypter + ): + """Test successful MFA setup data generation.""" + mock_get_settings.return_value = self.mfa_settings + mock_gen_secret.return_value = "NEWSECRET123" + mock_gen_qr.return_value = "data:image/png;base64,qrdata" + mock_encrypter.encrypt_token.return_value = "ENCRYPTED_SECRET" + + result = MFAService.generate_mfa_setup_data(self.account) + + assert result["secret"] == "NEWSECRET123" + assert result["qr_code"] == "data:image/png;base64,qrdata" + assert self.mfa_settings.secret == "ENCRYPTED_SECRET" + mock_encrypter.encrypt_token.assert_called_once_with("test-tenant-id", "NEWSECRET123") + mock_session.commit.assert_called_once() + + @patch("services.mfa_service.MFAService.get_or_create_mfa_settings") + def test_generate_mfa_setup_data_already_enabled(self, mock_get_settings): + """Test MFA setup data generation when already enabled.""" + self.mfa_settings.enabled = True + mock_get_settings.return_value = self.mfa_settings + + with pytest.raises(ValueError) as context: + MFAService.generate_mfa_setup_data(self.account) + + assert "already enabled" in str(context.value) + + +if __name__ == "__main__": + unittest.main() diff --git a/api/tests/unit_tests/services/test_mfa_service_security.py b/api/tests/unit_tests/services/test_mfa_service_security.py new file mode 100644 index 00000000000000..c475efa5f0df96 --- /dev/null +++ b/api/tests/unit_tests/services/test_mfa_service_security.py @@ -0,0 +1,189 @@ +import hashlib +import json +import unittest +from unittest.mock import Mock, patch + +import pytest + +from models.account import Account, AccountMFASettings +from services.mfa_service import MFAService + + +class TestMFAServiceSecurity(unittest.TestCase): + """Test MFA service security features including encryption and hashing.""" + + def setUp(self): + self.account = Mock(spec=Account) + self.account.id = "test-account-id" + self.account.email = "test@example.com" + self.account.current_tenant_id = "test-tenant-id" + + self.mfa_settings = Mock(spec=AccountMFASettings) + self.mfa_settings.account_id = self.account.id + self.mfa_settings.enabled = False + self.mfa_settings.secret = None + self.mfa_settings.backup_codes = None + self.mfa_settings.setup_at = None + + def test_hash_backup_code(self): + """Test backup code hashing.""" + code = "ABCD1234" + expected_hash = hashlib.sha256(code.upper().encode()).hexdigest() + + hashed = MFAService._hash_backup_code(code) + + assert hashed == expected_hash + # Verify the hash is consistent + assert hashed == MFAService._hash_backup_code(code.lower()) + + @patch("services.mfa_service.encrypter.encrypt_token") + @patch("services.mfa_service.MFAService.get_or_create_mfa_settings") + @patch("services.mfa_service.MFAService.generate_qr_code") + @patch("services.mfa_service.db.session") + def test_generate_mfa_setup_data_encrypts_secret( + self, mock_session, mock_generate_qr, mock_get_settings, mock_encrypt + ): + """Test that MFA setup encrypts the secret before storing.""" + mock_get_settings.return_value = self.mfa_settings + mock_encrypt.return_value = "encrypted_secret" + mock_generate_qr.return_value = "data:image/png;base64,fake_qr" + + result = MFAService.generate_mfa_setup_data(self.account) + + # Verify encryption was called with the tenant ID + mock_encrypt.assert_called_once() + call_args = mock_encrypt.call_args[0] + assert call_args[0] == "test-tenant-id" # tenant_id + assert len(call_args[1]) == 32 # Base32 secret + + # Verify encrypted secret was stored + assert self.mfa_settings.secret == "encrypted_secret" + mock_session.commit.assert_called_once() + + # Verify plain secret is returned (for user display) + assert "secret" in result + assert result["secret"] != "encrypted_secret" # Should be plain secret + assert len(result["secret"]) == 32 # Base32 length + + @patch("services.mfa_service.encrypter.decrypt_token") + @patch("pyotp.TOTP") + def test_verify_totp_with_encryption(self, mock_totp_class, mock_decrypt): + """Test TOTP verification with encrypted secret.""" + mock_decrypt.return_value = "decrypted_secret" + mock_totp = Mock() + mock_totp.verify.return_value = True + mock_totp_class.return_value = mock_totp + + result = MFAService.verify_totp("encrypted_secret", "123456", "test-tenant-id") + + # Verify decryption was called + mock_decrypt.assert_called_once_with("test-tenant-id", "encrypted_secret") + # Verify TOTP was created with decrypted secret + mock_totp_class.assert_called_once_with("decrypted_secret") + assert result + + @patch("services.mfa_service.db.session") + def test_verify_backup_code_with_hash(self, mock_session): + """Test backup code verification with hashed storage.""" + # Pre-hashed backup codes + hashed_codes = [ + MFAService._hash_backup_code("CODE1234"), + MFAService._hash_backup_code("CODE5678"), + ] + self.mfa_settings.backup_codes = json.dumps(hashed_codes) + + # Test valid code + result = MFAService.verify_backup_code(self.mfa_settings, "code1234") # Test case insensitive + + assert result + # Verify the code was removed + remaining = json.loads(self.mfa_settings.backup_codes) + assert len(remaining) == 1 + assert MFAService._hash_backup_code("CODE1234") not in remaining + + @patch("services.mfa_service.db.session") + def test_verify_backup_code_invalid_with_hash(self, mock_session): + """Test backup code verification fails with wrong code.""" + hashed_codes = [ + MFAService._hash_backup_code("CODE1234"), + MFAService._hash_backup_code("CODE5678"), + ] + self.mfa_settings.backup_codes = json.dumps(hashed_codes) + + result = MFAService.verify_backup_code(self.mfa_settings, "WRONGCODE") + + assert not result + # Verify no codes were removed + remaining = json.loads(self.mfa_settings.backup_codes) + assert len(remaining) == 2 + + @patch("services.mfa_service.encrypter.decrypt_token") + @patch("services.mfa_service.MFAService.get_or_create_mfa_settings") + @patch("pyotp.TOTP") + @patch("services.mfa_service.MFAService.generate_backup_codes") + @patch("services.mfa_service.db.session") + def test_setup_mfa_with_security_features( + self, mock_session, mock_gen_codes, mock_totp_class, mock_get_settings, mock_decrypt + ): + """Test MFA setup with both encryption and hashing.""" + mock_get_settings.return_value = self.mfa_settings + self.mfa_settings.secret = "encrypted_secret" + + # Setup decryption + mock_decrypt.return_value = "decrypted_secret" + + # Setup TOTP verification + mock_totp = Mock() + mock_totp.verify.return_value = True + mock_totp_class.return_value = mock_totp + + # Setup backup codes + mock_gen_codes.return_value = ["CODE1", "CODE2", "CODE3"] + + result = MFAService.setup_mfa(self.account, "123456") + + # Verify secret was decrypted for verification + mock_decrypt.assert_called_once_with("test-tenant-id", "encrypted_secret") + + # Verify backup codes were hashed before storage + stored_codes = json.loads(self.mfa_settings.backup_codes) + assert len(stored_codes) == 3 + assert stored_codes[0] == MFAService._hash_backup_code("CODE1") + assert stored_codes[1] == MFAService._hash_backup_code("CODE2") + assert stored_codes[2] == MFAService._hash_backup_code("CODE3") + + # Verify plain codes are returned to user + assert result["backup_codes"] == ["CODE1", "CODE2", "CODE3"] + + @patch("services.mfa_service.encrypter.decrypt_token") + @patch("services.mfa_service.db.session") + @patch("pyotp.TOTP") + def test_authenticate_with_mfa_encrypted(self, mock_totp_class, mock_session, mock_decrypt): + """Test authentication with encrypted MFA secret.""" + self.mfa_settings.enabled = True + self.mfa_settings.secret = "encrypted_secret" + self.mfa_settings.backup_codes = json.dumps([]) + + mock_session.query.return_value.filter_by.return_value.first.return_value = self.mfa_settings + mock_decrypt.return_value = "decrypted_secret" + + mock_totp = Mock() + mock_totp.verify.return_value = True + mock_totp_class.return_value = mock_totp + + result = MFAService.authenticate_with_mfa(self.account, "123456") + + assert result + mock_decrypt.assert_called_once_with("test-tenant-id", "encrypted_secret") + + def test_no_tenant_id_raises_error(self): + """Test that operations fail gracefully when no tenant ID is available.""" + self.account.current_tenant_id = None + + with patch("services.mfa_service.MFAService.get_or_create_mfa_settings") as mock_get_settings: + mock_get_settings.return_value = self.mfa_settings + + with pytest.raises(ValueError) as context: + MFAService.generate_mfa_setup_data(self.account) + + assert "No tenant associated" in str(context.value) diff --git a/api/uv.lock b/api/uv.lock index 7ce71cd2158f2a..5555288321db3e 100644 --- a/api/uv.lock +++ b/api/uv.lock @@ -1341,10 +1341,12 @@ dependencies = [ { name = "pydantic-extra-types" }, { name = "pydantic-settings" }, { name = "pyjwt" }, + { name = "pyotp" }, { name = "pypdfium2" }, { name = "python-docx" }, { name = "python-dotenv" }, { name = "pyyaml" }, + { name = "qrcode" }, { name = "readabilipy" }, { name = "redis", extra = ["hiredis"] }, { name = "resend" }, @@ -1535,10 +1537,12 @@ requires-dist = [ { name = "pydantic-extra-types", specifier = "~=2.10.3" }, { name = "pydantic-settings", specifier = "~=2.9.1" }, { name = "pyjwt", specifier = "~=2.10.1" }, + { name = "pyotp", specifier = "~=2.9.0" }, { name = "pypdfium2", specifier = "==4.30.0" }, { name = "python-docx", specifier = "~=1.1.0" }, { name = "python-dotenv", specifier = "==1.0.1" }, { name = "pyyaml", specifier = "~=6.0.1" }, + { name = "qrcode", specifier = "~=7.4.2" }, { name = "readabilipy", specifier = "~=0.3.0" }, { name = "redis", extras = ["hiredis"], specifier = "~=6.1.0" }, { name = "resend", specifier = "~=2.9.0" }, @@ -4827,6 +4831,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/83/7b/c103cca858de87476db5e7c7f0f386b429c3057a7291155c70560b15d951/pyobvector-0.2.16-py3-none-any.whl", hash = "sha256:0710272e5c807a6d0bdeee96972cdc9fdca04fc4b40c2d1260b08ff8b79190ef", size = 52664, upload-time = "2025-09-03T08:52:22.372Z" }, ] +[[package]] +name = "pyotp" +version = "2.9.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f3/b2/1d5994ba2acde054a443bd5e2d384175449c7d2b6d1a0614dbca3a63abfc/pyotp-2.9.0.tar.gz", hash = "sha256:346b6642e0dbdde3b4ff5a930b664ca82abfa116356ed48cc42c7d6590d36f63", size = 17763, upload-time = "2023-07-27T23:41:03.295Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c3/c0/c33c8792c3e50193ef55adb95c1c3c2786fe281123291c2dbf0eaab95a6f/pyotp-2.9.0-py3-none-any.whl", hash = "sha256:81c2e5865b8ac55e825b0358e496e1d9387c811e85bb40e71a3b29b288963612", size = 13376, upload-time = "2023-07-27T23:41:01.685Z" }, +] + [[package]] name = "pypandoc" version = "1.15" @@ -4880,6 +4893,15 @@ version = "0.48.9" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/c7/2c/94ed7b91db81d61d7096ac8f2d325ec562fc75e35f3baea8749c85b28784/PyPika-0.48.9.tar.gz", hash = "sha256:838836a61747e7c8380cd1b7ff638694b7a7335345d0f559b04b2cd832ad5378", size = 67259, upload-time = "2022-03-15T11:22:57.066Z" } +[[package]] +name = "pypng" +version = "0.20220715.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/93/cd/112f092ec27cca83e0516de0a3368dbd9128c187fb6b52aaaa7cde39c96d/pypng-0.20220715.0.tar.gz", hash = "sha256:739c433ba96f078315de54c0db975aee537cbc3e1d0ae4ed9aab0ca1e427e2c1", size = 128992, upload-time = "2022-07-15T14:11:05.301Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3e/b9/3766cc361d93edb2ce81e2e1f87dd98f314d7d513877a342d31b30741680/pypng-0.20220715.0-py3-none-any.whl", hash = "sha256:4a43e969b8f5aaafb2a415536c1a8ec7e341cd6a3f957fd5b5f32a4cfeed902c", size = 58057, upload-time = "2022-07-15T14:11:03.713Z" }, +] + [[package]] name = "pyproject-hooks" version = "1.2.0" @@ -5219,6 +5241,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/3a/fa/5abd82cde353f1009c068cca820195efd94e403d261b787e78ea7a9c8318/qdrant_client-1.9.0-py3-none-any.whl", hash = "sha256:ee02893eab1f642481b1ac1e38eb68ec30bab0f673bef7cc05c19fa5d2cbf43e", size = 229258, upload-time = "2024-04-22T13:35:46.81Z" }, ] +[[package]] +name = "qrcode" +version = "7.4.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "pypng" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/30/35/ad6d4c5a547fe9a5baf85a9edbafff93fc6394b014fab30595877305fa59/qrcode-7.4.2.tar.gz", hash = "sha256:9dd969454827e127dbd93696b20747239e6d540e082937c90f14ac95b30f5845", size = 535974, upload-time = "2023-02-05T22:11:46.548Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/24/79/aaf0c1c7214f2632badb2771d770b1500d3d7cbdf2590ae62e721ec50584/qrcode-7.4.2-py3-none-any.whl", hash = "sha256:581dca7a029bcb2deef5d01068e39093e80ef00b4a61098a2182eac59d01643a", size = 46197, upload-time = "2023-02-05T22:11:43.4Z" }, +] + [[package]] name = "rapidfuzz" version = "3.14.1" diff --git a/docker/docker-compose.test.yaml b/docker/docker-compose.test.yaml new file mode 100644 index 00000000000000..9feac2c309a6b1 --- /dev/null +++ b/docker/docker-compose.test.yaml @@ -0,0 +1,4 @@ +# Testing override - removes DB volume mount to avoid permission issues +services: + db: + volumes: [] diff --git a/docker/volumes/sandbox/dependencies/python-requirements.txt b/docker/volumes/sandbox/dependencies/python-requirements.txt index e69de29bb2d1d6..97ba7a98aed44e 100644 --- a/docker/volumes/sandbox/dependencies/python-requirements.txt +++ b/docker/volumes/sandbox/dependencies/python-requirements.txt @@ -0,0 +1,2 @@ +pyotp==2.9.0 +qrcode==7.4.2 \ No newline at end of file diff --git a/web/app/components/header/account-dropdown/index.tsx b/web/app/components/header/account-dropdown/index.tsx index ec0bc2f2665798..21eb97bbb30378 100644 --- a/web/app/components/header/account-dropdown/index.tsx +++ b/web/app/components/header/account-dropdown/index.tsx @@ -125,7 +125,9 @@ export default function AppSelector() {
setShowAccountSettingModal({ payload: 'members' })}> + )} onClick={() => { + setShowAccountSettingModal({ payload: 'members' }) + }}>
{t('common.userProfile.settings')}
diff --git a/web/app/components/header/account-setting/index.tsx b/web/app/components/header/account-setting/index.tsx index 8e71597e9caad1..ec749ecf580462 100644 --- a/web/app/components/header/account-setting/index.tsx +++ b/web/app/components/header/account-setting/index.tsx @@ -15,11 +15,14 @@ import { RiMoneyDollarCircleLine, RiPuzzle2Fill, RiPuzzle2Line, + RiShieldKeyholeFill, + RiShieldKeyholeLine, RiTranslate2, } from '@remixicon/react' import Button from '../../base/button' import MembersPage from './members-page' import LanguagePage from './language-page' +import MFAPage from './mfa-page' import ApiBasedExtensionPage from './api-based-extension-page' import DataSourcePage from './data-source-page-new' import ModelProviderPage from './model-provider-page' @@ -53,11 +56,14 @@ export default function AccountSetting({ onCancel, activeTab = 'members', }: IAccountSettingProps) { - const [activeMenu, setActiveMenu] = useState(activeTab) const { t } = useTranslation() const { enableBilling, enableReplaceWebAppLogo } = useProviderContext() const { isCurrentWorkspaceDatasetOperator } = useAppContext() + // Set appropriate default tab based on user role + const defaultTab = isCurrentWorkspaceDatasetOperator ? 'mfa' : activeTab + const [activeMenu, setActiveMenu] = useState(defaultTab) + const workplaceGroupItems = (() => { if (isCurrentWorkspaceDatasetOperator) return [] @@ -116,6 +122,12 @@ export default function AccountSetting({ key: 'account-group', name: t('common.settings.generalGroup'), items: [ + { + key: 'mfa', + name: t('common.settings.mfa'), + icon: , + activeIcon: , + }, { key: 'language', name: t('common.settings.language'), @@ -155,7 +167,7 @@ export default function AccountSetting({ { menuItems.map(menuItem => (
- {!isCurrentWorkspaceDatasetOperator && ( + {menuItem.items.length > 0 && (
{menuItem.name}
)}
@@ -219,6 +231,7 @@ export default function AccountSetting({ {activeMenu === 'data-source' && } {activeMenu === 'api-based-extension' && } {activeMenu === 'custom' && } + {activeMenu === 'mfa' && } {activeMenu === 'language' && }
diff --git a/web/app/components/header/account-setting/mfa-page.test.tsx b/web/app/components/header/account-setting/mfa-page.test.tsx new file mode 100644 index 00000000000000..3e8a26045833f8 --- /dev/null +++ b/web/app/components/header/account-setting/mfa-page.test.tsx @@ -0,0 +1,409 @@ +import React from 'react' +import { act, fireEvent, render, screen, waitFor } from '@testing-library/react' +import '@testing-library/jest-dom' + +// Mock the service base to avoid ky import issues +jest.mock('@/service/base', () => ({ + get: jest.fn(), + post: jest.fn(), + put: jest.fn(), + del: jest.fn(), +})) + +// Mock the translation hook +jest.mock('react-i18next', () => ({ + useTranslation: () => ({ + t: (key: string) => key, + }), +})) + +import MFAPage from './mfa-page' + +// Mock the Toast component +jest.mock('@/app/components/base/toast', () => ({ + __esModule: true, + default: { + notify: jest.fn(), + }, +})) + +// Mock Modal component +jest.mock('@/app/components/base/modal', () => ({ + __esModule: true, + default: ({ isShow, onClose, children }: any) => + isShow ?
{children}
: null, +})) + +import { QueryClient, QueryClientProvider } from '@tanstack/react-query' + +// Create a test wrapper component +const createWrapper = () => { + const queryClient = new QueryClient({ + defaultOptions: { + queries: { + retry: false, + gcTime: 0, + staleTime: 0, + }, + mutations: { + retry: false, + }, + }, + }) + + return ({ children }: { children: React.ReactNode }) => ( + + {children} + + ) +} + +describe('MFAPage Component', () => { + let wrapper: ReturnType + + beforeEach(() => { + jest.clearAllMocks() + wrapper = createWrapper() + }) + + test('renders loading state initially', () => { + const { get } = require('@/service/base') + get.mockImplementation(() => new Promise(() => { + // Never resolves - intentionally empty for testing loading state + })) // Never resolves + + const { container } = render(, { wrapper }) + + // Look for the loading spinner icon + const spinner = container.querySelector('.animate-spin') + expect(spinner).toBeInTheDocument() + }) + + test('renders enable button when MFA is disabled', async () => { + const { get } = require('@/service/base') + get.mockResolvedValue({ enabled: false }) + + render(, { wrapper }) + + await waitFor(() => { + expect(screen.getByText('mfa.enable')).toBeInTheDocument() + }) + }) + + test('renders disable button when MFA is enabled', async () => { + const { get } = require('@/service/base') + get.mockResolvedValue({ + enabled: true, + setup_at: '2025-01-01T12:00:00', + }) + + render(, { wrapper }) + + await waitFor(() => { + expect(screen.getByText('mfa.disable')).toBeInTheDocument() + }) + }) + + test('opens setup modal when enable button is clicked', async () => { + const { get, post } = require('@/service/base') + get.mockResolvedValue({ enabled: false }) + post.mockResolvedValue({ + secret: 'TEST_SECRET', + qr_code: 'data:image/png;base64,test', + }) + + render(, { wrapper }) + + await waitFor(() => { + expect(screen.getByText('mfa.enable')).toBeInTheDocument() + }) + + await act(async () => { + fireEvent.click(screen.getByText('mfa.enable')) + }) + + await waitFor(() => { + expect(screen.getByTestId('modal')).toBeInTheDocument() + }, { timeout: 10000 }) + + expect(screen.getByText('mfa.scanQRCode')).toBeInTheDocument() + }, 15000) + + test('completes MFA setup successfully', async () => { + const { get, post } = require('@/service/base') + const Toast = require('@/app/components/base/toast').default + + get.mockResolvedValue({ enabled: false }) + post.mockImplementation((url) => { + if (url.includes('/setup') && !url.includes('/complete')) { + return Promise.resolve({ + secret: 'TEST_SECRET', + qr_code: 'data:image/png;base64,test', + }) + } + else if (url.includes('/setup/complete')) { + return Promise.resolve({ + message: 'MFA setup successfully', + backup_codes: ['CODE1', 'CODE2', 'CODE3', 'CODE4', 'CODE5', 'CODE6', 'CODE7', 'CODE8'], + setup_at: '2025-01-01T12:00:00', + }) + } + }) + + render(, { wrapper }) + + // Wait for initial render + await waitFor(() => { + expect(screen.getByText('mfa.enable')).toBeInTheDocument() + }) + + // Click enable + fireEvent.click(screen.getByText('mfa.enable')) + + // Wait for QR code to be displayed + await waitFor(() => { + const qrCode = screen.queryByAltText('MFA QR Code') + expect(qrCode).toBeInTheDocument() + }, { timeout: 5000 }) + + // Click next button to go to verify step + fireEvent.click(screen.getByText('mfa.next')) + + // Wait for input field to appear + await waitFor(() => { + expect(screen.getByPlaceholderText('000000')).toBeInTheDocument() + }) + + // Enter TOTP code + const input = screen.getByPlaceholderText('000000') + fireEvent.change(input, { target: { value: '123456' } }) + + // Click verify button + const verifyButton = screen.getByRole('button', { name: /verify|mfa.verify/i }) + fireEvent.click(verifyButton) + + // Wait for backup codes to be displayed + await waitFor(() => { + expect(screen.getByText('mfa.backupCodesTitle')).toBeInTheDocument() + }) + + // Click done button + fireEvent.click(screen.getByText('mfa.done')) + + // Check that toast was called + expect(Toast.notify).toHaveBeenCalledWith({ + type: 'success', + message: 'mfa.setupSuccess', + }) + }, 15000) + + test('shows error when setup fails', async () => { + const { get, post } = require('@/service/base') + const Toast = require('@/app/components/base/toast').default + + get.mockResolvedValue({ enabled: false }) + post.mockImplementation((url) => { + if (url.includes('/setup') && !url.includes('/complete')) { + return Promise.resolve({ + secret: 'TEST_SECRET', + qr_code: 'data:image/png;base64,test', + }) + } + else if (url.includes('/setup/complete')) { + return Promise.reject(new Error('Invalid TOTP token')) + } + }) + + render(, { wrapper }) + + // Wait and click enable + await waitFor(() => { + expect(screen.getByText('mfa.enable')).toBeInTheDocument() + }) + + fireEvent.click(screen.getByText('mfa.enable')) + + // Wait for QR code + await waitFor(() => { + expect(screen.queryByAltText('MFA QR Code')).toBeInTheDocument() + }, { timeout: 5000 }) + + // Click next to go to verify step + fireEvent.click(screen.getByText('mfa.next')) + + // Wait for input + await waitFor(() => { + expect(screen.getByPlaceholderText('000000')).toBeInTheDocument() + }) + + // Enter wrong TOTP code + const input = screen.getByPlaceholderText('000000') + fireEvent.change(input, { target: { value: '000000' } }) + + // Click verify + const verifyButton = screen.getByRole('button', { name: /verify|mfa.verify/i }) + fireEvent.click(verifyButton) + + await waitFor(() => { + expect(Toast.notify).toHaveBeenCalledWith({ + type: 'error', + message: 'mfa.invalidToken', + }) + }, { timeout: 5000 }) + }, 15000) + + test('disables MFA successfully', async () => { + const { get, post } = require('@/service/base') + const Toast = require('@/app/components/base/toast').default + + get.mockResolvedValue({ + enabled: true, + setup_at: '2025-01-01T12:00:00', + }) + post.mockImplementation((url) => { + if (url.includes('/disable')) { + return Promise.resolve({ + success: true, + message: 'MFA disabled successfully', + }) + } + }) + + render(, { wrapper }) + + // Wait for disable button + await waitFor(() => { + expect(screen.getByText('mfa.disable')).toBeInTheDocument() + }) + + // Click disable + fireEvent.click(screen.getByText('mfa.disable')) + + // Modal should open + await waitFor(() => { + expect(screen.getByTestId('modal')).toBeInTheDocument() + }) + + // Find password input + const passwordInput = screen.getByPlaceholderText('common.account.password') + fireEvent.change(passwordInput, { target: { value: 'password123' } }) + + // Click disable button in modal + const disableButtons = screen.getAllByText('mfa.disable') + // The second one should be the button in the modal + fireEvent.click(disableButtons[1]) + + await waitFor(() => { + expect(Toast.notify).toHaveBeenCalledWith({ + type: 'success', + message: 'mfa.disabledSuccessfully', + }) + }, { timeout: 5000 }) + }, 15000) + + test('shows error when disable fails with wrong password', async () => { + const { get, post } = require('@/service/base') + const Toast = require('@/app/components/base/toast').default + + get.mockResolvedValue({ + enabled: true, + setup_at: '2025-01-01T12:00:00', + }) + post.mockImplementation((url) => { + if (url.includes('/disable')) + return Promise.reject(new Error('Invalid password')) + }) + + render(, { wrapper }) + + // Wait and click disable + await waitFor(() => { + expect(screen.getByText('mfa.disable')).toBeInTheDocument() + }) + + fireEvent.click(screen.getByText('mfa.disable')) + + // Wait for modal + await waitFor(() => { + expect(screen.getByTestId('modal')).toBeInTheDocument() + }) + + // Enter wrong password + const passwordInput = screen.getByPlaceholderText('common.account.password') + fireEvent.change(passwordInput, { target: { value: 'wrongpassword' } }) + + // Click disable button in modal + const disableButtons = screen.getAllByText('mfa.disable') + // The second one should be the button in the modal + fireEvent.click(disableButtons[1]) + + await waitFor(() => { + expect(Toast.notify).toHaveBeenCalledWith({ + type: 'error', + message: 'mfa.invalidPassword', + }) + }, { timeout: 5000 }) + }, 15000) + + test('handles backup codes display correctly', async () => { + const { get, post } = require('@/service/base') + + get.mockResolvedValue({ enabled: false }) + + // Mock immediate responses + post.mockImplementation((url) => { + if (url.includes('/setup') && !url.includes('/complete')) { + return Promise.resolve({ + secret: 'TEST_SECRET', + qr_code: 'data:image/png;base64,test', + }) + } + else if (url.includes('/setup/complete')) { + return Promise.resolve({ + message: 'MFA setup successfully', + backup_codes: ['ABCD1234', 'EFGH5678', 'IJKL9012', 'MNOP3456', 'QRST7890', 'UVWX1234', 'YZAB5678', 'CDEF9012'], + setup_at: '2025-01-01T12:00:00', + }) + } + }) + + render(, { wrapper }) + + // Wait for initial render + await waitFor(() => { + expect(screen.getByText('mfa.enable')).toBeInTheDocument() + }) + + // Setup MFA + fireEvent.click(screen.getByText('mfa.enable')) + + // Wait for QR code + await waitFor(() => { + const qrCode = screen.queryByAltText('MFA QR Code') + expect(qrCode).toBeInTheDocument() + }, { timeout: 10000 }) + + // Click next to go to verify step + fireEvent.click(screen.getByText('mfa.next')) + + // Wait for input + await waitFor(() => { + expect(screen.getByPlaceholderText('000000')).toBeInTheDocument() + }) + + // Enter TOTP code + const input = screen.getByPlaceholderText('000000') + fireEvent.change(input, { target: { value: '123456' } }) + + // Verify + const verifyButton = screen.getByRole('button', { name: /verify|mfa.verify/i }) + fireEvent.click(verifyButton) + + // Check backup codes are displayed + await waitFor(() => { + expect(screen.getByText('mfa.backupCodesTitle')).toBeInTheDocument() + expect(screen.getByText('ABCD1234')).toBeInTheDocument() + expect(screen.getByText('EFGH5678')).toBeInTheDocument() + }, { timeout: 5000 }) + }, 10000) // Increase test timeout +}) diff --git a/web/app/components/header/account-setting/mfa-page.tsx b/web/app/components/header/account-setting/mfa-page.tsx new file mode 100644 index 00000000000000..2a4d6b59bd443c --- /dev/null +++ b/web/app/components/header/account-setting/mfa-page.tsx @@ -0,0 +1,304 @@ +'use client' +import { useState } from 'react' +import { useTranslation } from 'react-i18next' +import { RiCheckboxCircleFill, RiLoader2Line, RiShieldKeyholeLine } from '@remixicon/react' +import Toast from '../../base/toast' +import Button from '../../base/button' +import Input from '../../base/input' +import Modal from '../../base/modal' +import { useMutation, useQuery, useQueryClient } from '@tanstack/react-query' +import { get, post } from '@/service/base' + +// API service functions +const mfaService = { + getStatus: async () => { + return get<{ + enabled: boolean + setup_at: string | null + }>('/account/mfa/status') + }, + + initSetup: async () => { + return post<{ + secret: string + qr_code: string + }>('/account/mfa/setup', { body: {} }) + }, + + completeSetup: async (totpToken: string) => { + return post<{ + message: string + backup_codes: string[] + setup_at: string + }>('/account/mfa/setup/complete', { + body: { totp_token: totpToken }, + }) + }, + + disable: async (password: string) => { + return post('/account/mfa/disable', { + body: { password }, + }) + }, +} + +export default function MFAPage() { + const { t } = useTranslation() + const queryClient = useQueryClient() + + // State + const [isSetupModalOpen, setIsSetupModalOpen] = useState(false) + const [isDisableModalOpen, setIsDisableModalOpen] = useState(false) + const [setupStep, setSetupStep] = useState<'qr' | 'verify' | 'backup'>('qr') + const [totpToken, setTotpToken] = useState('') + const [password, setPassword] = useState('') + const [qrData, setQrData] = useState<{ secret: string; qr_code: string } | null>(null) + const [backupCodes, setBackupCodes] = useState([]) + + // Query MFA status + const { data: mfaStatus, isLoading } = useQuery({ + queryKey: ['mfa-status'], + queryFn: mfaService.getStatus, + }) + + // Mutations + const initSetupMutation = useMutation({ + mutationFn: mfaService.initSetup, + onSuccess: (data) => { + setQrData(data) + setIsSetupModalOpen(true) + setSetupStep('qr') + }, + onError: () => { + Toast.notify({ type: 'error', message: t('common.somethingWentWrong') }) + }, + }) + + const completeSetupMutation = useMutation({ + mutationFn: ({ totpToken }: { totpToken: string }) => + mfaService.completeSetup(totpToken), + onSuccess: (data) => { + setBackupCodes(data.backup_codes) + setSetupStep('backup') + queryClient.invalidateQueries({ queryKey: ['mfa-status'] }) + }, + onError: () => { + Toast.notify({ type: 'error', message: t('mfa.invalidToken') }) + }, + }) + + const disableMutation = useMutation({ + mutationFn: mfaService.disable, + onSuccess: () => { + setIsDisableModalOpen(false) + queryClient.invalidateQueries({ queryKey: ['mfa-status'] }) + Toast.notify({ type: 'success', message: t('mfa.disabledSuccessfully') }) + }, + onError: () => { + Toast.notify({ type: 'error', message: t('mfa.invalidPassword') }) + }, + }) + + const handleSetupStart = () => { + initSetupMutation.mutate() + } + + const handleVerifyToken = () => { + if (totpToken.length !== 6) { + Toast.notify({ type: 'error', message: t('mfa.tokenLength') }) + return + } + completeSetupMutation.mutate({ totpToken }) + } + + const handleDisable = () => { + disableMutation.mutate(password) + } + + const handleCopyBackupCodes = () => { + const codesText = backupCodes.join('\n') + navigator.clipboard.writeText(codesText) + Toast.notify({ type: 'success', message: t('mfa.copied') }) + } + + if (isLoading) { + return ( +
+ +
+ ) + } + + return ( +
+
+
+ +
+
{t('mfa.description')}
+
+ {t('mfa.securityTip')} +
+
+ +
+
+
+
+ +
+
+
{t('mfa.authenticatorApp')}
+
{t('mfa.authenticatorDescription')}
+
+
+
+ {mfaStatus?.enabled && ( + + )} + +
+
+ + {mfaStatus?.enabled && mfaStatus?.setup_at && ( +
+ {t('mfa.enabledAt', { date: new Date(mfaStatus.setup_at).toLocaleDateString() })} +
+ )} +
+ + {/* Setup Modal */} + setIsSetupModalOpen(false)} + title={t('mfa.setupTitle')} + className="!max-w-md" + > + {setupStep === 'qr' && qrData && ( +
+

{t('mfa.scanQRCode')}

+
+ MFA QR Code +
+
+

{t('mfa.secretKey')}

+ {qrData.secret} +
+ +
+ )} + + {setupStep === 'verify' && ( +
+

{t('mfa.enterToken')}

+ setTotpToken(e.target.value)} + placeholder="000000" + maxLength={6} + className="text-center font-mono text-2xl" + /> + +
+ )} + + {setupStep === 'backup' && ( +
+
+

{t('mfa.backupCodesTitle')}

+

{t('mfa.backupCodesWarning')}

+
+
+
+ {backupCodes.map((code, index) => ( + {code} + ))} +
+
+
+ + +
+
+ )} + + + {/* Disable Modal */} + setIsDisableModalOpen(false)} + title={t('mfa.disableTitle')} + className="!max-w-md" + > +
+

{t('mfa.disableDescription')}

+ setPassword(e.target.value)} + placeholder={t('common.account.password')} + aria-label={t('mfa.enterYourPassword')} + /> +
+ + +
+
+ +
+ ) +} diff --git a/web/app/signin/components/mail-and-password-auth.tsx b/web/app/signin/components/mail-and-password-auth.tsx index aaadc0b197a321..32518fe553b81d 100644 --- a/web/app/signin/components/mail-and-password-auth.tsx +++ b/web/app/signin/components/mail-and-password-auth.tsx @@ -10,6 +10,7 @@ import { login } from '@/service/common' import Input from '@/app/components/base/input' import I18NContext from '@/context/i18n' import { noop } from 'lodash-es' +import MFAVerification from './mfa-verification' import { resolvePostLoginRedirect } from '../utils/post-login-redirect' import type { ResponseError } from '@/service/fetch' @@ -19,6 +20,8 @@ type MailAndPasswordAuthProps = { allowRegistration: boolean } +const passwordRegex = /^(?=.*[a-zA-Z])(?=.*\d).{8,}$/ + export default function MailAndPasswordAuth({ isInvite, isEmailSetup, allowRegistration }: MailAndPasswordAuthProps) { const { t } = useTranslation() const { locale } = useContext(I18NContext) @@ -28,6 +31,7 @@ export default function MailAndPasswordAuth({ isInvite, isEmailSetup, allowRegis const emailFromLink = decodeURIComponent(searchParams.get('email') || '') const [email, setEmail] = useState(emailFromLink) const [password, setPassword] = useState('') + const [showMFAVerification, setShowMFAVerification] = useState(false) const [isLoading, setIsLoading] = useState(false) const handleEmailPasswordLogin = async () => { @@ -61,7 +65,10 @@ export default function MailAndPasswordAuth({ isInvite, isEmailSetup, allowRegis url: '/login', body: loginData, }) - if (res.result === 'success') { + if (res.result === 'fail' && res.code === 'mfa_required') { + setShowMFAVerification(true) + } + else if (res.result === 'success') { if (isInvite) { router.replace(`/signin/invite-settings?${searchParams.toString()}`) } @@ -92,6 +99,18 @@ export default function MailAndPasswordAuth({ isInvite, isEmailSetup, allowRegis } } + if (showMFAVerification) { + return ( + + ) + } + return