kulnoor/drive
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
main
drive/bot_handlers.py
T
gosangam 732a06f65d first commit
2026-10-18 23:48:22 +05:30

543 lines
25 KiB
Python
Raw Permalink Blame History

"""Telegram control commands (admin only), the /upload conversation, and the notifier.
The /upload conversation is persisted in the ``upload_tasks`` table (statuses
AWAITING_START -> AWAITING_END -> AWAITING_CONFIRMATION), keyed by the admin's
Telegram user ID, so every administrator has an independent setup and a restart
does not lose it.
"""
from __future__ import annotations
import logging
import re
from datetime import datetime, timedelta
from database import iso, utcnow
from errors import BotError
from link_parser import EXAMPLE, LinkError, parse_message_link
from models import ACTIVE_TASK_STATUSES, SETUP_TASK_STATUSES, State, TaskStatus, UploadTask, VideoRecord
from progress import HtmlReply, esc
from queue_manager import QueueManager, public_error
log = logging.getLogger(__name__)
HELP = """Commands (administrators only):
/upload - process a range of channel messages (asks for start and end links)
/confirm - start the prepared upload task
/cancel - cancel your upload setup or your active task (/cancel ID for a specific task)
/taskstatus - status of upload tasks (/taskstatus ID for one task)
/uploadhistory - recent upload tasks and their results
/status - current processing state
/progress - processed / pending / failed counts
/run - process every discovered channel video
/pause - pause at the next safe checkpoint
/resume - continue after a pause or /stop
/stop - checkpoint the current video and stop
/sync - discover new channel videos and queue them
/retry - re-queue failed items
/failed - list failed items with error summaries
/folders - topic folders and uploaded files
/check ID - details for one Telegram message ID
/help - this message"""
START_PROMPT = (
"📥 <b>Start Upload Task</b>\n\n"
"Please send the Telegram link of the first message you want to process.\n\n"
"Example:\nhttps://t.me/c/3909337424/14"
)
END_PROMPT = (
"📌 <b>Enter the End Link</b>\n\n"
"Now send the Telegram link of the last message you want to process.\n\n"
"Example:\nhttps://t.me/c/3909337424/20"
)
_STEP = {
TaskStatus.AWAITING_START: "waiting for the start link",
TaskStatus.AWAITING_END: "waiting for the end link",
TaskStatus.AWAITING_CONFIRMATION: "waiting for /confirm",
}
_COMMAND_RE = re.compile(r"^/(?P<cmd>[A-Za-z_]+)(?:@\w+)?(?:\s+(?P<arg>.*))?$", re.DOTALL)
_LINK_HINT = re.compile(r"(?:t\.me|telegram\.me|telegram\.dog)/", re.IGNORECASE)
def _short(text: str | None, limit: int = 160) -> str:
if not text:
return "-"
text = " ".join(text.split())
# never echo filesystem paths back to chat
text = re.sub(r"(?:/[^\s/]+)+/([^\s/]+)", r"…/\1", text)
return text if len(text) <= limit else text[: limit - 1] + "…"
class CommandRouter:
def __init__(self, config, db, queue: QueueManager):
self.config = config
self.db = db
self.queue = queue
def is_admin(self, user_id: int | None) -> bool:
return user_id is not None and user_id in self.config.admin_ids
async def handle(self, user_id: int | None, text: str) -> str | None:
stripped = (text or "").strip()
match = _COMMAND_RE.match(stripped)
if not match:
if not self.is_admin(user_id):
return None # ignore chatter from non-admins
try:
return await self.on_text(user_id, stripped)
except Exception as exc:
log.exception("Handling a message from %s failed", user_id)
return f"Could not process that message: {public_error(exc)}"
if not self.is_admin(user_id):
log.warning("Rejected command from non-admin user %s", user_id)
return "You are not authorized to control this bot."
command = match.group("cmd").lower()
handler = getattr(self, f"cmd_{command}", None)
if handler is None:
return "Unknown command. Send /help for the list of commands."
try:
return await handler((match.group("arg") or "").strip(), user_id)
except Exception as exc:
log.exception("Command /%s failed", command)
return f"Command failed: {public_error(exc)}"
# -- /upload conversation --------------------------------------------------------------
def _active_setup(self, user_id: int) -> tuple[UploadTask | None, bool]:
"""This admin's unfinished setup, expiring it if it is older than the timeout."""
setups = self.db.tasks_with_status(*SETUP_TASK_STATUSES, admin_user_id=user_id)
if not setups:
return None, False
setup = setups[-1]
for stale in setups[:-1]: # defensive: only one setup per admin may exist
self.db.update_task(stale.id, status=TaskStatus.CANCELLED, last_error="superseded", completed_at=iso(utcnow()))
age = utcnow() - datetime.fromisoformat(setup.updated_at)
if age > timedelta(minutes=self.config.upload_setup_timeout):
self.db.update_task(
setup.id, status=TaskStatus.CANCELLED, last_error="setup expired", completed_at=iso(utcnow())
)
log.info("Upload setup #%s of admin %s expired", setup.id, user_id)
return None, True
return setup, False
async def on_text(self, user_id: int, text: str) -> str | None:
setup, expired = self._active_setup(user_id)
if setup is None:
if _LINK_HINT.search(text):
note = "Your previous upload setup expired. " if expired else ""
return f"{note}To process a message range, send /upload first."
return None
if setup.status == TaskStatus.AWAITING_START:
return await self._receive_start(setup, text)
if setup.status == TaskStatus.AWAITING_END:
return await self._receive_end(setup, text)
return f"Upload Task #{setup.id} is ready. Send /confirm to begin or /cancel to discard it."
async def _check_link(self, text: str, *, expected_channel_id: int | None = None):
"""Return (link, channel_id, title) or an error message."""
try:
link = parse_message_link(text)
except LinkError as exc:
return str(exc)
source = self.queue.source
if source is None:
return "Telegram access is not available right now."
try:
channel_id, title = await source.resolve_link(link)
except BotError as exc:
return exc.public
if expected_channel_id is not None and channel_id != expected_channel_id:
return "The end link belongs to a different channel than the start link. Both links must come from the same channel."
if (
expected_channel_id is None
and self.config.upload_restrict_to_source_channel
and channel_id != source.channel_id
):
return (
"This link is not from the configured source channel "
f"({source.channel_title or self.config.source_channel}). "
"Set UPLOAD_RESTRICT_TO_SOURCE_CHANNEL=false to allow other channels the bot can access."
)
try:
accessible = await source.message_accessible(channel_id, link.message_id)
except BotError as exc:
return exc.public
if not accessible:
return f"Message {link.message_id} was not found in that channel (deleted, or not visible to the bot's account)."
return link, channel_id, title
async def _receive_start(self, setup: UploadTask, text: str) -> str:
result = await self._check_link(text)
if isinstance(result, str):
return HtmlReply(f"❌ {esc(result)}\n\nPlease send a valid start link, for example:\n{EXAMPLE}")
link, channel_id, title = result
self.db.update_task(
setup.id,
status=TaskStatus.AWAITING_END,
channel_ref=link.channel_ref,
channel_id=channel_id,
channel_title=title,
start_message_id=link.message_id,
)
return HtmlReply(f"Start message: {link.message_id} ✔\n\n{END_PROMPT}")
async def _receive_end(self, setup: UploadTask, text: str) -> str:
result = await self._check_link(text, expected_channel_id=setup.channel_id)
if isinstance(result, str):
return HtmlReply(
f"❌ {esc(result)}\n\nThe start link (message {setup.start_message_id}) is kept. "
"Please send a valid end link."
)
link, _, _ = result
if link.message_id < setup.start_message_id:
return HtmlReply(
f"❌ The end message ID ({link.message_id}) must be greater than or equal to the start message ID "
f"({setup.start_message_id}). Please send a valid end link."
)
size = link.message_id - setup.start_message_id + 1
if size > self.config.upload_max_range:
return HtmlReply(
f"❌ That range covers {size} messages; the maximum is {self.config.upload_max_range} "
"(UPLOAD_MAX_RANGE). Please send a closer end link."
)
task = self.db.update_task(setup.id, status=TaskStatus.AWAITING_CONFIRMATION, end_message_id=link.message_id)
return self._confirmation(task)
def _channel_label(self, task: UploadTask) -> str:
source = self.queue.source
name = task.channel_title or str(task.channel_id)
if source is not None and task.channel_id == source.channel_id:
return f"{name} (configured source channel)"
return name
def _confirmation(self, task: UploadTask) -> HtmlReply:
return HtmlReply(
"✅ <b>Upload Task Ready</b>\n\n"
f"Start Message ID: {task.start_message_id}\n"
f"End Message ID: {task.end_message_id}\n"
f"Channel: {esc(self._channel_label(task))}\n\n"
"The bot will process eligible videos within this message range, organize them into Google Drive "
"topic folders, and clean up local temporary files after successful uploads.\n\n"
"Send /confirm to begin or /cancel to cancel this task."
)
# -- task commands ----------------------------------------------------------------------
async def cmd_upload(self, arg: str, user_id: int) -> str:
setup, expired = self._active_setup(user_id)
if setup is not None:
hint = {
TaskStatus.AWAITING_START: "Send the start link to continue.",
TaskStatus.AWAITING_END: f"Start message {setup.start_message_id} is set; send the end link to continue.",
TaskStatus.AWAITING_CONFIRMATION: "Send /confirm to begin.",
}[setup.status]
return (
f"⚠️ You already have an upload setup in progress (Task #{setup.id}, {_STEP[setup.status]}).\n"
f"{hint}\nOr send /cancel to discard it and start again."
)
self.db.create_task(user_id)
prefix = "Your previous unfinished setup expired.\n\n" if expired else ""
return HtmlReply(prefix + START_PROMPT)
async def cmd_confirm(self, arg: str, user_id: int) -> str:
setup, expired = self._active_setup(user_id)
if setup is None:
note = "Your previous setup expired. " if expired else ""
return f"{note}There is no upload task waiting for confirmation. Start one with /upload."
if setup.status != TaskStatus.AWAITING_CONFIRMATION:
return (
"Both the start and end links must be validated before /confirm. "
f"Task #{setup.id} is {_STEP[setup.status]}."
)
task = self.queue.submit_task(setup.id)
ahead = [t for t in self.db.tasks_with_status(*ACTIVE_TASK_STATUSES) if t.id < task.id]
lines = [f"🚀 <b>Upload Task #{task.id} confirmed</b>", f"Range: {task.range_text}"]
if self.queue.paused_flag:
lines.append("The queue is paused; send /resume to start processing.")
elif ahead:
lines.append("Queued after task(s) " + ", ".join(f"#{t.id}" for t in ahead) + ".")
else:
lines.append("Processing has started. Use /taskstatus to follow it.")
return HtmlReply("\n".join(lines))
async def cmd_cancel(self, arg: str, user_id: int) -> str:
if arg:
if not arg.isdigit():
return "Usage: /cancel or /cancel TASK_ID"
return self._cancel_reply(int(arg), await self.queue.cancel_task(int(arg)))
setup, _ = self._active_setup(user_id)
if setup is not None:
await self.queue.cancel_task(setup.id, reason="setup cancelled")
return f"Upload setup #{setup.id} cancelled. Nothing was downloaded or uploaded."
own = self.db.tasks_with_status(*ACTIVE_TASK_STATUSES, admin_user_id=user_id)
if own:
running = [t for t in own if t.status == TaskStatus.RUNNING]
target = running[0] if running else own[-1]
return self._cancel_reply(target.id, await self.queue.cancel_task(target.id))
others = self.db.tasks_with_status(*ACTIVE_TASK_STATUSES)
if others:
ids = ", ".join(f"#{t.id}" for t in others)
return f"You have nothing to cancel. Active task(s) {ids} belong to another administrator; use /cancel ID to cancel one explicitly."
return "Nothing to cancel."
@staticmethod
def _cancel_reply(task_id: int, outcome: str) -> str:
return {
"unknown": f"Task #{task_id} does not exist.",
"finished": f"Task #{task_id} has already finished.",
"setup_cancelled": f"Upload setup #{task_id} cancelled.",
"cancelled": (
f"🛑 Task #{task_id} cancelled at a safe checkpoint. Completed uploads are kept; "
"the remaining videos were not processed and their temporary files were removed."
),
}[outcome]
async def cmd_taskstatus(self, arg: str, user_id: int) -> str:
if arg:
if not arg.isdigit():
return "Usage: /taskstatus or /taskstatus TASK_ID"
task = self.db.get_task(int(arg))
return self._describe_task(task) if task else f"Task #{arg} does not exist."
parts = []
setup, expired = self._active_setup(user_id)
if setup is not None:
parts.append(
f"Your setup: Task #{setup.id}, {_STEP[setup.status]}"
+ (f" (start {setup.start_message_id})" if setup.start_message_id else "")
)
elif expired:
parts.append("Your previous setup expired.")
active = self.db.tasks_with_status(*ACTIVE_TASK_STATUSES)
parts += [self._describe_task(t) for t in active]
if not active:
recent = self.db.recent_tasks(1)
parts.append("No active upload tasks." + (f"\n\nLast task:\n{self._describe_task(recent[0])}" if recent else ""))
return HtmlReply("\n\n".join(parts))
def _describe_task(self, task: UploadTask) -> HtmlReply:
lines = [
f"<b>Task #{task.id}</b> — {esc(task.status.value)}",
f"Range: {task.range_text}",
f"Channel: {esc(self._channel_label(task)) if task.channel_id else '-'}",
f"Requested by: {task.admin_user_id}",
]
if task.status not in SETUP_TASK_STATUSES:
lines += [
f"Eligible videos: {task.eligible_count} · Uploaded: {task.uploaded_count} · "
f"Duplicates: {task.duplicate_count} · Failed: {task.failed_count}",
f"Non-video skipped: {task.non_video_count} · Other skipped: {task.skipped_count} · "
f"Missing: {task.missing_count}",
f"Temporary files deleted: {task.deleted_files_count}",
]
if task.current_message_id and task.status == TaskStatus.RUNNING:
lines.append(f"Current message: {task.current_message_id}")
board = self.queue.reporter.board if self.queue.reporter else None
if board and board.task and board.task.get("id") == task.id:
lines.append(f"Video: {esc(board.title)} · {esc(board.status)}")
if board.upload:
lines.append(f"Upload: {esc(board.upload.describe())}")
elif board.download:
lines.append(f"Download: {esc(board.download.describe())}")
for label, value in (("Created", task.created_at), ("Started", task.started_at), ("Finished", task.completed_at)):
if value:
lines.append(f"{label}: {value}")
if task.last_error:
lines.append(f"Note: {esc(_short(task.last_error))}")
return HtmlReply("\n".join(lines))
async def cmd_uploadhistory(self, arg: str, user_id: int) -> str:
tasks = self.db.recent_tasks(10)
if not tasks:
return "No upload tasks yet. Start one with /upload."
lines = ["<b>Recent upload tasks</b>"]
for t in tasks:
lines.append(
f"#{t.id} {t.range_text} {esc(t.status.value)} — uploaded {t.uploaded_count}, "
f"duplicates {t.duplicate_count}, failed {t.failed_count}, non-video {t.non_video_count} "
f"(by {t.admin_user_id})"
)
return HtmlReply("\n".join(lines))
# -- general commands ---------------------------------------------------------------------
async def cmd_start(self, arg: str, user_id: int) -> str:
return (
"Telegram -> Google Drive video organizer.\n"
"Videos are uploaded into Drive folders named after the caption Topic, "
"using the caption Title as the file name.\n\n"
"Process a message range: /upload, send the start and end links, then /confirm.\n"
"Process the whole channel: /sync then /run.\n\n" + HELP
)
async def cmd_help(self, arg: str, user_id: int) -> str:
return HELP
async def cmd_status(self, arg: str, user_id: int) -> str:
counts = self.db.counts()
lines = [f"Queue: {self.queue.status}"]
if self.queue.halt_reason:
lines.append(f"Halted because: {self.queue.halt_reason}")
active = self.db.tasks_with_status(*ACTIVE_TASK_STATUSES)
if active:
lines.append("Upload tasks: " + ", ".join(f"#{t.id} {t.status.value} ({t.range_text})" for t in active))
current = self.queue.current
if current and self.queue.reporter:
lines.append("")
lines.append(re.sub(r"</?b>", "", self.queue.reporter.board.render()))
elif current:
lines.append(f"Current: message {current.message_id} ({current.title or 'untitled'})")
lines.append("")
lines.append(
f"Completed {counts['COMPLETED']} · Pending {counts['PENDING'] + counts['RETRY_WAIT']} · "
f"Failed {counts['FAILED']} · Skipped {counts['SKIPPED']}"
)
lines.append(f"Whole-channel processing (/run): {'on' if self.queue.global_enabled else 'off'}")
lines.append(f"Live channel listener: {'on' if self.queue.listening else 'off'}")
return "\n".join(lines)
async def cmd_progress(self, arg: str, user_id: int) -> str:
c = self.db.counts()
total = sum(c.values())
processed = c["COMPLETED"]
remaining = self.queue.remaining()
considered = total - c["SKIPPED"]
pct = f"{100 * processed / considered:.1f}%" if considered else "n/a"
return (
f"Total discovered: {total}\n"
f"Processed (uploaded & verified): {processed} ({pct})\n"
f"Pending: {c['PENDING']}\n"
f"Waiting to retry: {c['RETRY_WAIT']}\n"
f"In progress: {c['DOWNLOADING'] + c['DOWNLOADED'] + c['UPLOADING']}\n"
f"Failed: {c['FAILED']}\n"
f"Skipped: {c['SKIPPED']}\n"
f"Remaining: {remaining}"
)
async def cmd_run(self, arg: str, user_id: int) -> str:
if self.queue.start():
pending = self.queue.remaining()
hint = "" if pending else "\nThe queue is empty - use /sync to discover channel videos."
return f"Processing started ({pending} item(s) waiting).{hint}"
return f"Processing is already {self.queue.status}."
async def cmd_pause(self, arg: str, user_id: int) -> str:
if self.queue.pause():
return "Pausing: the current transfer stops at the next resumable checkpoint. Use /resume to continue."
return f"Nothing to pause (queue is {self.queue.status})."
async def cmd_resume(self, arg: str, user_id: int) -> str:
if self.queue.resume():
return "Processing resumed."
return f"Nothing to resume (queue is {self.queue.status})."
async def cmd_stop(self, arg: str, user_id: int) -> str:
await self.queue.stop(timeout=60)
return (
"Stopped. Interrupted work was checkpointed; a running upload task was paused. "
"Use /resume to continue upload tasks or /run for whole-channel processing."
)
async def cmd_retry(self, arg: str, user_id: int) -> str:
count = self.queue.retry_failed()
if not count:
return "No failed items to retry."
started = self.queue.start()
return f"Re-queued {count} failed item(s)." + (" Processing started." if started else "")
async def cmd_failed(self, arg: str, user_id: int) -> str:
items = self.db.list_by_state(State.FAILED, limit=20)
if not items:
return "No failed items."
lines = [f"Failed items (showing {len(items)}):"]
for rec in items:
lines.append(f"• #{rec.message_id} {_short(rec.title, 60)} — {_short(rec.last_error)}")
lines.append("Use /retry to re-queue them, /check ID for details.")
return "\n".join(lines)
async def cmd_folders(self, arg: str, user_id: int) -> str:
topics = self.db.topic_counts()
counts = self.db.counts()
lines = [f"Drive folders created/reused: {self.db.folder_count()}", f"Files uploaded: {counts['COMPLETED']}"]
if topics:
lines.append("")
lines += [f"• {topic or self.config.fallback_topic}: {n}" for topic, n in topics]
return "\n".join(lines)
async def cmd_check(self, arg: str, user_id: int) -> str:
if not arg.lstrip("-").isdigit():
return "Usage: /check MESSAGE_ID"
records = self.db.find_by_message_id(int(arg))
if not records:
return f"Message {arg} is not in the database. Use /sync or /upload to discover messages."
return "\n\n".join(self._describe(r) for r in records)
async def cmd_sync(self, arg: str, user_id: int) -> str:
result = await self.queue.sync(rescan=arg.lower() == "all")
if self.queue.running:
self.queue.wake()
return "Sync complete.\n" + result.summary()
def _describe(self, rec: VideoRecord) -> str:
lines = [
f"Message {rec.message_id} — {rec.state}",
f"Title: {rec.title or '-'}",
f"Topic: {rec.topic or '- (uses ' + self.config.fallback_topic + ')'}",
f"Batch: {rec.batch or '-'}",
f"Index: {', '.join(rec.indexes) or '-'}",
f"Extracted by: {rec.extracted_by or '-'}",
f"Download: {rec.download_status or '-'} · Upload: {rec.upload_status or '-'}",
f"Size: {rec.expected_size or rec.local_size or '?'} bytes",
f"Attempts: {rec.attempts}",
]
if rec.task_id:
lines.append(f"Upload task: #{rec.task_id}")
if rec.final_name:
lines.append(f"Drive name: {rec.final_name}")
if rec.drive_file_id:
lines.append(f"Drive file ID: {rec.drive_file_id}")
if rec.next_retry_at and rec.state == State.RETRY_WAIT:
lines.append(f"Next retry: {rec.next_retry_at}")
if rec.skip_reason:
lines.append(f"Skip reason: {_short(rec.skip_reason)}")
if rec.last_error:
lines.append(f"Last error: {_short(rec.last_error)}")
if rec.warnings:
lines.append(f"Warnings: {_short('; '.join(rec.warnings), 300)}")
return "\n".join(lines)
def _parse_mode(text: str) -> str | None:
return "html" if isinstance(text, HtmlReply) else None
class TelethonNotifier:
def __init__(self, client):
self.client = client
async def send(self, chat_id: int, text: str) -> int | None:
message = await self.client.send_message(chat_id, str(text), parse_mode=_parse_mode(text), link_preview=False)
return message.id
async def edit(self, chat_id: int, message_id: int, text: str) -> None:
from telethon import errors
try:
await self.client.edit_message(chat_id, message_id, str(text), parse_mode=_parse_mode(text), link_preview=False)
except errors.MessageNotModifiedError:
pass
def register_command_handlers(client, router: CommandRouter) -> None:
from telethon import events
@client.on(events.NewMessage(incoming=True, func=lambda e: e.is_private))
async def _on_private_message(event) -> None: # pragma: no cover - needs a live Telegram connection
reply = await router.handle(event.sender_id, event.raw_text)
if not reply:
return
if isinstance(reply, HtmlReply):
await event.respond(str(reply), parse_mode="html", link_preview=False)
return
for start in range(0, len(reply), 4000):
await event.respond(reply[start : start + 4000], parse_mode=None, link_preview=False)
Reference in New Issue View Git Blame Copy Permalink