infrasynth-backend-kit/infrasynth/security/views.py

import secrets
from datetime import datetime, timedelta

from django.contrib.auth import authenticate, get_user_model, login, logout
from django.utils import timezone
from rest_framework import mixins, status, viewsets
from rest_framework.decorators import action
from rest_framework.exceptions import (
    AuthenticationFailed,
    NotAuthenticated,
    NotFound,
    PermissionDenied,
    ValidationError,
)
from rest_framework.permissions import AllowAny, IsAuthenticated
from rest_framework.response import Response
from rest_framework_simplejwt.tokens import RefreshToken

from infrasynth.shared.exceptions import NotFoundError
from infrasynth.shared.settings_utils import get_setting
from infrasynth.tenancy.services import TenantService

from .altcha.services import ALTCHAService
from .models import APIKey, Grant, Revoke, Role, TwoFactorConfig
from .permissions import IsAuthenticatedAndPermitted
from .serializers import (
    APIKeySerializer,
    GrantSerializer,
    LoginSerializer,
    RevokeSerializer,
    RoleSerializer,
)
from .services import AuthorizationService
from .signals import (
    api_key_created,
    api_key_rotated,
    login_failed,
    two_factor_verified,
    user_logged_in,
    user_logged_out,
)
from .throttling import LoginAttemptGuard
from .two_factor.services import RecoveryCodeService, TOTPService
from .two_factor.utils import generate_pre_auth_token

UserModel = get_user_model()


