discord-havoc/util/litterbox.py

"""
Litterbox (Catbox) Uploader (Async)
"""

import traceback
from typing import Optional, Union
from io import BufferedReader
from aiohttp import ClientSession, ClientTimeout, FormData
import magic
import logging
import random
import os

litterbox_api_url: str = "https://litterbox.catbox.moe/resources/internals/api.php"
http_headers: dict[str, str] = {
    "accept": "*/*",
    "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/101.0.4951.64 Safari/537.36 Edg/101.0.1210.53",
    "Accept-Language": "en-US,en;q=0.9,it;q=0.8,es;q=0.7",
    "referer": "https://www.google.com/",
    "cookie": "DSID=AAO-7r4OSkS76zbHUkiOpnI0kk-X19BLDFF53G8gbnd21VZV2iehu-w_2v14cxvRvrkd_NjIdBWX7wUiQ66f-D8kOkTKD1BhLVlqrFAaqDP3LodRK2I0NfrObmhV9HsedGE7-mQeJpwJifSxdchqf524IMh9piBflGqP0Lg0_xjGmLKEQ0F4Na6THgC06VhtUG5infEdqMQ9otlJENe3PmOQTC_UeTH5DnENYwWC8KXs-M4fWmDADmG414V0_X0TfjrYu01nDH2Dcf3TIOFbRDb993g8nOCswLMi92LwjoqhYnFdf1jzgK0",
}


class LitterboxAsync:
    def __init__(self) -> None:
        self.headers: dict[str, str] = http_headers
        self.api_path = litterbox_api_url

    def generateRandomFileName(self, fileExt: Optional[str] = None) -> str:
        """
        Generate Random Filename

        Args:
            fileExt (Optional[str]): File extension to use for naming
        Returns:
            str
        """
        if not fileExt:
            fileExt = "png"
        return f"{random.getrandbits(32)}.{fileExt}"

    async def upload(
        self, file: Union[str, bytes, BufferedReader], time="1h"
    ) -> Optional[str]:
        """
        Upload File to Litterbox

        Args:
            file (Union[str, bytes, BufferedReader]): File to upload (accepts either filepath or io.BufferedReader)
            time (str): Expiration time, default: 1h
        Returns:
            Optional[str]
        """
        try:
            if not file:
                return None
            if isinstance(file, str):
                if not os.path.exists(file):
                    logging.critical("Could not find %s", file)
                    return None
            if isinstance(file, BufferedReader):
                file = file.read()

            fileExt: str = "png"
            if isinstance(file, str):
                if file.find(".") > 0:
                    fileExt = "".join(file.split(".")[-1:])

                file = open(file, "rb").read()

            with magic.Magic(flags=magic.MAGIC_MIME) as m:
                if isinstance(file, BufferedReader):
                    content_type = str(m.id_buffer(file))
                else:
                    content_type = str(m.id_filename(file))

            post_data: FormData = FormData()
            post_data.add_field(name="reqtype", value="fileupload")
            post_data.add_field(name="userhash", value="")
            post_data.add_field(name="time", value=time)

            post_data.add_field(
                name="fileToUpload",
                value=file,
                filename=self.generateRandomFileName(fileExt),
                content_type=content_type,
            )
            async with ClientSession() as session:
                async with await session.post(
                    self.api_path,
                    headers=self.headers,
                    data=post_data,
                    timeout=ClientTimeout(connect=5, sock_read=70),
                ) as request:
                    request.raise_for_status()
                    return await request.text()
        except:
            traceback.print_exc()
            return None
        finally:
            if isinstance(file, BufferedReader):
                try:
                    file.close()
                except:
                    return None