main.py

from typing import List
from uuid import UUID

from fastapi import Depends, FastAPI, HTTPException, Request, UploadFile, status
from fastapi.middleware.cors import CORSMiddleware
from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm
from itsdangerous import BadSignature
from itsdangerous.url_safe import URLSafeTimedSerializer
from slowapi import Limiter, _rate_limit_exceeded_handler
from slowapi.errors import RateLimitExceeded
from slowapi.util import get_remote_address
from sqlalchemy.orm import Session

from . import schemas, utils
from .config import settings
from .database import SessionLocal, create_database

create_database()

app = FastAPI()

# CORS handling
origins = [settings.customer_url, settings.worker_url]
app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Rate Limiting for some endpoints
limiter = Limiter(key_func=get_remote_address)
app.state.limiter = limiter
app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)

# Authentication setup
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
oauth2_tokener = URLSafeTimedSerializer(settings.signing_key)


# DB Dependency
def get_db():
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()


# oauth2/token handling
@app.get("/token")
def check_token_validity(token: str = Depends(oauth2_scheme)):
    check_token(token, None)


# TODO: move this to utils.py?
def check_token(token: str, item_uuid: str):
    try:
        auth_data = oauth2_tokener.loads(token, max_age=settings.token_lifetime * 60)
        if auth_data == "all" or auth_data == item_uuid:
            return  # success
    except BadSignature:
        print("failed to load access token")
        print(item_uuid)

    raise HTTPException(
        status_code=status.HTTP_401_UNAUTHORIZED,
        detail="Invalid authentication credentials",
        headers={"WWW-Authenticate": "Bearer"},
    )


@app.post("/token")
def verify_supporter(form_data: OAuth2PasswordRequestForm = Depends()):
    if form_data.password != settings.shared_secret:
        raise HTTPException(status_code=401, detail="Incorrect username or password")
    return {
        "access_token": oauth2_tokener.dumps("all"),
        "token_type": "bearer",
    }


# TODO: merge this with the supporter login? we want a single endpoint for all of this.
@app.post("/login")
def verify_customer(
    login_data: schemas.LoginData, db: Session = Depends(get_db)
):  # item_uuid: str, signature: str):
    print(login_data)
    delivery = utils.get_delivery_by_uuid(db, UUID(login_data.delivery_uuid))
    if not delivery:
        raise HTTPException(status_code=404, detail="Item not found")
    if not utils.verify_signature(delivery, delivery.uuid, login_data.signature):
        raise HTTPException(status_code=400, detail="Invalid signature")
    return {
        "access_token": oauth2_tokener.dumps(str(delivery.uuid)),
        "token_type": "bearer",
    }


@app.post("/delivery", response_model=schemas.Delivery)
@limiter.limit("2/minute")
def prepare_delivery(
    request: Request, delivery: schemas.DeliveryBase, db: Session = Depends(get_db)
):
    return utils.prepare_delivery(db, delivery.verification)


@app.get("/deliveries", response_model=List[schemas.Delivery])
def list_deliveries(token: str = Depends(oauth2_scheme), db: Session = Depends(get_db)):
    check_token(token, None)
    return utils.get_deliveries(db)


@app.get("/delivery/{delivery_uuid}", response_model=schemas.Delivery)
def get_delivery_by_uuid(
    delivery_uuid: str,
    token: str = Depends(oauth2_scheme),
    db: Session = Depends(get_db),
):
    check_token(token, delivery_uuid)
    delivery = utils.get_delivery_by_uuid(db, UUID(delivery_uuid))
    if not delivery:
        raise HTTPException(status_code=404, detail="Delivery not found")
    return delivery


@app.get("/tag/{delivery_tag}", response_model=schemas.Delivery)
def get_delivery_by_tag(
    delivery_tag: str,
    token: str = Depends(oauth2_scheme),
    db: Session = Depends(get_db),
):
    delivery = utils.get_delivery_by_tag(db, delivery_tag)
    if not delivery:
        raise HTTPException(status_code=404, detail="Delivery not found")
    check_token(token, delivery.uuid)
    return delivery


@app.put("/delivery/{delivery_uuid}", response_model=schemas.Delivery)
def update_delivery(
    delivery_uuid: str,
    data: schemas.DeliveryUpdate,
    token: str = Depends(oauth2_scheme),
    db: Session = Depends(get_db),
):
    check_token(token, delivery_uuid)
    delivery = utils.get_delivery_by_uuid(db, UUID(delivery_uuid))
    if not delivery:
        raise HTTPException(status_code=404, detail="Delivery not found")
    delivery = utils.update_delivery_data(db, delivery, data)
    return delivery


@app.post("/item", response_model=schemas.Item)
def add_item(
    item: schemas.ItemAdd,
    token: str = Depends(oauth2_scheme),
    db: Session = Depends(get_db),
):
    check_token(token, None)
    delivery = utils.get_delivery_by_uuid(db, item.delivery_uuid)
    if not delivery:
        raise HTTPException(status_code=404, detail="Delivery not found")
    storage = utils.get_storage_by_name(db, item.storage_name)
    if not storage:
        raise HTTPException(status_code=404, detail="Storage not found")
    return utils.add_item_for_delivery_at_storage(db, delivery, storage)


@app.get("/item/{item_uuid}", response_model=schemas.Item)
def get_item(
    item_uuid: str, token: str = Depends(oauth2_scheme), db: Session = Depends(get_db)
):
    check_token(token, item_uuid)
    item = utils.get_item_by_uuid(db, UUID(item_uuid))
    if not item:
        raise HTTPException(status_code=404, detail="Item not found")
    return item


@app.delete("/item/{item_uuid}", response_model=schemas.Item)
def deploy_item(
    item_uuid: str, token: str = Depends(oauth2_scheme), db: Session = Depends(get_db)
):
    check_token(token, None)
    item = utils.get_item_by_uuid(db, UUID(item_uuid))
    if not item:
        raise HTTPException(status_code=404, detail="Item not found")
    return utils.deploy_item(db, item)


@app.get("/storage/{storage_name}", response_model=schemas.Storage)
def get_storage(
    storage_name: str,
    token: str = Depends(oauth2_scheme),
    db: Session = Depends(get_db),
):
    check_token(token, None)
    storage = utils.get_storage_by_name(db, storage_name)
    if not storage:
        raise HTTPException(status_code=404, detail="Storage not found")
    return utils.get_storage_by_name(db, storage_name)


@app.get("/storages", response_model=List[schemas.Storage])
def get_storages(
    token: str = Depends(oauth2_scheme),
    db: Session = Depends(get_db),
):
    check_token(token, None)
    return utils.get_storages(db)


@app.post("/item/register", response_model=schemas.Item)
def add_item_with_image(
    image: UploadFile,
    storage_name: str,
    token: str = Depends(oauth2_scheme),
    db: Session = Depends(get_db),
):
    check_token(token, None)
    print(image.file)
    storage = utils.get_storage_by_name(db, storage_name)
    if not storage:
        raise HTTPException(status_code=404, detail="Storage not found")
    return utils.add_item_with_image_at_storage(db, image, storage)