feat(mfa/plugin): add dispatch mfa plugin (#5175)

* feat(mfa/plugin): add dispatch mfa plugin

* feat(mfa/plugin): add dispatch mfa plugin

* feat(mfa/plugin): add dispatch mfa plugin

* feat: improve confirmation modal

* Update src/dispatch/plugins/dispatch_core/plugin.py

Co-authored-by: Marc Vilanova <39573146+mvilanova@users.noreply.github.com>

---------

Co-authored-by: Marc Vilanova <39573146+mvilanova@users.noreply.github.com>

setup.py

@@ -311,9 +311,7 @@ def _build_static(self):
         env["DISPATCH_STATIC_DIST_PATH"] = self.dispatch_static_dist_path
         env["NODE_ENV"] = "production"
         # TODO: Our JS builds should not require 4GB heap space
-        env["NODE_OPTIONS"] = (
-            (env.get("NODE_OPTIONS", "") + " --max-old-space-size=4096")
-        ).lstrip()
+        env["NODE_OPTIONS"] = (env.get("NODE_OPTIONS", "") + " --max-old-space-size=4096").lstrip()
         # self._run_npm_command(["webpack", "--bail"], env=env)
 
     def _write_version_file(self, version_info):
@@ -405,6 +403,7 @@ def run(self):
             "dispatch_atlassian_confluence = dispatch.plugins.dispatch_atlassian_confluence.plugin:ConfluencePagePlugin",
             "dispatch_atlassian_confluence_document = dispatch.plugins.dispatch_atlassian_confluence.docs.plugin:ConfluencePageDocPlugin",
             "dispatch_aws_sqs = dispatch.plugins.dispatch_aws.plugin:AWSSQSSignalConsumerPlugin",
+            "dispatch_auth_mfa = dispatch.plugins.dispatch_core.plugin:DispatchMfaPlugin",
             "dispatch_basic_auth = dispatch.plugins.dispatch_core.plugin:BasicAuthProviderPlugin",
             "dispatch_contact = dispatch.plugins.dispatch_core.plugin:DispatchContactPlugin",
             "dispatch_document_resolver = dispatch.plugins.dispatch_core.plugin:DispatchDocumentResolverPlugin",

src/dispatch/auth/models.py

@@ -2,6 +2,7 @@
 import secrets
 from typing import List
 from datetime import datetime, timedelta
+from uuid import uuid4
 
 import bcrypt
 from jose import jwt
@@ -10,6 +11,7 @@
 from pydantic.networks import EmailStr
 
 from sqlalchemy import DateTime, Column, String, LargeBinary, Integer, Boolean
+from sqlalchemy.dialects.postgresql import UUID
 from sqlalchemy.orm import relationship
 from sqlalchemy.sql.schema import ForeignKey
 from sqlalchemy_utils import TSVectorType
@@ -20,7 +22,7 @@
     DISPATCH_JWT_EXP,
 )
 from dispatch.database.core import Base
-from dispatch.enums import UserRoles
+from dispatch.enums import DispatchEnum, UserRoles
 from dispatch.models import OrganizationSlug, PrimaryKey, TimeStampMixin, DispatchBase, Pagination
 from dispatch.organization.models import Organization, OrganizationRead
 from dispatch.project.models import Project, ProjectRead
@@ -192,3 +194,31 @@ class UserRegisterResponse(DispatchBase):
 
 class UserPagination(Pagination):
     items: List[UserRead] = []
+
+
+class MfaChallengeStatus(DispatchEnum):
+    PENDING = "pending"
+    APPROVED = "approved"
+    DENIED = "denied"
+    EXPIRED = "expired"
+
+
+class MfaChallenge(Base, TimeStampMixin):
+    id = Column(Integer, primary_key=True, autoincrement=True)
+    valid = Column(Boolean, default=False)
+    reason = Column(String, nullable=True)
+    action = Column(String)
+    status = Column(String, default=MfaChallengeStatus.PENDING)
+    challenge_id = Column(UUID(as_uuid=True), default=uuid4, unique=True)
+    dispatch_user_id = Column(Integer, ForeignKey(DispatchUser.id), nullable=False)
+    dispatch_user = relationship(DispatchUser, backref="mfa_challenges")
+
+
+class MfaPayloadResponse(DispatchBase):
+    status: str
+
+
+class MfaPayload(DispatchBase):
+    action: str
+    project_id: int
+    challenge_id: str

src/dispatch/auth/views.py

@@ -1,3 +1,5 @@
+import logging
+
 from fastapi import APIRouter, Depends, HTTPException, status
 from pydantic.error_wrappers import ErrorWrapper, ValidationError
 
@@ -17,9 +19,13 @@
 from dispatch.database.service import CommonParameters, search_filter_sort_paginate
 from dispatch.enums import UserRoles
 from dispatch.models import OrganizationSlug, PrimaryKey
