Multimedika/Bot_Development

1import os2from dotenv import load_dotenv3 4from datetime import timedelta5from typing import Annotated6 7from fastapi import APIRouter, Depends, status8from fastapi.responses import JSONResponse9from fastapi.security import OAuth2PasswordRequestForm10from passlib.context import CryptContext11from sqlalchemy.orm import Session12 13from db.models import User14from db.database import get_db15from api.auth import get_current_user, create_access_token16from service.dto import CreateUserRequest, UserVerification, Token17 18from collections import Counter19from time import time20 21 22load_dotenv()23 24 25router = APIRouter(tags=["User"])26 27bcrypt_context = CryptContext(schemes=["bcrypt"], deprecated="auto")28 29db_dependency = Annotated[Session, Depends(get_db)]30user_dependency = Annotated[dict, Depends(get_current_user)]31 32ACCESS_TOKEN_EXPIRE_MINUTES = 4320033 34 35# Rate-limiting config36FAILED_ATTEMPT_LIMIT = 337BLOCK_TIME_SECONDS = 300  # Block for 5 minutes38 39# In-memory tracking for failed attempts40failed_attempts = Counter()41blocked_users = {}42 43 44@router.post("/login", response_model=Token)45async def login_for_access_token(46    login_data: Annotated[OAuth2PasswordRequestForm, Depends()],47    db: Session = Depends(get_db),48):49    username = login_data.username50 51    # Check if user is blocked52    if username in blocked_users:53        block_until = blocked_users[username]54        if time() < block_until:55            return JSONResponse(56                status_code=status.HTTP_403_FORBIDDEN,57                content=f"Too many failed attempts. Try again after {int(block_until - time())} seconds.",58            )59 60        else:61            # Unblock the user after the time period62            del blocked_users[username]63            del failed_attempts[username]64 65    user = db.query(User).filter(User.username == username).first()66 67    if not user:68        # Automatically register the user69        create_user_request = CreateUserRequest(70            name=login_data.username,71            username=login_data.username,72            email=login_data.username,73            password=os.getenv("USER_PASSWORD"),  # Replace with a generated or temporary password74            role_id=2,75        )76        registration_response = await register_user(db, create_user_request)77 78        if isinstance(registration_response, JSONResponse):79            return registration_response  # Return error response if registration failed80 81        # Retrieve the newly created user after successful registration82        user = db.query(User).filter(User.username == username).first()83        84        if not user:85            return JSONResponse(86                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,87                content="User registration failed unexpectedly."88            )89            90        correct_password = (91            bcrypt_context.verify(os.getenv("USER_PASSWORD"), user.password_hash) or92            bcrypt_context.verify(login_data.password, user.password_hash)93        )94        95        if not correct_password :96            failed_attempts[username] = failed_attempts.get(username, 0) + 197            if failed_attempts[username] >= FAILED_ATTEMPT_LIMIT:98                blocked_users[username] = time() + BLOCK_TIME_SECONDS99                failed_attempts.pop(username, None)  # Reset after blocking100                return JSONResponse(101                    status_code=status.HTTP_403_FORBIDDEN,102                    content="Too many failed attempts. You are temporarily blocked."103                )104 105            return JSONResponse(106                status_code=status.HTTP_401_UNAUTHORIZED,107                content="Invalid credentials."108            )109            110    failed_attempts.pop(username, None)111 112    try:113        access_token_expires = timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)114        access_token = create_access_token(115            user.username,116            user.name,117            user.id,118            user.role_id,119            access_token_expires,120            user.email,121        )122 123        return {"access_token": access_token, "token_type": "bearer"}124 125    except Exception as e:126        print(e)127        return JSONResponse(status_code=500, content="An error occurred during login")128 129 130@router.get("/login", response_model=dict)131async def get_user(user: user_dependency):132    if user is None:133        return JSONResponse(status_code=401, content="Authentication Failed")134    return {135        "username": user.get("username"),136        "name": user.get("name"),137        "id": user.get("id"),138        "email": user.get("email"),139        "role": user.get("role_id"),140    }141 142 143@router.get("/users", response_model=list[dict])144async def get_all_users(user: user_dependency, db: Session = Depends(get_db)):145    # Check if the current user has an admin role146    if user.get("role_id") != 1:  # Adjust this check based on how roles are represented147        return JSONResponse(status_code=401, content="Authentication Failed")148 149    # Query the database to retrieve all users150    users = db.query(151        User152    ).all()  # Assuming you have a User model with an SQLAlchemy session153    return [154        {155            "id": user.id,156            "username": user.username,157            "name": user.name,158            "email": user.email,159            "role": user.role_id,160        }161        for user in users162    ]163 164 165async def register_user(db: db_dependency, create_user_request: CreateUserRequest):166    existing_user = (167        db.query(User).filter(User.email == create_user_request.email).first()168    )169 170    if existing_user:171        return JSONResponse(status_code=400, content="Email is already registered")172 173    try:174        password_hash = bcrypt_context.hash(create_user_request.password)175 176        create_user_model = User(177            name=create_user_request.name,178            username=create_user_request.username,179            email=create_user_request.email,180            role_id=create_user_request.role_id,181            password_hash=password_hash,182        )183 184        db.add(create_user_model)185        db.commit()186        db.refresh(create_user_model)187 188        return {"message": "User created successfully", "user_id": create_user_model.id}189    except Exception as e:190        print(e)191        return JSONResponse(192            status_code=500, content="An error occuring when register user"193        )194 195 196# @router.post("/forgot_password")197# async def forget_password():198#     pass199 200 201# @router.post("/change_password")202# async def change_password(203#     user: user_dependency, db: db_dependency, user_verification: UserVerification204# ):205#     if user is None:206#         return JSONResponse(status_code=401, content="Authentication Failed")207#     user_model = db.query(User).filter(User.id == user.get("id")).first()208 209#     if not bcrypt_context.verify(210#         user_verification.password, user_model.hashed_password211#     ):212#         return JSONResponse(status_code=401, content="Error on password change")213 214#     user_model.hashed_password = bcrypt_context.hash(user_verification.new_password)215#     db.add(user_model)216#     db.commit()217#     db.refresh(user_model)218 219#     return {"message": "User's password successfully changed", "user_id": user_model.id}220