mirror of
https://github.com/flibusta-apps/telegram_files_server.git
synced 2025-12-06 12:35:39 +01:00
Update uploading/downloading
This commit is contained in:
@@ -1,61 +1,60 @@
|
||||
from io import BytesIO
|
||||
from typing import Optional
|
||||
|
||||
from telegram_files_storage import AiogramFilesStorage, TelethonFilesStorage
|
||||
|
||||
from app.services.storages import StoragesContainer
|
||||
from app.models import UploadBackends
|
||||
from app.services.storages import StoragesContainer, BotStorage, UserStorage
|
||||
|
||||
|
||||
class FileDownloader:
|
||||
_aiogram_storage_index = 0
|
||||
_telethon_storage_index = 0
|
||||
_bot_storage_index = 0
|
||||
_user_storage_index = 0
|
||||
|
||||
@classmethod
|
||||
@property
|
||||
def AIOGRAM_STORAGES(cls) -> list[AiogramFilesStorage]:
|
||||
return StoragesContainer.AIOGRAM_STORAGES
|
||||
def bot_storages(cls) -> list[BotStorage]:
|
||||
return StoragesContainer.BOT_STORAGES
|
||||
|
||||
@classmethod
|
||||
@property
|
||||
def TELETHON_STORAGES(cls) -> list[TelethonFilesStorage]:
|
||||
return StoragesContainer.TELETHON_STORAGES
|
||||
def user_storages(cls) -> list[UserStorage]:
|
||||
return StoragesContainer.USER_STORAGES
|
||||
|
||||
@classmethod
|
||||
def get_aiogram_storage(cls) -> AiogramFilesStorage:
|
||||
if not cls.AIOGRAM_STORAGES:
|
||||
def get_bot_storage(cls) -> BotStorage:
|
||||
if not cls.bot_storages:
|
||||
raise ValueError("Aiogram storage not exist!")
|
||||
|
||||
cls._aiogram_storage_index = (cls._aiogram_storage_index + 1) % len(
|
||||
cls.AIOGRAM_STORAGES
|
||||
)
|
||||
bot_storages: list[BotStorage] = cls.bot_storages # type: ignore
|
||||
|
||||
return cls.AIOGRAM_STORAGES[cls._aiogram_storage_index]
|
||||
cls._bot_storage_index = (cls._bot_storage_index + 1) % len(bot_storages)
|
||||
|
||||
return bot_storages[cls._bot_storage_index]
|
||||
|
||||
@classmethod
|
||||
def get_telethon_storage(cls) -> TelethonFilesStorage:
|
||||
if not cls.TELETHON_STORAGES:
|
||||
def get_user_storage(cls) -> UserStorage:
|
||||
if not cls.user_storages:
|
||||
raise ValueError("Telethon storage not exists!")
|
||||
|
||||
cls._telethon_storage_index = (cls._telethon_storage_index + 1) % len(
|
||||
cls.TELETHON_STORAGES
|
||||
)
|
||||
user_storages: list[UserStorage] = cls.user_storages # type: ignore
|
||||
|
||||
return cls.TELETHON_STORAGES[cls._telethon_storage_index]
|
||||
cls._user_storage_index = (cls._user_storage_index + 1) % len(user_storages)
|
||||
|
||||
return user_storages[cls._user_storage_index]
|
||||
|
||||
@classmethod
|
||||
async def download_by_file_id(cls, file_id: str) -> Optional[BytesIO]:
|
||||
if not cls.AIOGRAM_STORAGES:
|
||||
return None
|
||||
|
||||
storage = cls.get_aiogram_storage()
|
||||
|
||||
return await storage.download(file_id)
|
||||
|
||||
@classmethod
|
||||
async def download_by_message_id(cls, message_id: int) -> Optional[BytesIO]:
|
||||
if not cls.TELETHON_STORAGES:
|
||||
return None
|
||||
|
||||
storage = cls.get_telethon_storage()
|
||||
async def _download_via(cls, message_id: int, storage_type: UploadBackends):
|
||||
if storage_type == UploadBackends.bot:
|
||||
storage = cls.get_bot_storage()
|
||||
else:
|
||||
storage = cls.get_user_storage()
|
||||
|
||||
return await storage.download(message_id)
|
||||
|
||||
@classmethod
|
||||
async def download_by_message_id(cls, message_id: int):
|
||||
if not cls.bot_storages and not cls.user_storages:
|
||||
raise ValueError("Files storage not exist!")
|
||||
|
||||
if (
|
||||
data := await cls._download_via(message_id, UploadBackends.bot)
|
||||
) is not None:
|
||||
return data
|
||||
|
||||
return await cls._download_via(message_id, UploadBackends.user)
|
||||
|
||||
@@ -3,25 +3,23 @@ from typing import Optional
|
||||
|
||||
from fastapi import UploadFile
|
||||
|
||||
from telegram_files_storage import AiogramFilesStorage, TelethonFilesStorage
|
||||
|
||||
from app.models import UploadedFile, UploadBackends
|
||||
from app.services.storages import StoragesContainer
|
||||
from app.services.storages import StoragesContainer, BotStorage, UserStorage
|
||||
|
||||
|
||||
class FileUploader:
|
||||
_aiogram_storage_index = 0
|
||||
_telethon_storage_index = 0
|
||||
_bot_storage_index = 0
|
||||
_user_storage_index = 0
|
||||
|
||||
@classmethod
|
||||
@property
|
||||
def AIOGRAM_STORAGES(cls) -> list[AiogramFilesStorage]:
|
||||
return StoragesContainer.AIOGRAM_STORAGES
|
||||
def bot_storages(cls) -> list[BotStorage]:
|
||||
return StoragesContainer.BOT_STORAGES
|
||||
|
||||
@classmethod
|
||||
@property
|
||||
def TELETHON_STORAGES(cls) -> list[TelethonFilesStorage]:
|
||||
return StoragesContainer.TELETHON_STORAGES
|
||||
def user_storages(cls) -> list[UserStorage]:
|
||||
return StoragesContainer.USER_STORAGES
|
||||
|
||||
def __init__(self, file: UploadFile, caption: Optional[str] = None) -> None:
|
||||
self.file = file
|
||||
@@ -31,38 +29,16 @@ class FileUploader:
|
||||
self.upload_backend: Optional[UploadBackends] = None
|
||||
|
||||
async def _upload(self) -> bool:
|
||||
if not self.AIOGRAM_STORAGES and not self.TELETHON_STORAGES:
|
||||
if not self.bot_storages and not self.user_storages:
|
||||
raise ValueError("Files storage not exist!")
|
||||
|
||||
if await self._upload_via_aiogram():
|
||||
if await self._upload_via(UploadBackends.bot):
|
||||
return True
|
||||
|
||||
return await self._upload_via_telethon()
|
||||
return await self._upload_via(UploadBackends.user)
|
||||
|
||||
async def _upload_via_aiogram(self) -> bool:
|
||||
if not self.AIOGRAM_STORAGES:
|
||||
return False
|
||||
|
||||
data = await self.file.read()
|
||||
|
||||
if isinstance(data, str):
|
||||
data = data.encode()
|
||||
|
||||
if len(data) > 50 * 1000 * 1000:
|
||||
return False
|
||||
|
||||
bytes_io = BytesIO(data)
|
||||
bytes_io.name = self.file.filename
|
||||
|
||||
storage = self.get_aiogram_storage()
|
||||
|
||||
self.upload_data = await storage.upload(bytes_io, self.caption) # type: ignore
|
||||
self.upload_backend = UploadBackends.aiogram
|
||||
|
||||
return True
|
||||
|
||||
async def _upload_via_telethon(self) -> bool:
|
||||
if not self.TELETHON_STORAGES:
|
||||
async def _upload_via(self, storage_type: UploadBackends) -> bool:
|
||||
if not self.bot_storages:
|
||||
return False
|
||||
|
||||
data = await self.file.read()
|
||||
@@ -73,12 +49,18 @@ class FileUploader:
|
||||
bytes_io = BytesIO(data)
|
||||
bytes_io.name = self.file.filename
|
||||
|
||||
storage = self.get_telethon_storage()
|
||||
if storage_type == UploadBackends.bot:
|
||||
storage = self.get_bot_storage()
|
||||
else:
|
||||
storage = self.get_user_storage()
|
||||
|
||||
self.upload_data = await storage.upload(
|
||||
bytes_io, caption=self.caption
|
||||
) # type: ignore
|
||||
self.upload_backend = UploadBackends.telethon
|
||||
data = await storage.upload(bytes_io, caption=self.caption) # type: ignore
|
||||
|
||||
if not data:
|
||||
return False
|
||||
|
||||
self.upload_data = {"chat_id": data[0], "message_id": data[1]}
|
||||
self.upload_backend = storage_type
|
||||
|
||||
return True
|
||||
|
||||
@@ -89,26 +71,26 @@ class FileUploader:
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_aiogram_storage(cls) -> AiogramFilesStorage:
|
||||
if not cls.AIOGRAM_STORAGES:
|
||||
def get_bot_storage(cls) -> BotStorage:
|
||||
if not cls.bot_storages:
|
||||
raise ValueError("Aiogram storage not exist!")
|
||||
|
||||
cls._aiogram_storage_index = (cls._aiogram_storage_index + 1) % len(
|
||||
cls.AIOGRAM_STORAGES
|
||||
)
|
||||
bot_storages: list[BotStorage] = cls.bot_storages # type: ignore
|
||||
|
||||
return cls.AIOGRAM_STORAGES[cls._aiogram_storage_index]
|
||||
cls._bot_storage_index = (cls._bot_storage_index + 1) % len(bot_storages)
|
||||
|
||||
return bot_storages[cls._bot_storage_index]
|
||||
|
||||
@classmethod
|
||||
def get_telethon_storage(cls) -> TelethonFilesStorage:
|
||||
if not cls.TELETHON_STORAGES:
|
||||
def get_user_storage(cls) -> UserStorage:
|
||||
if not cls.user_storages:
|
||||
raise ValueError("Telethon storage not exists!")
|
||||
|
||||
cls._telethon_storage_index = (cls._telethon_storage_index + 1) % len(
|
||||
cls.TELETHON_STORAGES
|
||||
)
|
||||
user_storages: list[UserStorage] = cls.user_storages # type: ignore
|
||||
|
||||
return cls.TELETHON_STORAGES[cls._telethon_storage_index]
|
||||
cls._user_storage_index = (cls._user_storage_index + 1) % len(user_storages)
|
||||
|
||||
return user_storages[cls._user_storage_index]
|
||||
|
||||
@classmethod
|
||||
async def upload(
|
||||
|
||||
@@ -1,23 +1,116 @@
|
||||
from telegram_files_storage import AiogramFilesStorage, TelethonFilesStorage
|
||||
import abc
|
||||
from io import BytesIO
|
||||
from typing import AsyncIterator, Union, Optional
|
||||
|
||||
import telethon.client
|
||||
import telethon.errors
|
||||
import telethon.tl.types
|
||||
|
||||
from core.config import env_config
|
||||
|
||||
|
||||
class BaseStorage(abc.ABC):
|
||||
def __init__(
|
||||
self, channel_id: Union[str, int], app_id: int, api_hash: str, session: str
|
||||
):
|
||||
self.channel_id = channel_id
|
||||
|
||||
self.client = telethon.client.TelegramClient(session, app_id, api_hash)
|
||||
|
||||
self.ready = False
|
||||
|
||||
async def prepare(self):
|
||||
...
|
||||
|
||||
async def upload(
|
||||
self, file: BytesIO, caption: Optional[str] = None
|
||||
) -> Optional[tuple[Union[str, int], int]]:
|
||||
message = await self.client.send_file(
|
||||
self.channel_id, file=file, caption=caption
|
||||
)
|
||||
|
||||
if not message.media:
|
||||
return None
|
||||
|
||||
return self.channel_id, message.id
|
||||
|
||||
async def download(self, message_id: int) -> Optional[AsyncIterator[bytes]]:
|
||||
messages = await self.client.get_messages(self.channel_id, ids=[message_id])
|
||||
|
||||
if not messages:
|
||||
return None
|
||||
|
||||
message: telethon.tl.types.Message = messages[0]
|
||||
|
||||
if message.media is None:
|
||||
return None
|
||||
|
||||
return self.client.iter_download(message.media)
|
||||
|
||||
|
||||
class UserStorage(BaseStorage):
|
||||
async def prepare(self):
|
||||
if self.ready:
|
||||
return
|
||||
|
||||
await self.client.start() # type: ignore
|
||||
|
||||
if not await self.client.is_user_authorized():
|
||||
await self.client.sign_in()
|
||||
try:
|
||||
await self.client.sign_in(code=input("Enter code: "))
|
||||
except telethon.errors.SessionPasswordNeededError:
|
||||
await self.client.sign_in(password=input("Enter password: "))
|
||||
|
||||
self.ready = True
|
||||
|
||||
|
||||
class BotStorage(BaseStorage):
|
||||
def __init__(
|
||||
self,
|
||||
channel_id: Union[str, int],
|
||||
app_id: int,
|
||||
api_hash: str,
|
||||
session: str,
|
||||
token: str,
|
||||
) -> None:
|
||||
super().__init__(channel_id, app_id, api_hash, session)
|
||||
|
||||
self.token = token
|
||||
|
||||
async def prepare(self):
|
||||
if self.ready:
|
||||
return
|
||||
|
||||
await self.client.start(bot_token=self.token) # type: ignore
|
||||
|
||||
self.ready = True
|
||||
|
||||
|
||||
class StoragesContainer:
|
||||
AIOGRAM_STORAGES: list[AiogramFilesStorage] = []
|
||||
TELETHON_STORAGES: list[TelethonFilesStorage] = []
|
||||
BOT_STORAGES: list[BotStorage] = []
|
||||
USER_STORAGES: list[UserStorage] = []
|
||||
|
||||
@classmethod
|
||||
async def prepare(cls):
|
||||
if not env_config.TELETHON_APP_CONFIG:
|
||||
return
|
||||
|
||||
if env_config.BOT_TOKENS:
|
||||
cls.AIOGRAM_STORAGES: list[AiogramFilesStorage] = [
|
||||
AiogramFilesStorage(env_config.TELEGRAM_CHAT_ID, token)
|
||||
cls.BOT_STORAGES: list[BotStorage] = [
|
||||
BotStorage(
|
||||
env_config.TELEGRAM_CHAT_ID,
|
||||
env_config.TELETHON_APP_CONFIG.APP_ID,
|
||||
env_config.TELETHON_APP_CONFIG.API_HASH,
|
||||
token.split(":")[0],
|
||||
token,
|
||||
)
|
||||
for token in env_config.BOT_TOKENS
|
||||
]
|
||||
|
||||
if env_config.TELETHON_APP_CONFIG and env_config.TELETHON_SESSIONS:
|
||||
cls.TELETHON_STORAGES: list[TelethonFilesStorage] = [
|
||||
TelethonFilesStorage(
|
||||
if env_config.TELETHON_SESSIONS:
|
||||
cls.USER_STORAGES: list[UserStorage] = [
|
||||
UserStorage(
|
||||
env_config.TELEGRAM_CHAT_ID,
|
||||
env_config.TELETHON_APP_CONFIG.APP_ID,
|
||||
env_config.TELETHON_APP_CONFIG.API_HASH,
|
||||
@@ -26,5 +119,5 @@ class StoragesContainer:
|
||||
for session in env_config.TELETHON_SESSIONS
|
||||
]
|
||||
|
||||
for storage in [*cls.AIOGRAM_STORAGES, *cls.TELETHON_STORAGES]:
|
||||
for storage in [*cls.BOT_STORAGES, *cls.USER_STORAGES]:
|
||||
await storage.prepare()
|
||||
|
||||
Reference in New Issue
Block a user