Upload files in single request with images

Author: IAmTomahawkx

config-template.json

@@ -20,6 +20,10 @@
         "google_application_secret": "",
         "github_bot_token": ""
     },
+    "bunny_cdn": {
+        "hostname": "mystbin",
+        "token": ""
+    },
     "site": {
         "frontend_site": "https://mysite.com",
         "backend_site": "https://api.mysite.com",
@@ -29,6 +33,7 @@
     "paste": {
         "character_limit": 300000,
         "file_limit": 5,
+        "filesize_limit": "8mb",
         "log_ip": true
     },
     "debug": {

mystbin/backend/models/payloads.py

@@ -27,7 +27,14 @@ class PasteFile(BaseModel):
     filename: str
 
     class Config:
-        schema_extra = {"content": "explosions everywhere", "filename": "kaboom.txt"}
+        schema_extra = {"example": {"content": "explosions everywhere", "filename": "kaboom.txt"}}
+
+
+class RichPasteFile(PasteFile):
+    attachment: str | None
+
+    class Config:
+        schema_extra = {"example": {"content": "explosions everywhere", "filename": "kaboom.txt", "attachment": "image1.png"}}
 
 
 class PastePut(BaseModel):
@@ -43,13 +50,35 @@ class Config:
                 "files": [
                     {"content": "import this", "filename": "foo.py"},
                     {
-                        "content": "doc.md",
-                        "filename": "**do not use this in production**",
+                        "filename": "doc.md",
+                        "content": "**do not use this in production**",
                     },
                 ],
             }
         }
 
+class RichPastePost(PastePut):
+    files: List[RichPasteFile]
+
+    class Config:
+        schema_extra = {
+            "example": {
+                "expires": "2020-11-16T13:46:49.215Z",
+                "password": None,
+                "files": [
+                    {
+                        "content": "import this",
+                        "filename": "foo.py",
+                        "attachment": None
+                    },
+                    {
+                        "filename": "doc.md",
+                        "content": "**do not use this in production**",
+                        "attachment": "image2.jpeg"
+                    },
+                ],
+            }
+        }
 
 class PastePatch(BaseModel):
     new_expires: Optional[datetime.datetime] = None

mystbin/backend/models/responses.py

@@ -27,8 +27,7 @@ class File(BaseModel):
     content: str
     loc: int
     charcount: int
-    tab_id: int | None
-    image: str | None
+    attachment: str | None
 
     class Config:
         schema_extra = {
@@ -37,8 +36,7 @@ class Config:
                 "content": "import datetime\\nprint(datetime.datetime.utcnow())",
                 "loc": 2,
                 "charcount": 49,
-                "tab_id": 0,
-                "url": "...",
+                "attachment": "https://mystbin.b-cdn.com/umbra_sucks.jpeg",
             }
         }
 

mystbin/backend/routers/pastes.py