+from dispatch.plugin import service as plugin_service
+from dispatch.plugins.dispatch_core.exceptions import MfaException
 from dispatch.organization.models import OrganizationRead
 
 from .models import (
+    MfaPayload,
+    MfaPayloadResponse,
     UserLogin,
     UserLoginResponse,
     UserOrganization,
@@ -33,6 +39,8 @@
 from .service import get, get_by_email, update, create
 
 
+log = logging.getLogger(__name__)
+
 auth_router = APIRouter()
 user_router = APIRouter()
 
@@ -246,6 +254,49 @@ def register_user(
     return user
 
 
+@auth_router.post("/mfa", response_model=MfaPayloadResponse)
+def mfa_check(
+    payload_in: MfaPayload,
+    current_user: CurrentUser,
+    db_session: DbSession,
+):
+    log.info(f"MFA check initiated for user: {current_user.email}")
+    log.debug(f"Payload received: {payload_in.dict()}")
+
+    try:
+        log.info(f"Attempting to get active MFA plugin for project: {payload_in.project_id}")
+        mfa_auth_plugin = plugin_service.get_active_instance(
+            db_session=db_session, project_id=payload_in.project_id, plugin_type="auth-mfa"
+        )
+
+        if not mfa_auth_plugin:
+            log.error(f"MFA plugin not enabled for project: {payload_in.project_id}")
+            raise HTTPException(
+                status_code=400, detail="MFA plugin is not enabled for the project."
+            )
+
+        log.info(f"MFA plugin found: {mfa_auth_plugin.__class__.__name__}")
+
+        log.info("Validating MFA token")
+        status = mfa_auth_plugin.instance.validate_mfa_token(payload_in, current_user, db_session)
+
+        log.info("MFA token validation successful")
+        return MfaPayloadResponse(status=status)
+
+    except MfaException as e:
+        log.error(f"MFA Exception occurred: {str(e)}")
+        log.debug(f"MFA Exception details: {type(e).__name__}", exc_info=True)
+        raise HTTPException(status_code=400, detail=str(e)) from e
+
+    except Exception as e:
+        log.critical(f"Unexpected error in MFA check: {str(e)}")
+        log.exception("Full traceback:")
+        raise HTTPException(status_code=500, detail="An unexpected error occurred") from e
+
+    finally:
+        log.info("MFA check completed")
+
+
 if DISPATCH_AUTH_REGISTRATION_ENABLED:
     register_user = auth_router.post("/register", response_model=UserRegisterResponse)(
         register_user

src/dispatch/database/revisions/tenant/versions/2024-09-06_51eacaf1f62c.py

@@ -0,0 +1,45 @@
+"""Adds mfa_challenge to track challenges against core Dispatch MFA plugin.
+
+Revision ID: 51eacaf1f62c
+Revises: 71cd7ed999c4
+Create Date: 2024-08-09 12:59:54.631968
+
+"""
+from alembic import op
+import sqlalchemy as sa
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision = "51eacaf1f62c"
+down_revision = "d6b3853be8e4"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table(
+        "mfa_challenge",
+        sa.Column("id", sa.Integer(), nullable=False, autoincrement=True),
+        sa.Column("valid", sa.Boolean(), nullable=True, server_default=sa.text("true")),
+        sa.Column("reason", sa.String(), nullable=True),
+        sa.Column("action", sa.String(), nullable=True),
+        sa.Column("challenge_id", postgresql.UUID(as_uuid=True), nullable=True),
+        sa.Column("dispatch_user_id", sa.Integer(), nullable=False),
+        sa.Column("status", sa.String(), nullable=True),
+        sa.Column("updated_at", sa.DateTime(), nullable=True),
+        sa.Column("created_at", sa.DateTime(), nullable=True),
+        sa.ForeignKeyConstraint(
+            ["dispatch_user_id"],
+            ["dispatch_core.dispatch_user.id"],
+        ),
+        sa.PrimaryKeyConstraint("id"),
+        sa.UniqueConstraint("challenge_id"),
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_table("mfa_challenge")
+    # ### end Alembic commands ###

src/dispatch/plugins/bases/auth_mfa.py

@@ -13,3 +13,9 @@ class MultiFactorAuthenticationPlugin(Plugin):
 
     def send_push_notification(self, items, **kwargs):
         raise NotImplementedError
+
+    def validate_mfa(self, items, **kwargs):
+        raise NotImplementedError
+
+    def create_mfa_challenge(self, items, **kwargs):
+        raise NotImplementedError

src/dispatch/plugins/dispatch_core/exceptions.py

@@ -0,0 +1,34 @@
+class MfaException(Exception):
+    """Base exception for MFA-related errors."""
+
+    pass
+
+
+class InvalidChallengeError(MfaException):
+    """Raised when the challenge is invalid."""
+
+    pass
+
+
+class UserMismatchError(MfaException):
+    """Raised when the challenge doesn't belong to the current user."""
+
+    pass
+
+
+class ActionMismatchError(MfaException):
+    """Raised when the action doesn't match the challenge."""
+
+    pass
+
+
+class ExpiredChallengeError(MfaException):
+    """Raised when the challenge is no longer valid."""
+
+    pass
+
+
+class InvalidChallengeStateError(MfaException):
+    """Raised when the challenge is in an invalid state."""
+
+    pass