288 lines
10 KiB
Python
288 lines
10 KiB
Python
import contextlib
|
|
import logging
|
|
from datetime import datetime
|
|
from io import BytesIO
|
|
from os import makedirs, remove, sep
|
|
from pathlib import Path
|
|
from shutil import rmtree
|
|
from time import time
|
|
from traceback import format_exc
|
|
from typing import Dict, List, Tuple, Union
|
|
|
|
import aiofiles
|
|
from aiohttp import ClientSession
|
|
from apscheduler.schedulers.asyncio import AsyncIOScheduler
|
|
from bson import ObjectId
|
|
from libbot import json_write
|
|
from libbot.i18n.sync import _
|
|
from libbot.pyrogram.classes import PyroClient
|
|
from photosapi_client.errors import UnexpectedStatus
|
|
from pyrogram.errors import bad_request_400
|
|
from pyrogram.types import Message, User
|
|
from pytimeparse.timeparse import timeparse
|
|
from ujson import dumps, loads
|
|
|
|
from classes.enums.submission_types import SubmissionType
|
|
from classes.exceptions import (
|
|
SubmissionDuplicatesError,
|
|
SubmissionUnavailableError,
|
|
SubmissionUnsupportedError,
|
|
)
|
|
from classes.pyrouser import PyroUser
|
|
from modules.api_client import (
|
|
BodyPhotoUpload,
|
|
BodyVideoUpload,
|
|
File,
|
|
Photo,
|
|
Video,
|
|
client,
|
|
photo_upload,
|
|
video_upload,
|
|
)
|
|
from modules.database import col_submitted, col_users
|
|
from modules.http_client import http_session
|
|
from modules.sender import send_content
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class PyroClient(PyroClient):
|
|
def __init__(self, scheduler: AsyncIOScheduler):
|
|
super().__init__(locales_root=Path("locale"), scheduler=scheduler)
|
|
|
|
self.version: float = 0.3
|
|
|
|
self.owner: int = self.config["bot"]["owner"]
|
|
self.admins: List[int] = self.config["bot"]["admins"] + [
|
|
self.config["bot"]["owner"]
|
|
]
|
|
|
|
self.sender_session: Union[ClientSession, None] = None
|
|
|
|
self.scopes_placeholders: Dict[str, int] = {
|
|
"owner": self.owner,
|
|
"comments": self.config["posting"]["comments"],
|
|
}
|
|
|
|
async def start(self):
|
|
await super().start()
|
|
|
|
if self.sender_session is None:
|
|
self.sender_session = ClientSession()
|
|
|
|
if self.config["reports"]["update"]:
|
|
try:
|
|
async with ClientSession(
|
|
json_serialize=dumps,
|
|
) as http_session:
|
|
check_update = await http_session.get(
|
|
"https://git.end-play.xyz/api/v1/repos/profitroll/TelegramPoster/releases?page=1&limit=1"
|
|
)
|
|
|
|
response = await check_update.json()
|
|
|
|
if len(response) == 0:
|
|
raise ValueError("No bot releases on git found.")
|
|
|
|
if float(response[0]["tag_name"].replace("v", "")) > self.version:
|
|
logger.info(
|
|
"New version %s found (current %s)",
|
|
response[0]["tag_name"].replace("v", ""),
|
|
self.version,
|
|
)
|
|
await self.send_message(
|
|
self.owner,
|
|
self._(
|
|
"update_available",
|
|
"message",
|
|
).format(
|
|
response[0]["tag_name"],
|
|
response[0]["html_url"],
|
|
response[0]["body"],
|
|
),
|
|
)
|
|
else:
|
|
logger.info("No updates found, bot is up to date.")
|
|
except bad_request_400.PeerIdInvalid:
|
|
logger.warning(
|
|
"Could not send startup message to bot owner. Perhaps user has not started the bot yet."
|
|
)
|
|
except Exception as exc:
|
|
logger.exception("Update check failed due to %s: %s", exc, format_exc())
|
|
|
|
if self.config["mode"]["post"]:
|
|
if self.config["posting"]["use_interval"]:
|
|
self.scheduler.add_job(
|
|
send_content,
|
|
"interval",
|
|
seconds=timeparse(self.config["posting"]["interval"]),
|
|
args=[self, self.sender_session],
|
|
)
|
|
else:
|
|
for entry in self.config["posting"]["time"]:
|
|
dt_obj = datetime.strptime(entry, "%H:%M")
|
|
self.scheduler.add_job(
|
|
send_content,
|
|
"cron",
|
|
hour=dt_obj.hour,
|
|
minute=dt_obj.minute,
|
|
args=[self, self.sender_session],
|
|
)
|
|
|
|
async def stop(self):
|
|
makedirs(self.config["locations"]["cache"], exist_ok=True)
|
|
await json_write(
|
|
{"timestamp": time()},
|
|
Path(f"{self.config['locations']['cache']}/shutdown_time"),
|
|
)
|
|
|
|
await http_session.close()
|
|
|
|
if self.sender_session is not None:
|
|
await self.sender_session.close()
|
|
|
|
await super().stop()
|
|
|
|
async def submit_media(
|
|
self, id: str, purge_caption: bool = False
|
|
) -> Tuple[Union[Message, None], Union[str, None]]:
|
|
db_entry = await col_submitted.find_one({"_id": ObjectId(id)})
|
|
submission = None
|
|
|
|
if db_entry is None:
|
|
raise SubmissionUnavailableError()
|
|
|
|
if db_entry["temp"]["uuid"] is None:
|
|
try:
|
|
submission = await self.get_messages(
|
|
db_entry["user"], db_entry["telegram"]["msg_id"]
|
|
)
|
|
filepath = await self.download_media(
|
|
submission, file_name=self.config["locations"]["tmp"] + sep
|
|
)
|
|
except Exception as exc:
|
|
raise SubmissionUnavailableError() from exc
|
|
|
|
elif not Path(
|
|
f"{self.config['locations']['data']}/submissions/{db_entry['temp']['uuid']}/{db_entry['temp']['file']}",
|
|
).exists():
|
|
raise SubmissionUnavailableError()
|
|
else:
|
|
filepath = Path(
|
|
f"{self.config['locations']['data']}/submissions/{db_entry['temp']['uuid']}/{db_entry['temp']['file']}",
|
|
)
|
|
with contextlib.suppress(Exception):
|
|
submission = await self.get_messages(
|
|
db_entry["user"], db_entry["telegram"]["msg_id"]
|
|
)
|
|
|
|
async with aiofiles.open(str(filepath), "rb") as fh:
|
|
media_bytes = BytesIO(await fh.read())
|
|
|
|
try:
|
|
if db_entry["type"] == SubmissionType.PHOTO.value:
|
|
response = await photo_upload(
|
|
self.config["posting"]["api"]["album"],
|
|
client=client,
|
|
body=BodyPhotoUpload(
|
|
File(media_bytes, filepath.name, "image/jpeg")
|
|
),
|
|
ignore_duplicates=self.config["submission"]["allow_duplicates"],
|
|
compress=False,
|
|
caption="queue",
|
|
)
|
|
elif db_entry["type"] == SubmissionType.VIDEO.value:
|
|
response = await video_upload(
|
|
self.config["posting"]["api"]["album"],
|
|
client=client,
|
|
body=BodyVideoUpload(File(media_bytes, filepath.name, "video/*")),
|
|
caption="queue",
|
|
)
|
|
# elif db_entry["type"] == SubmissionType.ANIMATION.value:
|
|
# response = await video_upload(
|
|
# self.config["posting"]["api"]["album"],
|
|
# client=client,
|
|
# body=BodyVideoUpload(
|
|
# File(media_bytes, filepath.name, "video/*")
|
|
# ),
|
|
# caption="queue",
|
|
# )
|
|
except UnexpectedStatus as exc:
|
|
raise SubmissionUnsupportedError(str(filepath)) from exc
|
|
|
|
response_dict = (
|
|
{}
|
|
if not hasattr(response, "content")
|
|
else loads(response.content.decode("utf-8"))
|
|
)
|
|
|
|
if "duplicates" in response_dict and len(response_dict["duplicates"]) > 0:
|
|
duplicates = []
|
|
for index, duplicate in enumerate(response_dict["duplicates"]): # type: ignore
|
|
if response_dict["access_token"] is None:
|
|
duplicates.append(
|
|
f"`{duplicate['id']}`:\n{self.config['posting']['api']['address_external']}/photos/{duplicate['id']}"
|
|
)
|
|
else:
|
|
duplicates.append(
|
|
f"`{duplicate['id']}`:\n{self.config['posting']['api']['address_external']}/token/photo/{response_dict['access_token']}?id={index}"
|
|
)
|
|
raise SubmissionDuplicatesError(str(filepath), duplicates)
|
|
|
|
db_update = (
|
|
{"$set": {"done": True, "caption": None}}
|
|
if purge_caption
|
|
else {"$set": {"done": True}}
|
|
)
|
|
|
|
await col_submitted.update_one({"_id": ObjectId(id)}, db_update)
|
|
|
|
try:
|
|
if db_entry["temp"]["uuid"] is not None:
|
|
rmtree(
|
|
Path(
|
|
f"{self.config['locations']['data']}/submissions/{db_entry['temp']['uuid']}",
|
|
),
|
|
ignore_errors=True,
|
|
)
|
|
else:
|
|
remove(str(filepath))
|
|
except (FileNotFoundError, NotADirectoryError):
|
|
logger.error("Could not delete '%s' on submission accepted", filepath)
|
|
|
|
return (
|
|
submission,
|
|
response.parsed.id if hasattr(response, "parsed") else response.id,
|
|
)
|
|
|
|
async def find_user(self, user: Union[int, User]) -> PyroUser:
|
|
"""Find User by it's ID or User object
|
|
|
|
### Args:
|
|
* user (`Union[int, User]`): ID or User object to extract ID from
|
|
|
|
### Returns:
|
|
* `PyroUser`: PyroUser object
|
|
"""
|
|
if (
|
|
await col_users.find_one(
|
|
{"id": user.id if isinstance(user, User) else user}
|
|
) # type: ignore
|
|
is None
|
|
):
|
|
await col_users.insert_one(
|
|
{
|
|
"id": user.id if isinstance(user, User) else user,
|
|
"locale": user.language_code if isinstance(user, User) else None,
|
|
"banned": False,
|
|
"cooldown": datetime(1970, 1, 1, 0, 0),
|
|
"subscription": {"expires": datetime(1970, 1, 1, 0, 0)},
|
|
}
|
|
) # type: ignore
|
|
|
|
db_record = await col_users.find_one(
|
|
{"id": user.id if isinstance(user, User) else user}
|
|
) # type: ignore
|
|
|
|
return PyroUser(**db_record)
|