244 lines
9.3 KiB
Python
244 lines
9.3 KiB
Python
import asyncio
|
|
from glob import iglob
|
|
from os import getcwd, makedirs, path, remove
|
|
from shutil import disk_usage, rmtree
|
|
from traceback import format_exc
|
|
from uuid import uuid4
|
|
from zipfile import ZipFile
|
|
import aiofiles
|
|
|
|
from convopyro import listen_message
|
|
from pyrogram import filters
|
|
from pyrogram.types import Message
|
|
|
|
from classes.poster_client import PosterClient
|
|
from modules.app import app, users_with_context
|
|
from modules.logger import logWrite
|
|
from modules.utils import configGet, extract_and_save, locale
|
|
from modules.api_client import (
|
|
client,
|
|
photo_upload,
|
|
photo_delete,
|
|
HTTPValidationError,
|
|
BodyPhotoUploadAlbumsAlbumPhotosPost,
|
|
File,
|
|
)
|
|
|
|
|
|
@app.on_message(~filters.scheduled & filters.command(["import"], prefixes=["", "/"]))
|
|
async def cmd_import(app: PosterClient, msg: Message):
|
|
if msg.from_user.id in app.admins:
|
|
global users_with_context
|
|
if msg.from_user.id not in users_with_context:
|
|
users_with_context.append(msg.from_user.id)
|
|
else:
|
|
return
|
|
await msg.reply_text(
|
|
locale("import_request", "message", locale=msg.from_user.language_code)
|
|
)
|
|
answer = await listen_message(app, msg.chat.id, timeout=600)
|
|
users_with_context.remove(msg.from_user.id)
|
|
if answer is None:
|
|
await msg.reply_text(
|
|
locale("import_ignored", "message", locale=msg.from_user.language_code),
|
|
quote=True,
|
|
)
|
|
return
|
|
if answer.text == "/cancel":
|
|
await answer.reply_text(
|
|
locale("import_abort", "message", locale=msg.from_user.language_code)
|
|
)
|
|
return
|
|
if answer.document is None:
|
|
await answer.reply_text(
|
|
locale(
|
|
"import_invalid_media",
|
|
"message",
|
|
locale=msg.from_user.language_code,
|
|
),
|
|
quote=True,
|
|
)
|
|
return
|
|
if answer.document.mime_type != "application/zip":
|
|
await answer.reply_text(
|
|
locale(
|
|
"import_invalid_mime", "message", locale=msg.from_user.language_code
|
|
),
|
|
quote=True,
|
|
)
|
|
return
|
|
if disk_usage(getcwd())[2] < (answer.document.file_size) * 3:
|
|
await msg.reply_text(
|
|
locale(
|
|
"import_too_big", "message", locale=msg.from_user.language_code
|
|
).format(
|
|
answer.document.file_size // (2**30),
|
|
disk_usage(getcwd())[2] // (2**30),
|
|
)
|
|
)
|
|
return
|
|
tmp_dir = str(uuid4())
|
|
logWrite(
|
|
f"Importing '{answer.document.file_name}' file {answer.document.file_size} bytes big (TMP ID {tmp_dir})"
|
|
)
|
|
makedirs(path.join(configGet("tmp", "locations"), tmp_dir), exist_ok=True)
|
|
tmp_path = path.join(configGet("tmp", "locations"), answer.document.file_id)
|
|
downloading = await answer.reply_text(
|
|
locale("import_downloading", "message", locale=msg.from_user.language_code),
|
|
quote=True,
|
|
)
|
|
await app.download_media(answer, file_name=tmp_path)
|
|
await downloading.edit(
|
|
locale("import_unpacking", "message", locale=msg.from_user.language_code)
|
|
)
|
|
try:
|
|
with ZipFile(tmp_path, "r") as handle:
|
|
tasks = [
|
|
extract_and_save(
|
|
handle, name, path.join(configGet("tmp", "locations"), tmp_dir)
|
|
)
|
|
for name in handle.namelist()
|
|
]
|
|
_ = await asyncio.gather(*tasks)
|
|
except Exception as exp:
|
|
logWrite(
|
|
f"Could not import '{answer.document.file_name}' due to {exp}: {format_exc}"
|
|
)
|
|
await answer.reply_text(
|
|
locale(
|
|
"import_unpack_error", "message", locale=msg.from_user.language_code
|
|
).format(exp, format_exc())
|
|
)
|
|
return
|
|
logWrite(f"Downloaded '{answer.document.file_name}' - awaiting upload")
|
|
await downloading.edit(
|
|
locale("import_uploading", "message", locale=msg.from_user.language_code)
|
|
)
|
|
remove(tmp_path)
|
|
|
|
for filename in iglob(
|
|
path.join(configGet("tmp", "locations"), tmp_dir) + "**/**", recursive=True
|
|
):
|
|
if not path.isfile(filename):
|
|
continue
|
|
# upload filename
|
|
async with aiofiles.open(filename, "rb") as f:
|
|
file_bytes = await f.read()
|
|
uploaded = await photo_upload(
|
|
album=configGet("album", "posting", "api"),
|
|
client=client,
|
|
multipart_data=BodyPhotoUploadAlbumsAlbumPhotosPost(
|
|
File(file_bytes, path.basename(filename), "image/jpeg")
|
|
),
|
|
ignore_duplicates=configGet("allow_duplicates", "submission"),
|
|
caption="queue",
|
|
compress=False,
|
|
)
|
|
if uploaded is None or isinstance(uploaded, HTTPValidationError):
|
|
await msg.reply_text(
|
|
locale(
|
|
"import_upload_error_other",
|
|
"message",
|
|
locale=msg.from_user.language_code,
|
|
).format(path.basename(filename)),
|
|
disable_notification=True,
|
|
)
|
|
return
|
|
if not hasattr(uploaded, "id"):
|
|
if hasattr(uploaded, "duplicates"):
|
|
logWrite(
|
|
f"Could not upload '{filename}' from '{path.join(configGet('tmp', 'locations'), tmp_dir)}'. Duplicates: {str(uploaded.duplicates)}",
|
|
debug=True,
|
|
)
|
|
await msg.reply_text(
|
|
locale(
|
|
"import_upload_error_duplicate",
|
|
"message",
|
|
locale=msg.from_user.language_code,
|
|
).format(path.basename(filename)),
|
|
disable_notification=True,
|
|
)
|
|
else:
|
|
await msg.reply_text(
|
|
locale(
|
|
"import_upload_error_other",
|
|
"message",
|
|
locale=msg.from_user.language_code,
|
|
).format(path.basename(filename)),
|
|
disable_notification=True,
|
|
)
|
|
else:
|
|
logWrite(
|
|
f"Uploaded '{filename}' from '{path.join(configGet('tmp', 'locations'), tmp_dir)}' and got ID {uploaded.id}",
|
|
debug=True,
|
|
)
|
|
|
|
await downloading.delete()
|
|
logWrite(
|
|
f"Removing '{path.join(configGet('tmp', 'locations'), tmp_dir)}' after uploading",
|
|
debug=True,
|
|
)
|
|
rmtree(path.join(configGet("tmp", "locations"), tmp_dir), ignore_errors=True)
|
|
await answer.reply_text(
|
|
locale("import_finished", "message", locale=msg.from_user.language_code),
|
|
quote=True,
|
|
)
|
|
return
|
|
|
|
|
|
@app.on_message(~filters.scheduled & filters.command(["export"], prefixes=["", "/"]))
|
|
async def cmd_export(app: PosterClient, msg: Message):
|
|
if msg.from_user.id in app.admins:
|
|
pass
|
|
|
|
|
|
@app.on_message(~filters.scheduled & filters.command(["remove"], prefixes=["", "/"]))
|
|
async def cmd_remove(app: PosterClient, msg: Message):
|
|
if msg.from_user.id in app.admins:
|
|
global users_with_context
|
|
if msg.from_user.id not in users_with_context:
|
|
users_with_context.append(msg.from_user.id)
|
|
else:
|
|
return
|
|
await msg.reply_text(
|
|
locale("remove_request", "message", locale=msg.from_user.language_code)
|
|
)
|
|
answer = await listen_message(app, msg.chat.id, timeout=600)
|
|
users_with_context.remove(msg.from_user.id)
|
|
if answer is None:
|
|
await msg.reply_text(
|
|
locale("remove_ignored", "message", locale=msg.from_user.language_code),
|
|
quote=True,
|
|
)
|
|
return
|
|
if answer.text == "/cancel":
|
|
await answer.reply_text(
|
|
locale("remove_abort", "message", locale=msg.from_user.language_code)
|
|
)
|
|
return
|
|
try:
|
|
response = await photo_delete(id=answer.text, client=client)
|
|
logWrite(
|
|
f"Removed '{answer.text}' by request of user {answer.from_user.id}"
|
|
)
|
|
await answer.reply_text(
|
|
locale(
|
|
"remove_success", "message", locale=msg.from_user.language_code
|
|
).format(answer.text)
|
|
)
|
|
except:
|
|
logWrite(
|
|
f"Could not remove '{answer.text}' by request of user {answer.from_user.id}"
|
|
)
|
|
await answer.reply_text(
|
|
locale(
|
|
"remove_failure", "message", locale=msg.from_user.language_code
|
|
).format(answer.text)
|
|
)
|
|
|
|
|
|
@app.on_message(~filters.scheduled & filters.command(["purge"], prefixes=["", "/"]))
|
|
async def cmd_purge(app: PosterClient, msg: Message):
|
|
if msg.from_user.id in app.admins:
|
|
pass
|