class _AuthSupport:
    """Cookie/JWT/2FA helpers shared by the auth and 2FA viewsets."""

    @staticmethod
    def _mint_tokens(user, tenant, *, two_factor: bool = False):
        refresh = RefreshToken.for_user(user)
        if tenant is not None:
            refresh["tenant"] = str(tenant.pk)
        # ``2fa`` lets per-endpoint gates require a step-up second factor.
        refresh["2fa"] = bool(two_factor)
        return str(refresh.access_token), str(refresh)

    def _cookie_config(self):
        return {
            "secure": get_setting("INFRASYNTH_SECURITY", "COOKIE_SECURE", True),
            "httponly": get_setting("INFRASYNTH_SECURITY", "COOKIE_HTTPONLY", True),
            "samesite": get_setting("INFRASYNTH_SECURITY", "COOKIE_SAMESITE", "Lax"),
            "access_cookie": get_setting("INFRASYNTH_SECURITY", "ACCESS_COOKIE_NAME", "access_token"),
            "refresh_cookie": get_setting("INFRASYNTH_SECURITY", "REFRESH_COOKIE_NAME", "refresh_token"),
            "pre_auth_cookie": get_setting("INFRASYNTH_SECURITY", "PRE_AUTH_COOKIE_NAME", "pre_auth_token"),
        }

    def _encrypt(self, value):
        from infrasynth.shared.crypto import encrypt as crypto_encrypt

        try:
            return crypto_encrypt(value)
        except Exception:
            return value

    def _decrypt(self, value):
        from infrasynth.shared.crypto import decrypt as crypto_decrypt

        try:
            return crypto_decrypt(value)
        except Exception:
            return value

    def _set_auth_cookies(self, response, access_token, refresh_token):
        cfg = self._cookie_config()
        response.set_cookie(
            cfg["access_cookie"],
            self._encrypt(access_token),
            max_age=3600,
            secure=cfg["secure"],
            httponly=cfg["httponly"],
            samesite=cfg["samesite"],
            path="/",
        )
        response.set_cookie(
            cfg["refresh_cookie"],
            self._encrypt(refresh_token),
            max_age=86400 * 7,
            secure=cfg["secure"],
            httponly=cfg["httponly"],
            samesite=cfg["samesite"],
            path="/",
        )

    def _set_refresh_cookie(self, response, refresh_token):
        cfg = self._cookie_config()
        response.set_cookie(
            cfg["refresh_cookie"],
            self._encrypt(refresh_token),
            max_age=86400 * 7,
            secure=cfg["secure"],
            httponly=cfg["httponly"],
            samesite=cfg["samesite"],
            path="/",
        )

    def _user_from_refresh_cookie(self, request):
        """Returns ``(user, claims)`` from the refresh cookie, or ``(None, {})``."""
        raw = request.COOKIES.get(get_setting("INFRASYNTH_SECURITY", "REFRESH_COOKIE_NAME", "refresh_token"))
        if not raw:
            return None, {}
        try:
            refresh = RefreshToken(self._decrypt(raw))  # type: ignore[arg-type]
            user_id = refresh.get("user_id")
            claims = {"2fa": bool(refresh.get("2fa"))}
        except Exception:
            return None, {}
        return UserModel.objects.filter(pk=user_id).first(), claims

    def _complete_login(self, request, user, *, two_factor: bool = False):
        """Issues tenant-bound tokens or a workspace-picker challenge."""
        memberships = TenantService().get_active_memberships(user)
        if len(memberships) == 1:
            tenant = memberships[0].tenant
            access, refresh = self._mint_tokens(user, tenant, two_factor=two_factor)
            response = Response({"detail": "Login successful.", "tenant": str(tenant.pk)})
            self._set_auth_cookies(response, access, refresh)
            return response
        if len(memberships) > 1:
            _, refresh = self._mint_tokens(user, None, two_factor=two_factor)
            response = Response(
                {
                    "detail": "Select a workspace.",
                    "workspaces": [
                        {"id": str(m.tenant.pk), "slug": m.tenant.slug, "name": m.tenant.name} for m in memberships
                    ],
                }
            )
            self._set_refresh_cookie(response, refresh)
            return response
        raise PermissionDenied("No active workspace membership for this user.")

    @staticmethod
    def _requires_two_factor(user):
        return TwoFactorConfig.objects.filter(user=user, is_enabled=True, is_configured=True).first()

    def _start_pre_auth(self, request, user, config):
        """Stores the pre-auth session and cookie; returns the challenge response."""
        token = generate_pre_auth_token()
        request.session["_2fa_user_id"] = user.pk
        request.session["_2fa_pre_auth_token"] = token
        request.session["_2fa_started_at"] = timezone.now().isoformat()
        cfg = self._cookie_config()
        response = Response(
            {
                "detail": "Two-factor authentication required.",
                "requires2fa": True,
                "method": config.method,
            },
            status=status.HTTP_200_OK,
        )
        lifetime = int(get_setting("INFRASYNTH_SECURITY", "PRE_AUTH_TOKEN_LIFETIME_MINUTES", 5))
        response.set_cookie(
            cfg["pre_auth_cookie"],
            token,
            max_age=lifetime * 60,
            secure=cfg["secure"],
            httponly=cfg["httponly"],
            samesite=cfg["samesite"],
            path="/",
        )
        return response

    @staticmethod
    def _valid_pre_auth(request) -> bool:
        user_id = request.session.get("_2fa_user_id")
        token = request.session.get("_2fa_pre_auth_token")
        if not user_id or not token:
            return False
        started = request.session.get("_2fa_started_at")
        lifetime = int(get_setting("INFRASYNTH_SECURITY", "PRE_AUTH_TOKEN_LIFETIME_MINUTES", 5))
        if started:
            try:
                started_at = datetime.fromisoformat(started)
            except ValueError:
                return False
            if timezone.now() - started_at > timedelta(minutes=lifetime):
                return False
        return True

    def _finish_two_factor(self, request, user):
        login(request, user)
        request.session["_2fa_verified"] = True
        request.session.pop("_2fa_user_id", None)
        request.session.pop("_2fa_pre_auth_token", None)
        request.session.pop("_2fa_started_at", None)
        two_factor_verified.send(sender=self.__class__, user=user, method="totp")
        return self._complete_login(request, user, two_factor=True)


