diff --git a/app/main.py b/app/main.py index b4b767d..6692cfd 100644 --- a/app/main.py +++ b/app/main.py @@ -1,5 +1,3 @@ -from fastapi import FastAPI, File, UploadFile - import dotenv import os import random @@ -8,6 +6,13 @@ import string import asyncio import uvicorn +from datetime import datetime +from fastapi import FastAPI, Response, status, Request + +from schemas.api_schemas import * +from s3_worker import S3Worker +from redis_worker import RedisWorker + dotenv.load_dotenv() disable_docs = os.getenv("DISABLE_DOCS", "true").lower() == "true" app: FastAPI = FastAPI(title="DropMeFiles analog") @@ -18,21 +23,31 @@ app: FastAPI = FastAPI(title="DropMeFiles analog") # docs_url=None if disable_docs else "/docs", # redoc_url=None if disable_docs else "/redoc", # openapi_url=None if disable_docs else "/openapi.json") +s3_worker = S3Worker() +redis_worker = RedisWorker() -data = dict() -@app.post("/upload_file") -async def create_upload_file(file: UploadFile = File(...)): - random_string = ''.join(random.choices(string.ascii_letters + string.digits, k=6)) - data[random_string] = file.filename - return {"filename": file.filename, "content_type": file.content_type, "id": random_string} +@app.get("/upload_token", status_code=200) +async def get_upload_token(file_name: str, file_type: str, file_size: int, response: Response, request: Request) -> UploadToken | BaseResponse: + if file_size > int(os.getenv('MAX_FILES_SIZE')): + response.status_code = status.HTTP_413_CONTENT_TOO_LARGE + return BaseResponse(result="The uploaded file is too large", error=True) + elif file_size <= 0: + response.status_code = status.HTTP_400_BAD_REQUEST + return BaseResponse(result="The file you are uploading is less than 1 byte, WTF?", error=True) -@app.get("/{item_id}") -async def read_item(item_id: str): - if item_id in data.keys(): - return {"return": data[item_id]} - else: - return {"return": "NO DATA!"} + user_ip = request.client.host + file_uuid = ''.join(random.choices(string.ascii_letters + string.digits, k=6)) + async with s3_worker as worker: + try: + post_data = await worker.generate_upload_post(file_name, content_type=file_type) + except Exception as exception: + response.status_code = status.HTTP_500_INTERNAL_SERVER_ERROR + return BaseResponse(result="Error generating S3 access token. Error: " + str(exception), error=True) + + redis_worker.create_record(user_ip, file_name, file_uuid, file_type, datetime.now().isoformat(), file_size) + print(post_data) + return UploadToken.model_validate(post_data) async def main(): config = uvicorn.Config("main:app", port=5000, host="0.0.0.0", log_level="debug") diff --git a/app/pyproject.toml b/app/pyproject.toml index 0251a7e..63551fe 100644 --- a/app/pyproject.toml +++ b/app/pyproject.toml @@ -9,6 +9,8 @@ dependencies = [ "dotenv>=0.9.9", "fastapi>=0.121.1", "minio>=7.2.20", + "pydantic>=2.12.5", "python-multipart>=0.0.20", + "redis>=7.1.0", "uvicorn>=0.38.0", ] diff --git a/app/redis_worker.py b/app/redis_worker.py index 4a2812b..16edcb6 100644 --- a/app/redis_worker.py +++ b/app/redis_worker.py @@ -17,15 +17,17 @@ class RedisWorker: ) self.files_ttl = int(os.getenv('FILES_TTL')) - def create_record(self, user_ip: str, file_name: str, file_uuid: str, file_type: str, add_date): + def create_record(self, user_ip: str, file_name: str, file_uuid: str, file_type: str, add_date:str, file_size:int): self.client.hset(file_uuid, mapping={ "file_name": file_name, + "file_size": file_size, "user_ip": user_ip, "file_type": file_type, "add_date": add_date }) self.client.expire(file_uuid, self.files_ttl) + def get_record(self, file_uuid:str): return self.client.hgetall(file_uuid) diff --git a/app/s3_worker.py b/app/s3_worker.py index f691b65..1903c97 100644 --- a/app/s3_worker.py +++ b/app/s3_worker.py @@ -2,7 +2,6 @@ import os import asyncio -from anyio import Condition from dotenv import load_dotenv from aiobotocore.session import get_session @@ -75,9 +74,9 @@ class S3Worker: ExpiresIn=expires_in, ) - +s3_worker = S3Worker() async def test_run(): - async with S3Worker() as worker: + async with s3_worker as worker: await worker.upload_file("test.txt", b"hello") file = await worker.download_file("test.txt") file_text = await file["Body"].read() diff --git a/app/schemas/api_schemas.py b/app/schemas/api_schemas.py new file mode 100644 index 0000000..8d0dd1b --- /dev/null +++ b/app/schemas/api_schemas.py @@ -0,0 +1,21 @@ +import typing + +from pydantic import BaseModel, Field, HttpUrl, PositiveInt + + +class BaseResponse(BaseModel): + result: typing.Any + error: bool = False + +class UploadFields(BaseModel): + content_type: str = Field(alias="Content-Type") + acl: str + key: str + aws_access_key_id: str = Field(alias="AWSAccessKeyId") + policy: str + signature: str + + +class UploadToken(BaseModel): + url: HttpUrl + fields: UploadFields \ No newline at end of file