@@ -17,15 +17,18 @@
 along with MystBin.  If not, see <https://www.gnu.org/licenses/>.
 """
 from __future__ import annotations
+import asyncio
 
 import datetime
 import json
 import pathlib
 import re
+import uuid
 from random import sample
-from typing import Dict, List, Optional, Union
+from typing import Coroutine, cast, Dict, List, Optional, Union
 
 from asyncpg import Record
+import pydantic
 from fastapi import APIRouter, File, Response, UploadFile
 from fastapi.responses import UJSONResponse
 from fastapi.encoders import jsonable_encoder
@@ -35,6 +38,13 @@
 from utils.db import _recursive_hook as recursive_hook
 from utils.ratelimits import limit
 
+try:
+    import ujson
+    _loads = ujson.loads
+except ModuleNotFoundError:
+    import json
+    _loads = json.loads
+
 
 _WORDS_LIST = open(pathlib.Path("utils/words.txt")).readlines()
 word_list = [word.title() for word in _WORDS_LIST if len(word) > 3]
@@ -54,9 +64,9 @@
 del __p, __f  # micro-opt, don't keep unneeded variables in-ram
 
 
-def generate_paste_id():
+def generate_paste_id(n: int = 3):
     """Generate three random words."""
-    word_samples = sample(word_list, 3)
+    word_samples = sample(word_list, n)
     return "".join(word_samples).replace("\n", "")
 
 
@@ -137,7 +147,7 @@ async def find_discord_tokens(request: MystbinRequest, pastes: payloads.PastePut
     response_model=responses.PastePostResponse,
     responses={
         201: {"model": responses.PastePostResponse},
-        400: {"content": {"application/json": {"example": {"error": "files.length: You have provided a bad paste"}}}},
+        400: {"content": {"application/json": {"example": {"error": "files.length: You have provided too many files"}}}},
     },
     status_code=201,
     name="Create paste",
@@ -178,42 +188,87 @@ async def put_pastes(
     return UJSONResponse(paste, status_code=201)
 
 
-@router.put(
-    "/images/upload/{paste_id}",
+@router.post(
+    "/rich-paste",
     tags=["pastes"],
+    response_model=responses.PastePostResponse,
     responses={
         201: {"model": responses.PastePostResponse},
-        401: {"model": errors.Unauthorized},
-        404: {"model": errors.NotFound},
+        400: {"content": {"application/json": {"example": {"error": "files.length: You have provided too many files"}}}}
     },
-    include_in_schema=False,
+    status_code=201,
+    name="Create Rich Paste"
 )
 @limit("postpastes")
-async def get_image_upload_link(
-    request: MystbinRequest, paste_id: str, password: Optional[str] = None, images: List[UploadFile] = File(...)
+async def post_rich_paste(
+    request: MystbinRequest,
+    data: bytes = File(None, max_length=(__config["paste"]["character_limit"] * __config["paste"]["file_limit"]) + 500),
+    images: List[UploadFile] = File(None),
 ):
-    """user = request.state.user
-    if not user:
-        return UJSONResponse({"error": "Unauthorized", "notice": "You must be signed in to use this route"}, status_code=401)"""
-
-    paste = await request.app.state.db.get_paste(paste_id, password)
-    if paste is None:
-        return UJSONResponse({"error": "Not Found"}, status_code=404)
+            
+    reads = data
+    try:
+        payload = payloads.RichPastePost.parse_raw(reads, content_type="application/json")
+        print(payload)
+    except pydantic.ValidationError as e:
+        return UJSONResponse({"detail": e.errors()}, status_code=422)
+    except:
+        return UJSONResponse({"error": f"multipart.section.data: Invalid JSON"})
+    
+    paste_id = generate_paste_id()
+
+    if images:
+        async def _partial(target, spool: UploadFile):
+            data = await spool.read()
+            await request.app.state.client.put(target, data=data, headers=headers) # TODO figure out how to pass spooled object instead of load into memory
+        
+        image_idx = {}
+        headers = {"Content-Type": "application/octet-stream", "AccessKey": f"{__config['bunny_cdn']['token']}"}
+        partials: list[Coroutine] = []
+
+        for image in images: # TODO honour config filesize limit
+            origin = image.filename.split(".")[-1]
+            new_name = f"{('%032x' % uuid.uuid4().int)[:8]}-{paste_id}.{origin}"
+            url = f"https://{__config['bunny_cdn']['hostname']}.b-cdn.net/images/{new_name}"
+            image_idx[image.filename] = url
+            target = f"https://storage.bunnycdn.com/{__config['bunny_cdn']['hostname']}/images/{new_name}"
+            partials.append(_partial(target, image))
+
+        for n, file in enumerate(payload.files):
+            if file.attachment is not None:
+                if file.attachment not in image_idx:
+                    return UJSONResponse({"error": f"files.{n}.attachment: Unkown attachment '{file.attachment}'"})
+                
+                file.__dict__["attachment"] = image_idx[file.attachment]
+        
+        await asyncio.wait(partials, return_when=asyncio.ALL_COMPLETED)
+    
+    if err := enforce_multipaste_limit(request.app, payload):
+        return err
 
-    headers = {"Content-Type": "application/octet-stream", "AccessKey": f"{__config['bunny_cdn']['token']}"}
+    notice = None
 
-    for image in images:
-        i = image.filename[0]
-        await request.app.state.db.update_paste_with_files(
-            paste_id=paste_id, tab_id=i, url=f"https://mystbin.b-cdn.net/images/{image.filename}"
-        )
+    if tokens := await find_discord_tokens(request, payload):
+        gist = await upload_to_gist(request, "\n".join(tokens))
+        notice = f"Discord tokens have been found and uploaded to {gist['html_url']}"
 
-        URL = f'https://storage.bunnycdn.com/{__config["bunny_cdn"]["hostname"]}/images/{image.filename}'
-        data = await image.read()
+    author: Optional[int] = request.state.user["id"] if request.state.user else None
 
-        await request.app.state.client.put(URL, headers=headers, data=data)
+    paste = await request.app.state.db.put_paste(
+        paste_id=paste_id,
+        pages=payload.files,
+        expires=payload.expires,
+        author=author,
+        password=payload.password,
+        origin_ip=request.headers.get("x-forwarded-for", request.client.host)
+        if request.app.config["paste"]["log_ip"]
+        else None,
+    )
 
-    return Response(status_code=201)
+    paste["notice"] = notice
+    paste = responses.PastePostResponse(**paste)  # type: ignore # this is a problem for future me #TODO
+    paste = recursive_hook(paste.dict())
+    return UJSONResponse(paste, status_code=201)
 
 
 desc = f"""Get a paste by ID.