class AuthViewSet(_AuthSupport, viewsets.GenericViewSet):
    permission_classes = [AllowAny]

    @action(detail=False, methods=["post"])
    def login(self, request):
        serializer = LoginSerializer(data=request.data)
        serializer.is_valid(raise_exception=True)
        credential = serializer.validated_data["username"]
        guard = LoginAttemptGuard()
        guard.check(request, credential)

        user = authenticate(
            request,
            username=credential,
            password=serializer.validated_data["password"],
        )
        if not user:
            guard.record_failure(request, credential)
            login_failed.send(
                sender=self.__class__,
                credentials_key=credential,
                ip=request.META.get("REMOTE_ADDR"),
                reason="invalid_credentials",
            )
            raise AuthenticationFailed("Invalid credentials.")
        if not user.is_active:
            raise AuthenticationFailed("User is disabled.")

        guard.clear(request, credential)

        if get_setting("INFRASYNTH_SECURITY", "ALTCHA_PROTECT_LOGIN", False):
            altcha_ok = ALTCHAService().verify(
                request.data.get("altchaChallengeId") or request.data.get("altcha_challenge_id"),
                request.data.get("altchaSolution") or request.data.get("altcha_solution"),
                request.data.get("altchaNumber") or request.data.get("altcha_number"),
            )
            if not altcha_ok:
                raise ValidationError({"altcha": "A valid ALTCHA proof-of-work solution is required."})

        config = self._requires_two_factor(user)
        if config is not None:
            return self._start_pre_auth(request, user, config)

        login(request, user)
        response = self._complete_login(request, user)
        user_logged_in.send(
            sender=self.__class__,
            user=user,
            ip=request.META.get("REMOTE_ADDR"),
            user_agent=request.META.get("HTTP_USER_AGENT", ""),
        )
        return response

    @action(detail=False, methods=["post"], url_path="select-workspace")
    def select_workspace(self, request):
        user, claims = self._user_from_refresh_cookie(request)
        if user is None:
            raise AuthenticationFailed("No pending workspace selection.")
        tenant_id = request.data.get("tenantId") or request.data.get("tenant_id")
        if not tenant_id:
            raise ValidationError({"tenantId": "This field is required."})
        try:
            tenant = TenantService().select_tenant(user, tenant_id)
        except NotFoundError as exc:
            raise NotFound(str(exc)) from exc
        access, refresh = self._mint_tokens(user, tenant, two_factor=claims.get("2fa", False))
        response = Response({"detail": "Workspace selected.", "tenant": str(tenant.pk)})
        self._set_auth_cookies(response, access, refresh)
        return response

    @action(detail=False, methods=["post"], url_path="switch-workspace")
    def switch_workspace(self, request):
        if not request.user or not request.user.is_authenticated:
            raise NotAuthenticated("Authentication required.")
        if not get_setting("INFRASYNTH_SECURITY", "ENABLE_WORKSPACE_SWITCHING", True):
            raise NotFound("Workspace switching is disabled.")
        tenant_id = request.data.get("tenantId") or request.data.get("tenant_id")
        if not tenant_id:
            raise ValidationError({"tenantId": "This field is required."})
        from infrasynth.tenancy.context import get_current_tenant

        previous = get_current_tenant()
        try:
            tenant = TenantService().switch_tenant(
                request.user,
                tenant_id,
                previous_tenant_id=getattr(previous, "pk", None),
            )
        except NotFoundError as exc:
            raise NotFound(str(exc)) from exc
        auth = getattr(request, "auth", None)
        two_factor = (
            bool(auth.get("2fa"))
            if auth is not None and hasattr(auth, "get")
            else bool(getattr(request, "session", None) and request.session.get("_2fa_verified"))
        )
        access, refresh = self._mint_tokens(request.user, tenant, two_factor=two_factor)
        response = Response({"detail": "Workspace switched.", "tenant": str(tenant.pk)})
        self._set_auth_cookies(response, access, refresh)
        return response

    @action(detail=False, methods=["post"])
    def logout(self, request):
        logout(request)
        cfg = self._cookie_config()
        response = Response({"detail": "Logout successful."})
        response.delete_cookie(cfg["access_cookie"], path="/")
        response.delete_cookie(cfg["refresh_cookie"], path="/")
        response.delete_cookie(cfg["pre_auth_cookie"], path="/")
        user_logged_out.send(sender=self.__class__, user=getattr(request, "user", None))
        return response

    @action(detail=False, methods=["post"])
    def refresh(self, request):
        cfg = self._cookie_config()
        raw_refresh = request.COOKIES.get(cfg["refresh_cookie"])
        if not raw_refresh:
            raise AuthenticationFailed("Refresh token not found.")

        try:
            refresh = RefreshToken(self._decrypt(raw_refresh))  # type: ignore[arg-type]
            access_token = str(refresh.access_token)
            if get_setting("INFRASYNTH_SECURITY", "ROTATE_REFRESH_TOKENS", True):
                new_refresh = str(refresh)
                if get_setting("INFRASYNTH_SECURITY", "BLACKLIST_AFTER_ROTATION", True):
                    from rest_framework_simplejwt.exceptions import TokenError

                    try:
                        refresh.blacklist()
                    except (AttributeError, TokenError):  # blacklist app not installed
                        pass
            else:
                new_refresh = raw_refresh
        except Exception:
            raise AuthenticationFailed("Invalid or expired refresh token.")

        response = Response({"detail": "Token refreshed."})
        self._set_auth_cookies(response, access_token, new_refresh)
        return response

    @action(detail=False, methods=["get"])
    def check(self, request):
        if not request.user or not request.user.is_authenticated:
            return Response({"authenticated": False}, status=status.HTTP_401_UNAUTHORIZED)
        authz = AuthorizationService()
        effective_permissions = list(authz.get_effective_permissions(request.user))
        user = request.user
        from infrasynth.tenancy.context import get_current_tenant

        tenant = get_current_tenant()
        return Response(
            {
                "id": user.pk,
                "email": getattr(user, "email", ""),
                "name": getattr(user, "get_full_name", lambda: str(user))(),
                "effective_permissions": effective_permissions,
                "tenant": str(tenant.pk) if tenant is not None else None,
            }
        )


