Create auth.py

auth.py

@@ -0,0 +1,276 @@
+import os
+import uuid
+import logging
+from datetime import datetime, timedelta
+from urllib.parse import quote_plus
+from typing import Optional
+
+from dotenv import load_dotenv
+from fastapi import APIRouter, HTTPException, Depends, Request, status, UploadFile, File, Form
+from fastapi.responses import JSONResponse
+from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm
+from slowapi import Limiter, _rate_limit_exceeded_handler
+from slowapi.util import get_remote_address
+from jose import JWTError, jwt
+from passlib.context import CryptContext
+from pydantic import BaseModel, EmailStr, Field, validator
+from pymongo import MongoClient
+import os.path
+
+load_dotenv()
+
+# Setup logging
+logger = logging.getLogger("uvicorn")
+logger.setLevel(logging.INFO)
+
+# MongoDB setup for user management
+password = quote_plus(os.getenv("MONGO_PASSWORD"))
+MONGO_URL = os.getenv("MONGO_ENDPOINT").replace("${PASSWORD}", password)
+client = MongoClient(MONGO_URL)
+db = client.users_database
+users_collection = db.users
+
+# Rate limiter (applied to auth endpoints)
+limiter = Limiter(key_func=get_remote_address, default_limits=["200 per day", "50 per hour"])
+
+# OAuth2 setup
+oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
+
+# Create an APIRouter instance
+router = APIRouter()
+
+# Pydantic models
+class User(BaseModel):
+    name: str = Field(..., min_length=3, max_length=50)
+    email: EmailStr
+    password: str
+
+    @validator("password")
+    def validate_password(cls, value):
+        if len(value) < 8:
+            raise ValueError("Password must be at least 8 characters long.")
+        if not any(char.isdigit() for char in value):
+            raise ValueError("Password must include at least one number.")
+        if not any(char.isupper() for char in value):
+            raise ValueError("Password must include at least one uppercase letter.")
+        if not any(char.islower() for char in value):
+            raise ValueError("Password must include at least one lowercase letter.")
+        if not any(char in "!@#$%^&*()-_+=<>?/" for char in value):
+            raise ValueError("Password must include at least one special character.")
+        return value
+
+class UserUpdate(BaseModel):
+    name: Optional[str] = Field(None, min_length=3, max_length=50)
+    email: Optional[EmailStr]
+    password: Optional[str]
+
+    @validator("password")
+    def validate_password(cls, value):
+        if value is not None:
+            if len(value) < 8:
+                raise ValueError("Password must be at least 8 characters long.")
+            if not any(char.isdigit() for char in value):
+                raise ValueError("Password must include at least one number.")
+            if not any(char.isupper() for char in value):
+                raise ValueError("Password must include at least one uppercase letter.")
+            if not any(char.islower() for char in value):
+                raise ValueError("Password must include at least one lowercase letter.")
+            if not any(char in "!@#$%^&*()-_+=<>?/" for char in value):
+                raise ValueError("Password must include at least one special character.")
+        return value
+
+class Token(BaseModel):
+    access_token: str
+    refresh_token: str
+    token_type: str
+
+class LoginResponse(Token):
+    name: str
+    avatar: Optional[str] = None
+
+class TokenData(BaseModel):
+    email: Optional[str] = None
+
+# Password hashing
+pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
+
+def verify_password(plain_password: str, hashed_password: str) -> bool:
+    return pwd_context.verify(plain_password, hashed_password)
+
+def get_password_hash(password: str) -> str:
+    return pwd_context.hash(password)
+
+def get_user(email: str) -> Optional[dict]:
+    return users_collection.find_one({"email": email})
+
+def authenticate_user(email: str, password: str) -> Optional[dict]:
+    user = get_user(email)
+    if not user or not verify_password(password, user["hashed_password"]):
+        return None
+    return user
+
+def create_token(data: dict, expires_delta: timedelta = None) -> str:
+    to_encode = data.copy()
+    expire = datetime.utcnow() + (expires_delta or timedelta(minutes=15))
+    to_encode.update({"exp": expire})
+    secret_key = os.getenv("SECRET_KEY")
+    algorithm = "HS256"
+    return jwt.encode(to_encode, secret_key, algorithm=algorithm)
+
+def create_access_token(email: str) -> str:
+    return create_token({"sub": email}, timedelta(minutes=int(os.getenv("ACCESS_TOKEN_EXPIRE_MINUTES", "90"))))
+
+def create_refresh_token(email: str) -> str:
+    return create_token({"sub": email}, timedelta(days=int(os.getenv("REFRESH_TOKEN_EXPIRE_DAYS", "7"))))
+
+def get_current_user(token: str = Depends(oauth2_scheme)) -> dict:
+    secret_key = os.getenv("SECRET_KEY")
+    try:
+        payload = jwt.decode(token, secret_key, algorithms=["HS256"])
+        email: str = payload.get("sub")
+        if not email:
+            raise HTTPException(status_code=401, detail="Invalid credentials")
+        user = get_user(email)
+        if not user:
+            raise HTTPException(status_code=401, detail="User not found")
+        return user
+    except JWTError:
+        raise HTTPException(status_code=401, detail="Invalid token")
+
+# Setup for avatar file saving
+AVATAR_DIR = "avatars"
+if not os.path.exists(AVATAR_DIR):
+    os.makedirs(AVATAR_DIR)
+
+def save_avatar_file(file: UploadFile) -> str:
+    allowed_types = ["image/jpeg", "image/png", "image/gif"]
+    if file.content_type not in allowed_types:
+        logger.error(f"Unsupported file type: {file.content_type}")
+        raise HTTPException(
+            status_code=400,
+            detail="Invalid image format. Only JPEG, PNG, and GIF are accepted."
+        )
+    file_extension = os.path.splitext(file.filename)[1]
+    unique_filename = f"{uuid.uuid4()}{file_extension}"
+    file_path = os.path.join(AVATAR_DIR, unique_filename)
+    try:
+        contents = file.file.read()
+        with open(file_path, "wb") as f:
+            f.write(contents)
+        logger.info(f"Avatar saved as {file_path}")
+    except Exception as e:
+        logger.exception("Failed to save avatar file")
+        raise HTTPException(status_code=500, detail="Could not save avatar file.")
+    finally:
+        file.file.close()
+    return file_path
+
+# ----- Auth Endpoints -----
+
+@router.post("/signup", response_model=Token)
+@limiter.limit("5/minute")
+async def signup(
+    request: Request,
+    name: str = Form(...),
+    email: EmailStr = Form(...),
+    password: str = Form(...),
+    avatar: Optional[UploadFile] = File(None)
+):
+    # Validate input using the User model
+    try:
+        _ = User(name=name, email=email, password=password)
+    except Exception as e:
+        logger.error(f"Validation error during signup: {e}")
+        raise HTTPException(status_code=400, detail=str(e))
+    if get_user(email):
+        logger.warning(f"Attempt to register already existing email: {email}")
+        raise HTTPException(status_code=400, detail="Email already registered")
+    hashed_password = get_password_hash(password)
+    user_data = {
+        "name": name,
+        "email": email,
+        "hashed_password": hashed_password,
+        "chat_histories": []  # Initialize an empty array for chat histories and future data.
+    }
+    if avatar:
+        avatar_path = save_avatar_file(avatar)
+        user_data["avatar"] = f"/avatars/{os.path.basename(avatar_path)}"
+    users_collection.insert_one(user_data)
+    logger.info(f"New user registered: {email}")
+    return {
+        "access_token": create_access_token(email),
+        "refresh_token": create_refresh_token(email),
+        "token_type": "bearer"
+    }
+
+@router.post("/login", response_model=LoginResponse)
+@limiter.limit("10/minute")
+async def login(request: Request, form_data: OAuth2PasswordRequestForm = Depends()):
+    user = authenticate_user(form_data.username, form_data.password)
+    if not user:
+        logger.warning(f"Failed login attempt for: {form_data.username}")
+        raise HTTPException(status_code=401, detail="Incorrect username or password")
+    logger.info(f"User logged in: {user['email']}")
+    return {
+        "access_token": create_access_token(user["email"]),
+        "refresh_token": create_refresh_token(user["email"]),
+        "token_type": "bearer",
+        "name": user["name"],
+        "avatar": user.get("avatar")
+    }
+
+@router.get("/user/data")
+@limiter.limit("20/minute")
+async def get_user_data(request: Request, current_user: dict = Depends(get_current_user)):
+    return {
+        "name": current_user["name"],
+        "email": current_user["email"],
+        "avatar": current_user.get("avatar"),
+        "chat_histories": current_user.get("chat_histories", [])
+    }
+
+@router.put("/user/update")
+@limiter.limit("10/minute")
+async def update_user(
+    request: Request,
+    name: Optional[str] = Form(None),
+    email: Optional[EmailStr] = Form(None),
+    password: Optional[str] = Form(None),
+    avatar: Optional[UploadFile] = File(None),
+    current_user: dict = Depends(get_current_user)
+):
+    update_data = {}
+    if name is not None:
+        update_data["name"] = name
+    if email is not None:
+        update_data["email"] = email
+    if password is not None:
+        try:
+            _ = User(name=current_user["name"], email=current_user["email"], password=password)
+        except Exception as e:
+            logger.error(f"Password validation error during update: {e}")
+            raise HTTPException(status_code=400, detail=str(e))
+        update_data["hashed_password"] = get_password_hash(password)
+    if avatar:
+        avatar_path = save_avatar_file(avatar)
+        update_data["avatar"] = f"/avatars/{os.path.basename(avatar_path)}"
+    if not update_data:
+        logger.info("No update parameters provided")
+        raise HTTPException(status_code=400, detail="No update parameters provided")
+    users_collection.update_one({"email": current_user["email"]}, {"$set": update_data})
+    logger.info(f"User updated: {current_user['email']}")
+    return {"message": "User updated successfully"}
+
+@router.post("/logout")
+@limiter.limit("20/minute")
+async def logout(request: Request, current_user: dict = Depends(get_current_user)):
+    logger.info(f"User logged out: {current_user['email']}")
+    return {"message": "User logged out successfully"}
+
+@router.exception_handler(Exception)
+async def generic_exception_handler(request: Request, exc):
+    logger.exception("Unhandled exception occurred in auth module")
+    return JSONResponse(
+        status_code=500,
+        content={"detail": "An internal server error occurred."}
+    )