diff --git a/modules/api_client.py b/modules/api_client.py index e8fba39..f5ac191 100644 --- a/modules/api_client.py +++ b/modules/api_client.py @@ -7,11 +7,13 @@ except ModuleNotFoundError: import asyncio from base64 import b64decode, b64encode +from traceback import print_exc +import aiofiles from ujson import dumps from os import makedirs, path, sep from random import choice from typing import Tuple, Union -from aiohttp import ClientSession +from aiohttp import ClientSession, FormData from classes.exceptions import SubmissionUploadError @@ -25,8 +27,8 @@ http_session = ClientSession(json_serialize=dumps, ) async def authorize() -> str: makedirs(configGet("cache", "locations"), exist_ok=True) if path.exists(configGet("cache", "locations")+sep+"api_access") is True: - with open(configGet("cache", "locations")+sep+"api_access", "rb") as file: - token = b64decode(file.read()).decode("utf-8") + async with aiofiles.open(configGet("cache", "locations")+sep+"api_access", "rb") as file: + token = b64decode(await file.read()).decode("utf-8") if (await http_session.get(configGet("address", "posting", "api")+"/users/me/", headers={"Authorization": f"Bearer {token}"})).status == 200: return token payload = { @@ -36,11 +38,11 @@ async def authorize() -> str: "password": configGet("password", "posting", "api") } response = await http_session.post(configGet("address", "posting", "api")+"/token", data=payload) - if response.status != 200: + if not response.ok: logWrite(f'Incorrect API credentials! Could not login into "{configGet("address", "posting", "api")}" using login "{configGet("username", "posting", "api")}": HTTP {response.status}') raise ValueError - with open(configGet("cache", "locations")+sep+"api_access", "wb") as file: - file.write(b64encode((await response.json())["access_token"].encode("utf-8"))) + async with aiofiles.open(configGet("cache", "locations")+sep+"api_access", "wb") as file: + await file.write(b64encode((await response.json())["access_token"].encode("utf-8"))) return (await response.json())["access_token"] async def random_pic(token: Union[str, None] = None) -> Tuple[str, str]: @@ -66,8 +68,17 @@ async def upload_pic(filepath: str, allow_duplicates: bool = False, token: Union token = await authorize() if token is None else token try: pic_name = path.basename(filepath) - files = {'file': (pic_name, open(filepath, 'rb'), 'image/jpeg')} - response = await http_session.post(f'{configGet("address", "posting", "api")}/albums/{configGet("album", "posting", "api")}/photos', params={"caption": "queue", "compress": False, "ignore_duplicates": allow_duplicates}, headers={"Authorization": f"Bearer {token}"}, files=files) + logWrite(f"Uploading {pic_name} to the API...", debug=True) + async with aiofiles.open(filepath, "rb") as f: + file_bytes = await f.read() + formdata = FormData() + formdata.add_field('file', file_bytes, filename=pic_name, content_type='image/jpeg') + response = await http_session.post( + f'{configGet("address", "posting", "api")}/albums/{configGet("album", "posting", "api")}/photos', + params={"caption": "queue", "compress": "false", "ignore_duplicates": str(allow_duplicates).lower()}, + headers={"Authorization": f"Bearer {token}"}, + data=formdata + ) if response.status != 200 and response.status != 409: logWrite(f"Could not upload '{filepath}' to API: HTTP {response.status} with message '{response.content}'") raise SubmissionUploadError(str(filepath), response.status, response.content) @@ -79,7 +90,8 @@ async def upload_pic(filepath: str, allow_duplicates: bool = False, token: Union else: duplicates.append(f'`{duplicate["id"]}`:\n{configGet("address_external", "posting", "api")}/token/photo/{(await response.json())["access_token"]}?id={index}') return True, duplicates - except: + except Exception as exp: + print_exc() return False, [] async def find_pic(name: str, caption: Union[str, None] = None, token: Union[str, None] = None) -> Union[dict, None]: diff --git a/modules/sender.py b/modules/sender.py index bd6b425..16136c4 100644 --- a/modules/sender.py +++ b/modules/sender.py @@ -1,9 +1,10 @@ from datetime import datetime from os import makedirs, path -from shutil import copyfileobj, rmtree +from shutil import rmtree from traceback import format_exc from uuid import uuid4 from PIL import Image +import aiofiles from classes.poster_client import PosterClient @@ -35,7 +36,7 @@ async def send_content(app: PosterClient) -> None: await app.send_message(app.owner, locale("api_queue_error", "message", locale=configGet("locale"))) return - response = await http_session.get(f'{configGet("address", "posting", "api")}/photos/{pic[0]}', headers={"Authorization": f"Bearer {token}"}, stream=True) + response = await http_session.get(f'{configGet("address", "posting", "api")}/photos/{pic[0]}', headers={"Authorization": f"Bearer {token}"}) if response.status != 200: logWrite(locale("post_invalid_pic", "console", locale=configGet("locale")).format(response.status, str(response.json()))) @@ -48,8 +49,8 @@ async def send_content(app: PosterClient) -> None: tmp_path = path.join(tmp_dir, pic[1]) - with open(path.join(configGet("tmp", "locations"), tmp_path), 'wb') as out_file: - out_file.write(await response.read()) + async with aiofiles.open(path.join(configGet("tmp", "locations"), tmp_path), 'wb') as out_file: + await out_file.write(await response.read()) logWrite(f'Candidate {pic[1]} ({pic[0]}) is {path.getsize(path.join(configGet("tmp", "locations"), tmp_path))} bytes big', debug=True) diff --git a/requirements.txt b/requirements.txt index 3bab7d8..27f8e1a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,6 +1,7 @@ apscheduler~=3.10.0 pyrogram~=2.0.100 aiohttp~=3.8.4 +aiofiles~=23.1.0 psutil~=5.9.4 pymongo~=4.3.3 pillow~=9.4.0