class TwoFactorViewSet(_AuthSupport, viewsets.GenericViewSet):
    permission_classes = [IsAuthenticated]

    def get_permissions(self):
        # ``verify``/``recovery`` are the pre-auth (not-yet-logged-in) steps; the
        # router-less ``as_view`` mapping does not apply per-action overrides.
        if self.action in ("verify", "recovery"):
            return [AllowAny()]
        return [IsAuthenticated()]

    @action(detail=False, methods=["post"])
    def setup(self, request):  # type: ignore[override]
        totp = TOTPService()
        secret = totp.generate_secret()
        qr_base64 = totp.generate_qr_base64(secret, request.user.email)
        provisioning_uri = totp.get_provisioning_uri(secret, request.user.email)
        request.session["_2fa_setup_secret"] = secret
        return Response(
            {
                "secret": secret,
                "qr_base64": qr_base64,
                "provisioning_uri": provisioning_uri,
            }
        )

    @action(detail=False, methods=["post"])
    def verify_setup(self, request):
        secret = request.session.get("_2fa_setup_secret")
        if not secret:
            raise AuthenticationFailed("No setup in progress.")
        code = request.data.get("code")
        if not code:
            raise AuthenticationFailed("Code is required.")
        totp = TOTPService()
        if not totp.verify(secret, code):
            raise AuthenticationFailed("Invalid code.")
        rcs = RecoveryCodeService()
        recovery_codes = rcs.generate_codes()
        config, _ = TwoFactorConfig.objects.get_or_create(user=request.user)
        config.is_enabled = True
        config.is_configured = True
        config.secret_key_encrypted = totp.encrypt_secret(secret)
        config.recovery_codes_encrypted = rcs.encrypt_codes(recovery_codes)
        config.save()
        del request.session["_2fa_setup_secret"]
        return Response({"recovery_codes": recovery_codes})

    @action(detail=False, methods=["post"], permission_classes=[AllowAny])
    def verify(self, request):
        if not self._valid_pre_auth(request):
            raise AuthenticationFailed("No pre-auth session.")
        user_id = request.session["_2fa_user_id"]
        try:
            config = TwoFactorConfig.objects.get(user_id=user_id, is_enabled=True)
        except TwoFactorConfig.DoesNotExist:
            raise AuthenticationFailed("2FA not configured.")
        code = request.data.get("code")
        if not code:
            raise AuthenticationFailed("Code is required.")
        totp = TOTPService()
        secret = totp.decrypt_secret(config.secret_key_encrypted)
        if not totp.verify(secret, code):
            raise AuthenticationFailed("Invalid code.")
        return self._finish_two_factor(request, UserModel.objects.get(pk=user_id))

    @action(detail=False, methods=["post"])
    def disable(self, request):
        try:
            config = TwoFactorConfig.objects.get(user=request.user)
            config.is_enabled = False
            config.is_configured = False
            config.secret_key_encrypted = None
            config.recovery_codes_encrypted = None
            config.save()
        except TwoFactorConfig.DoesNotExist:
            pass
        return Response({"detail": "2FA disabled."})

    @action(detail=False, methods=["post"], permission_classes=[AllowAny])
    def recovery(self, request):
        recovery_code = request.data.get("recovery_code")
        if not recovery_code:
            raise AuthenticationFailed("Recovery code is required.")
        if not self._valid_pre_auth(request):
            raise AuthenticationFailed("No pre-auth session.")
        user_id = request.session["_2fa_user_id"]
        try:
            config = TwoFactorConfig.objects.get(user_id=user_id, is_enabled=True)
        except TwoFactorConfig.DoesNotExist:
            raise AuthenticationFailed("2FA not configured.")
        rcs = RecoveryCodeService()
        if not rcs.verify_code(recovery_code, config.recovery_codes_encrypted):
            raise AuthenticationFailed("Invalid recovery code.")
        updated = rcs.remove_used_code(recovery_code, config.recovery_codes_encrypted)
        if updated:
            config.recovery_codes_encrypted = updated
            config.save(update_fields=["recovery_codes_encrypted"])
        return self._finish_two_factor(request, UserModel.objects.get(pk=user_id))


