add extra endpoints and develop security with adding permissions to endpoints

Author: RustamovAkrom

src/api/v1/user.py

@@ -1,5 +1,7 @@
 from fastapi import APIRouter, Depends
 
+from db.dependencies.auth import get_current_user, require_roles
+from db.models.users import User, UserRole
 from schemas.auth import (
     LoginOutScheme,
     LoginScheme,
@@ -12,6 +14,26 @@
 router = APIRouter()
 
 
+@router.get("/me")
+async def get_me(current_user: User = Depends(get_current_user)):
+    return {
+        "id": current_user.id,
+        "email": current_user.email,
+        "username": current_user.username,
+        "role": current_user.role,
+    }
+
+
+@router.get("/dashboard")
+async def dashboard(current_user: User = Depends(get_current_user)):
+    return {"message": f"Welcome, {current_user.username}"}
+
+
+@router.get("/admin/stats")
+async def admin_stats(admin_user: User = Depends(require_roles(UserRole.admin.value))):
+    return {"status": "ok", "admin": admin_user.username}
+
+
 @router.post("/register", response_model=UserCreateScheme)
 async def register_user(
     data: RegistrationScheme,

src/core/security.py

@@ -1,46 +1,109 @@
+import uuid
 from datetime import datetime, timedelta
+from typing import Any, Dict, Optional, Union
 
-from jose import jwt
+from jose import ExpiredSignatureError, JWTError, jwt
 from passlib.context import CryptContext
 
 from core.config import settings
 
 pwd_context = CryptContext(schemes=["argon2"], deprecated="auto")
 
 
+# Hashing helpers
 def hash_password(password: str) -> str:
     return pwd_context.hash(password)
 
 
+# Verifying helpers
 def verify_password(plain_password: str, hashed_password: str) -> bool:
     return pwd_context.verify(plain_password, hashed_password)
 
 
-def create_access_token(data: dict, expires_delta: timedelta | None = None) -> str:
-    to_encode = data.copy()
+# helpers
+def _now() -> datetime:
+    return datetime.utcnow()
+
+
+def _jti() -> str:
+    return str(uuid.uuid4())
+
+
+# Token factories
+def create_access_token(
+    subject: Union[str, int, Dict[str, Any]],
+    extra: Optional[Dict[str, Any]] = None,
+    expires_delta: timedelta | None = None,
+) -> str:
+    """
+    Create access token.
+    """
+
+    if isinstance(subject, dict):
+        payload_access: Dict[str, Any] = subject.copy()
+    else:
+        payload_access: Dict[str, Any] = {"sub": str(subject)}  # type: ignore
+
+    payload_access.setdefault("type", "access")
+    payload_access.setdefault("jti", _jti())
+    payload_access.setdefault("iat", int(_now().timestamp()))
+
+    if extra:
+        payload_access.update(extra)
+
     if expires_delta:
-        expire = datetime.utcnow() + expires_delta
+        exp = _now() + expires_delta
     else:
-        expire = datetime.utcnow() + timedelta(
-            minutes=settings.JWT_ACCESS_TOKEN_EXPIRE_MINUTES
-        )
-    to_encode.update({"exp": expire})
-    encoded_jwt = jwt.encode(
-        to_encode, settings.SECRET_KEY, algorithm=settings.JWT_ALGORITHM
+        exp = _now() + timedelta(minutes=settings.JWT_ACCESS_TOKEN_EXPIRE_MINUTES or 15)
+
+    payload_access["exp"] = int(exp.timestamp())
+
+    return jwt.encode(
+        payload_access, settings.SECRET_KEY, algorithm=settings.JWT_ALGORITHM
     )
-    return encoded_jwt
 
 
-def create_refresh_token(data: dict, expires_delta: timedelta | None = None) -> str:
-    to_encode = data.copy()
+def create_refresh_token(
+    subject: Union[str, int, Dict[str, Any]],
+    extra: Optional[Dict[str, Any]] = None,
+    expires_delta: timedelta | None = None,
+) -> str:
+    """
+    Create refresh token.
+    """
+
+    if isinstance(subject, dict):
+        payload_refresh: Dict[str, Any] = subject.copy()
+    else:
+        payload_refresh: Dict[str, Any] = {"sub": str(subject)}  # type: ignore
+
+    payload_refresh.setdefault("type", "refresh")
+    payload_refresh.setdefault("jti", _jti())
+    payload_refresh.setdefault("iat", int(_now().timestamp()))
+
+    if extra:
+        payload_refresh.update(extra)
+
     if expires_delta:
-        expire = datetime.utcnow() + expires_delta
+        exp = _now() + expires_delta
     else:
-        expire = datetime.utcnow() + timedelta(
-            days=settings.JWT_REFRESH_TOKEN_EXPIRES_DAYS
-        )
-    to_encode.update({"exp": expire})
-    encoded_jwt = jwt.encode(
-        to_encode, settings.SECRET_KEY, algorithm=settings.JWT_ALGORITHM
+        exp = _now() + timedelta(days=settings.JWT_REFRESH_TOKEN_EXPIRES_DAYS or 30)
+
+    payload_refresh["exp"] = int(exp.timestamp())
+
+    return jwt.encode(
+        payload_refresh, settings.SECRET_KEY, algorithm=settings.JWT_ALGORITHM
     )
-    return encoded_jwt
+
+
+def decode_token(token: str) -> dict:
+    try:
+        payload = jwt.decode(
+            token, settings.SECRET_KEY, algorithms=[settings.JWT_ALGORITHM]
+        )
+        return payload
+    except ExpiredSignatureError:
+        raise ExpiredSignatureError("Token has expired")
+
+    except JWTError:
+        raise JWTError("Invalid token")

src/db/crud/category.py

@@ -4,7 +4,7 @@
 from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncSession
 
-from src.db.dependencies import get_db_session
+from src.db.dependencies.sessions import get_db_session
 from src.db.models.category import Category
 from src.schemas.category import CategoryCreateScheme, CategoryUpdateScheme
 

src/db/crud/token.py

@@ -0,0 +1,39 @@
+from datetime import datetime
+from typing import Optional
+
+from sqlalchemy import delete, select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from db.models.revoked_token import RevokedToken
+
+
+class TokenCRUD:
+    def __init__(self, sessino: AsyncSession) -> None:
+        self.sessino = sessino
+
+    async def add(
+        self,
+        jti: str,
+        token_type: str,
+        expires_at: datetime,
+        user_id: Optional[int] = None,
+    ) -> None:
+        revoked = RevokedToken(
+            jti=jti, token_type=token_type, expires_at=expires_at, user_id=user_id
+        )
+        self.sessino.add(revoked)
+        await self.sessino.commit()
+
+    async def exists(self, jti: str) -> bool:
+        stmt = await self.sessino.execute(
+            select(RevokedToken).where(RevokedToken.jti == jti)
+        )
+        return stmt.scalar_one_or_none() is not None
+
+    async def cleanup_expired(self) -> int:
+        now = datetime.utcnow()
+        stmt = delete(RevokedToken).where(RevokedToken.expires_at < now)
+        res = await self.sessino.execute(stmt)
+        await self.sessino.commit()
+        # The result.rowcount may be driver-dependent; we return 0/1+ best-effort
+        return getattr(res, "rowcount", 0)

src/db/crud/user.py

@@ -2,7 +2,7 @@
 from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncSession
 
-from db.dependencies import get_db_session
+from db.dependencies.sessions import get_db_session
 from db.models.users import User
 from schemas.auth import RegistrationScheme
 

src/db/dependencies/__init__.py

@@ -0,0 +1,67 @@
+# src/dependencies/auth.py
+from fastapi import Depends, HTTPException, status
+from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from core.security import decode_token
+from db.crud.token import TokenCRUD
+from db.crud.user import UserCRUD
+from db.dependencies.sessions import get_db_session
+
+http_bearer = HTTPBearer(auto_error=False)
+
+
+async def get_current_user(
+    credentials: HTTPAuthorizationCredentials = Depends(http_bearer),
+    session: AsyncSession = Depends(get_db_session),
+):
+    if credentials is None:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED, detail="Not authenticated"
+        )
+
+    token = credentials.credentials
+
+    try:
+        payload = decode_token(token)
+    except ValueError as exc:
+        reason = str(exc)
+        if reason == "token_expired":
+            raise HTTPException(
+                status_code=status.HTTP_401_UNAUTHORIZED, detail="Token expired"
+            )
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid token"
+        )
+
+    if payload.get("type") != "access":
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED, detail="Token is not access token"
+        )
+
+    jti = payload.get("jti")
+    token_crud = TokenCRUD(session)
+    if await token_crud.exists(jti):
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED, detail="Token revoked"
+        )
+
+    user_crud = UserCRUD(session)
+    user = await user_crud.get_by_id(int(payload.get("sub")))
+    if not user:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED, detail="User not found"
+        )
+
+    return user
+
+
+def require_roles(*roles: str):
+    async def _inner(user=Depends(get_current_user)):
+        if getattr(user, "role", None) not in roles:
+            raise HTTPException(
+                status_code=status.HTTP_403_FORBIDDEN, detail="Forbidden"
+            )
+        return user
+
+    return _inner

src/db/dependencies/auth.py

@@ -0,0 +1,21 @@
+from collections.abc import AsyncGenerator
+
+from sqlalchemy.ext.asyncio import AsyncSession
+from starlette.requests import Request
+
+from core.database import get_async_session_maker
+
+
+async def get_db_session(request: Request) -> AsyncGenerator[AsyncSession, None]:
+    """Dependency that provides a database session."""
+    async_session_maker = get_async_session_maker()
+    session = async_session_maker()
+
+    try:
+        yield session
+        await session.commit()
+    except Exception:
+        await session.rollback()
+        raise
+    finally:
+        await session.close()