| 1 | # Create new file: app/middleware/customer_access.py |
| 2 | from typing import List |
| 3 | from typing import Optional |
| 4 | |
| 5 | from fastapi import Depends |
| 6 | from fastapi import HTTPException |
| 7 | from sqlalchemy import select |
| 8 | from sqlalchemy.ext.asyncio import AsyncSession |
| 9 | |
| 10 | from app.auth.models.users import RoleEnum |
| 11 | from app.auth.models.users import User |
| 12 | from app.auth.models.users import UserCustomerAccess |
| 13 | from app.auth.utils import AuthHandler |
| 14 | from app.db.db_session import get_db |
| 15 | |
| 16 | |
| 17 | class CustomerAccessHandler: |
| 18 | async def get_user_accessible_customers(self, user: User, session: AsyncSession) -> List[str]: |
| 19 | """Get all customer codes accessible to a user""" |
| 20 | # Admin and analyst users have access to all customers |
| 21 | if user.role_id in [RoleEnum.admin, RoleEnum.analyst]: |
| 22 | return ["*"] # Wildcard for all customers |
| 23 | |
| 24 | # Customer users only see their assigned customers |
| 25 | if user.role_id == RoleEnum.customer_user: |
| 26 | result = await session.execute(select(UserCustomerAccess.customer_code).where(UserCustomerAccess.user_id == user.id)) |
| 27 | return result.scalars().all() |
| 28 | |
| 29 | return [] # No access by default |
| 30 | |
| 31 | async def check_customer_access(self, user: User, customer_code: str, session: AsyncSession) -> bool: |
| 32 | """Check if user has access to specific customer""" |
| 33 | accessible_customers = await self.get_user_accessible_customers(user, session) |
| 34 | |
| 35 | # Wildcard access (admin/analyst) |
| 36 | if "*" in accessible_customers: |
| 37 | return True |
| 38 | |
| 39 | # Specific customer access |
| 40 | return customer_code in accessible_customers |
| 41 | |
| 42 | async def resolve_effective_customers( |
| 43 | self, |
| 44 | user: User, |
| 45 | requested_customers: Optional[List[str]], |
| 46 | session: AsyncSession, |
| 47 | ) -> List[str]: |
| 48 | """Resolve the customer codes a query should be filtered to. |
| 49 | |
| 50 | Combines the user's *accessible* customers with an optional *requested* |
| 51 | subset (e.g. a portal customer filter), so a caller can narrow the view |
| 52 | without ever escaping their own access scope. |
| 53 | |
| 54 | Returns either: |
| 55 | - ``["*"]`` — no filtering needed (wildcard access and no requested subset), or |
| 56 | - a concrete list of customer codes to filter on. An empty list means the |
| 57 | requested subset resolved to nothing the user may see, and callers should |
| 58 | treat it as "match no rows" (``column.in_([])``). |
| 59 | """ |
| 60 | accessible_customers = await self.get_user_accessible_customers(user, session) |
| 61 | |
| 62 | # No subset requested -> preserve existing behaviour (may be ["*"]). |
| 63 | if not requested_customers: |
| 64 | return accessible_customers |
| 65 | |
| 66 | # Wildcard access (admin/analyst): any requested subset is allowed as-is. |
| 67 | if "*" in accessible_customers: |
| 68 | return list(requested_customers) |
| 69 | |
| 70 | # Scoped user: only honor requested codes they actually have access to. |
| 71 | return [code for code in requested_customers if code in accessible_customers] |
| 72 | |
| 73 | async def filter_query_by_customer_access( |
| 74 | self, |
| 75 | user: User, |
| 76 | session: AsyncSession, |
| 77 | base_query, |
| 78 | customer_code_field, |
| 79 | requested_customers: Optional[List[str]] = None, |
| 80 | ): |
| 81 | """Filter any query by user's customer access. |
| 82 | |
| 83 | When ``requested_customers`` is provided, the query is further narrowed to |
| 84 | that subset (intersected with the user's access — see |
| 85 | ``resolve_effective_customers``). |
| 86 | """ |
| 87 | accessible_customers = await self.resolve_effective_customers(user, requested_customers, session) |
| 88 | |
| 89 | # Admin/analyst see everything (no subset requested) |
| 90 | if "*" in accessible_customers: |
| 91 | return base_query |
| 92 | |
| 93 | # Customer users (or anyone with a requested subset) see only matching data |
| 94 | if accessible_customers: |
| 95 | return base_query.where(customer_code_field.in_(accessible_customers)) |
| 96 | |
| 97 | # No access / requested subset resolved to nothing - return empty result |
| 98 | return base_query.where(False) |
| 99 | |
| 100 | def require_customer_access(self, customer_code: Optional[str] = None): |
| 101 | """FastAPI dependency to enforce customer access""" |
| 102 | |
| 103 | async def _check_access(current_user: User = Depends(AuthHandler().get_current_user), session: AsyncSession = Depends(get_db)): |
| 104 | if customer_code: |
| 105 | if not await self.check_customer_access(current_user, customer_code, session): |
| 106 | raise HTTPException(status_code=403, detail=f"Access denied to customer {customer_code}") |
| 107 | return current_user |
| 108 | |
| 109 | return _check_access |
| 110 | |
| 111 | |
| 112 | # Create a singleton instance |
| 113 | customer_access_handler = CustomerAccessHandler() |