79 lines
2.3 KiB
Python
79 lines
2.3 KiB
Python
from pyrogram import Client
|
|
from pyrogram.errors import (
|
|
ChannelInvalid,
|
|
ChannelPrivate,
|
|
FileIdInvalid,
|
|
FileReferenceExpired,
|
|
FileReferenceInvalid,
|
|
PeerIdInvalid,
|
|
)
|
|
from pyrogram.types import ChatPhoto
|
|
|
|
from userbot.modules.avatars.repository import get_avatar_file, mark_avatar_downloaded
|
|
from userbot.modules.download import download_bytes
|
|
from userbot.modules.jobs.context import JobContext
|
|
from userbot.modules.jobs.registry import register
|
|
|
|
STALE_FILE_ID = (FileIdInvalid, FileReferenceExpired, FileReferenceInvalid)
|
|
UNKNOWN_OWNER = (PeerIdInvalid, ChannelInvalid, ChannelPrivate)
|
|
|
|
|
|
async def _fresh_file_id(client: Client, owner_id: int, unique_id: str) -> str | None:
|
|
try:
|
|
chat = await client.get_chat(owner_id)
|
|
except UNKNOWN_OWNER:
|
|
return None
|
|
photo = chat.photo if chat else None
|
|
if photo and photo.big_photo_unique_id == unique_id:
|
|
return photo.big_file_id
|
|
async for item in client.get_chat_photos(owner_id):
|
|
if isinstance(item, ChatPhoto) and item.big_photo_unique_id == unique_id:
|
|
return item.big_file_id
|
|
return None
|
|
|
|
|
|
async def _download(
|
|
client: Client, owner_id: int, unique_id: str, file_id: str
|
|
) -> bytes | None:
|
|
try:
|
|
return await download_bytes(client, file_id)
|
|
except STALE_FILE_ID:
|
|
fresh = await _fresh_file_id(client, owner_id, unique_id)
|
|
if fresh is None:
|
|
return None
|
|
return await download_bytes(client, fresh)
|
|
|
|
|
|
@register("fetch_avatar")
|
|
async def fetch_avatar(ctx: JobContext) -> None:
|
|
client = ctx.client
|
|
if client is None:
|
|
return
|
|
capture = getattr(client, "capture", None)
|
|
if capture is None:
|
|
return
|
|
owner_kind = ctx.job.params["owner_kind"]
|
|
owner_id = ctx.job.params["owner_id"]
|
|
unique_id = ctx.job.params["unique_id"]
|
|
found = await get_avatar_file(
|
|
ctx.pool, ctx.account_id, owner_kind, owner_id, unique_id
|
|
)
|
|
if found is None:
|
|
return
|
|
file_id, downloaded = found
|
|
if downloaded or file_id is None:
|
|
return
|
|
data = await _download(client, owner_id, unique_id, file_id)
|
|
if data is None:
|
|
return
|
|
storage_key = capture.storage.put(data)
|
|
await mark_avatar_downloaded(
|
|
ctx.pool,
|
|
ctx.account_id,
|
|
owner_kind,
|
|
owner_id,
|
|
unique_id,
|
|
storage_key,
|
|
len(data),
|
|
)
|