class ALTCHAViewSet(viewsets.GenericViewSet):
    permission_classes = [AllowAny]

    @action(detail=False, methods=["post"])
    def challenge(self, request):
        return Response(ALTCHAService().create_challenge())

    @action(detail=False, methods=["post"])
    def verify(self, request):
        result = ALTCHAService().verify(
            request.data.get("challenge_id"),
            request.data.get("solution"),
            request.data.get("number"),
        )
        return Response({"verified": result})


class APIKeyViewSet(viewsets.ModelViewSet):
    queryset = APIKey.objects.order_by("-id").all()
    serializer_class = APIKeySerializer
    permission_classes = [IsAuthenticatedAndPermitted]
    required_permissions = ["security.manage_api_keys"]

    def get_queryset(self):
        return APIKey.objects.order_by("-id").all()

    @staticmethod
    def _generate_key():
        from django.contrib.auth.hashers import make_password

        prefix_chars = int(get_setting("INFRASYNTH_SECURITY", "API_KEY_PREFIX_LENGTH", 8))
        hasher = get_setting("INFRASYNTH_SECURITY", "API_KEY_HASH_ALGORITHM", "pbkdf2_sha256")
        prefix = secrets.token_hex(max(1, prefix_chars // 2))
        secret = secrets.token_hex(32)
        try:
            key_hash = make_password(secret, hasher=hasher)
        except ValueError:  # hasher not registered in this project — use the default
            key_hash = make_password(secret)
        return prefix, secret, f"{prefix}.{secret}", key_hash

    def _default_expiry(self):
        days = int(get_setting("INFRASYNTH_SECURITY", "API_KEY_DEFAULT_EXPIRY_DAYS", 365))
        if days <= 0:
            return None
        return timezone.now() + timedelta(days=days)

    def perform_create(self, serializer):
        prefix, _secret, full_key, key_hash = self._generate_key()
        instance = serializer.save(
            prefix=prefix,
            key_hash=key_hash,
            created_by=self.request.user,
            expires_at=serializer.validated_data.get("expires_at") or self._default_expiry(),
        )
        instance.key = full_key
        api_key_created.send(sender=APIKey, key_name=instance.name, created_by=self.request.user)
        return instance

    @action(detail=True, methods=["post"])
    def rotate(self, request, pk=None):
        """Issues a new secret for the key and deactivates the old one."""
        old = self.get_object()
        prefix, _secret, full_key, key_hash = self._generate_key()
        new = APIKey.all_objects.create(
            tenant=old.tenant,
            name=old.name,
            prefix=prefix,
            key_hash=key_hash,
            scopes=old.scopes,
            created_by=request.user,
            expires_at=old.expires_at or self._default_expiry(),
            rotated_from=old,
        )
        old.is_active = False
        old.save(update_fields=["is_active"])
        api_key_rotated.send(sender=APIKey, key_name=new.name, rotated_by=request.user)
        data = dict(APIKeySerializer(new).data)
        data["key"] = full_key
        return Response(data, status=status.HTTP_201_CREATED)


class RoleViewSet(viewsets.ModelViewSet):
    queryset = Role.objects.order_by("name").all()
    serializer_class = RoleSerializer
    permission_classes = [IsAuthenticatedAndPermitted]
    required_permissions = ["security.manage_roles"]
    lookup_field = "slug"

    def get_queryset(self):
        return Role.objects.order_by("name").all()

    def perform_destroy(self, instance):
        if instance.is_system:
            raise PermissionDenied("System roles cannot be deleted.")
        instance.delete()


class GrantViewSet(
    mixins.CreateModelMixin,
    mixins.ListModelMixin,
    mixins.DestroyModelMixin,
    viewsets.GenericViewSet,
):
    queryset = Grant.objects.all()
    serializer_class = GrantSerializer
    permission_classes = [IsAuthenticatedAndPermitted]
    required_permissions = ["security.manage_grants"]

    def get_queryset(self):
        return Grant.objects.select_related("user", "granted_by").order_by("-id").all()

    def perform_create(self, serializer):
        serializer.save(granted_by=self.request.user)


class RevokeViewSet(
    mixins.CreateModelMixin,
    mixins.ListModelMixin,
    mixins.DestroyModelMixin,
    viewsets.GenericViewSet,
):
    queryset = Revoke.objects.all()
    serializer_class = RevokeSerializer
    permission_classes = [IsAuthenticatedAndPermitted]
    required_permissions = ["security.manage_grants"]

    def get_queryset(self):
        return Revoke.objects.select_related("user", "revoked_by").order_by("-id").all()

    def perform_create(self, serializer):
        serializer.save(revoked_by=self.request.user)


class UserPermissionViewSet(viewsets.GenericViewSet):
    """Effective permissions and role assignment for a user (``/auth/users/<id>/…``)."""

    permission_classes = [IsAuthenticatedAndPermitted]
    queryset = UserModel.objects.all()

    def initial(self, request, *args, **kwargs):
        if request.method in ("PUT", "PATCH", "POST"):
            self.required_permissions = ["security.manage_roles"]
        else:
            self.required_permissions = ["security.view_permissions"]
        super().initial(request, *args, **kwargs)

    def _get_user(self, pk):
        user = UserModel.objects.filter(pk=pk).first()
        if user is None:
            raise NotFound("User not found.")
        return user

    @action(detail=True, methods=["get"], url_path="permissions")
    def permissions(self, request, pk=None):
        user = self._get_user(pk)
        perms = sorted(AuthorizationService().get_effective_permissions(user))
        return Response({"user": user.pk, "permissions": perms})

    @action(detail=True, methods=["get", "put"], url_path="roles")
    def roles(self, request, pk=None):
        user = self._get_user(pk)
        if request.method == "PUT":
            slugs = request.data.get("roles", [])
            if not isinstance(slugs, list):
                raise ValidationError({"roles": "Expected a list of role slugs."})
            resolved = list(Role.objects.filter(slug__in=slugs))
            missing = set(slugs) - {role.slug for role in resolved}
            if missing:
                raise ValidationError({"roles": f"Unknown role(s): {', '.join(sorted(missing))}"})
            user.roles.set(resolved)
        return Response({"user": user.pk, "roles": list(user.roles.values_list("slug", flat=True))})