mystbin/backend/utils/db.py

@@ -182,16 +182,6 @@ async def get_paste(self, paste_id: str, password: Optional[str] = None) -> Opti
                 contents = cast(List[asyncpg.Record], await self._do_query(query, paste_id, conn=conn))
                 resp = dict(resp[0])
                 resp["files"] = [{a: str(b) for a, b in x.items()} for x in contents]
-
-                images = await self.get_images(paste_id=paste_id)
-
-                for index, file in enumerate(resp["files"]):
-                    try:
-                        file["image"] = images[index]["url"]
-                        file["tab_id"] = images[index]["tab_id"]
-                    except IndexError:
-                        pass
-
                 return resp
             else:
                 return None
@@ -232,7 +222,7 @@ async def put_paste(
         *,
         paste_id: str,
         origin_ip: Optional[str],
-        pages: List[payloads.PasteFile],
+        pages: List[payloads.RichPasteFile] | List[payloads.PasteFile],
         expires: Optional[datetime.datetime] = None,
         author: Optional[int] = None,
         password: Optional[str] = None,
@@ -281,13 +271,14 @@ async def put_paste(
                         page.content,
                         page.filename,
                         page.content.count("\n") + 1,  # add an extra for line 1
+                        getattr(page, "attachment", None)
                     )
                 )
 
             files_query = """
-                          INSERT INTO files (parent_id, content, filename, loc)
-                          VALUES ($1, $2, $3, $4)
-                          RETURNING index, filename, loc, charcount, content
+                          INSERT INTO files (parent_id, content, filename, loc, attachment)
+                          VALUES ($1, $2, $3, $4, $5)
+                          RETURNING index, filename, loc, charcount, content, attachment
                           """
             inserted = []
             async with conn.transaction():
@@ -301,18 +292,6 @@ async def put_paste(
 
             return resp
 
-    async def update_paste_with_files(self, *, paste_id: str, tab_id: str, url: str) -> None:
-        query = """INSERT INTO images VALUES($1, $2, $3)"""
-
-        async with self.pool.acquire() as conn:
-            await self._do_query(query, paste_id, int(tab_id), url)
-
-    async def get_images(self, *, paste_id: str):
-
-        query = """SELECT * FROM images WHERE parent_id = $1"""
-        async with self.pool.acquire() as conn:
-            return [dict(x) for x in await self._do_query(query, paste_id)]
-
     @wrapped_hook_callback
     async def edit_paste(
         self,

mystbin/database/schema.sql

@@ -32,18 +32,10 @@ CREATE TABLE IF NOT EXISTS files (
     loc INTEGER NOT NULL,
     charcount INTEGER GENERATED ALWAYS AS (LENGTH(content)) STORED,
     index SERIAL NOT NULL,
+    attachment TEXT,
     PRIMARY KEY (parent_id, index)
 );
 
-CREATE TABLE IF NOT EXISTS images (
-    parent_id TEXT REFERENCES pastes(id) ON DELETE CASCADE,
-    tab_id BIGINT,
-    url TEXT,
-    index SERIAL NOT NULL,
-    PRIMARY KEY (parent_id, index)
-
-);
-
 CREATE TABLE IF NOT EXISTS bookmarks (
     userid bigint not null references users(id) ON DELETE CASCADE,
     paste text not null references pastes(id) ON DELETE CASCADE,