"""Message handling. Independent of the WhatsApp connection: `io` is supplied by the caller. io must provide: await io.reply(text) -> key (a new message) await io.edit(key, text) -> None (edit an earlier message) await io.send_file(item) -> key (the file message) await io.schedule(keys, when) -> None (delete these keys for everyone at `when`) A key is a tuple (jid, message_id). """ from __future__ import annotations import asyncio import logging import re import secrets import time from datetime import datetime, timedelta, timezone from .messages import MSG SINHALA = re.compile(r"[඀-෿]") CODE_RE = re.compile(r"^(?:[.!/]?(?:get|dl|code)\s+)?([A-Za-z0-9_-]{3,20})$", re.IGNORECASE) EXPLICIT_RE = re.compile(r"^[.!/]?(get|dl|code)\s", re.IGNORECASE) COMMAND_RE = re.compile(r"^/(\w+)\s*(.*)$", re.DOTALL) ROLES = ("free", "premium", "admin") log = logging.getLogger("bot") def t(key: str, lang: str, **vars) -> str: entry = MSG[key] text = entry.get(lang) or entry["en"] return re.sub(r"\{(\w+)\}", lambda m: str(vars.get(m.group(1), "")), text) def limit_for(user: dict, limits: dict) -> float: """Downloads allowed per 24 hours. The user's own dailyLimit wins. Infinity means unlimited.""" if user.get("dailyLimit") is not None: return user["dailyLimit"] if user["role"] not in limits: return 0 value = limits[user["role"]] return float("inf") if value is None else value class Bot: def __init__(self, db, cooldown_seconds: float = 5.0, max_concurrent: int = 3, catalog_url: str | None = None): self.db = db self.cooldown_seconds = cooldown_seconds self.catalog_url = catalog_url self._cooldown: dict[str, float] = {} self._slots = asyncio.Semaphore(max_concurrent) async def _safe_edit(self, io, key, text): try: await io.edit(key, text) return key except Exception: return await io.reply(text) async def _admin_command(self, name: str, args: str, io) -> None: parts = args.strip().split() phone = re.sub(r"\D", "", parts[0]) if parts else "" if name == "stats": s = await self.db.stats() await io.reply( f"Users: {s['users']} (premium {s['premium']}, admins {s['admins']}, blocked {s['blocked']})\n" f"Downloads today: {s['downloads_today']} · total: {s['downloads_total']}\n" f"Files: {s['files']}" ) return if name == "grant": role = parts[1].lower() if len(parts) > 1 else "" if not phone or role not in ROLES: await io.reply("Usage: /grant free|premium|admin") return user = await self.db.update_user(phone, {"role": role}) await io.reply(f"✅ {phone} is now {role}." if user else f"User {phone} not found. They must message the bot first.") return if name in ("block", "unblock"): if not phone: await io.reply(f"Usage: /{name} ") return user = await self.db.update_user(phone, {"blocked": name == "block"}) if user: await io.reply(f"✅ {phone} {'blocked' if name == 'block' else 'unblocked'}.") else: await io.reply(f"User {phone} not found.") return await io.reply("Admin commands: /stats · /grant free|premium|admin · /block · /unblock ") async def handle(self, user_id: str, text: str, io) -> None: clean = (text or "").strip() if not clean: return user, is_new = await self.db.touch_user(user_id, "si" if SINHALA.search(clean) else None) lang = user["lang"] if is_new: catalog = t("catalogLine", lang, url=self.catalog_url) if self.catalog_url else "" await io.reply(t("welcome", lang, catalog=catalog)) if user["blocked"]: await io.reply(t("blocked", lang)) return cmd = COMMAND_RE.match(clean) if cmd: name = cmd.group(1).lower() args = cmd.group(2) if name == "lang": value = args.strip().lower() if value not in ("si", "en"): await io.reply(t("help", lang)) return await self.db.update_user(user_id, {"lang": value}) await io.reply(t("langSet", value)) return if name == "help": await io.reply(t("help", lang)) return if user["role"] == "admin": await self._admin_command(name, args, io) return match = CODE_RE.match(clean) if not match: return explicit = bool(EXPLICIT_RE.match(clean)) item = await self.db.find_file(match.group(1)) if not item: if explicit: await io.reply(t("notFound", lang)) return now = time.monotonic() if now - self._cooldown.get(user_id, -1e9) < self.cooldown_seconds: return if len(self._cooldown) > 1000: self._cooldown.clear() self._cooldown[user_id] = now settings = await self.db.get_settings() limit = limit_for(user, settings["limits"]) used = 0 if limit != float("inf"): used = await self.db.downloads_today(user_id) if used >= limit: await io.reply(t("limit", lang, limit=int(limit))) return progress = await io.reply(t("preparing", lang, title=item["title"])) async with self._slots: try: file_key = await io.send_file(item) await self.db.record_download(user_id, item["code"]) await self.db.record_send(item["code"]) left = "" if limit == float("inf") else t("leftToday", lang, n=int(limit) - used - 1) hours = settings["autoDeleteHours"] note = t("deleteNote", lang, hours=hours) if hours > 0 else "" done_key = await self._safe_edit(io, progress, t("done", lang, title=item["title"], left=left, note=note)) if hours > 0: when = datetime.now(timezone.utc) + timedelta(hours=hours) try: await io.schedule([file_key, done_key], when) except Exception as e: log.error("Could not schedule auto-delete: %s", e) except Exception as e: ref = f"E-{secrets.randbelow(900) + 100}" log.error("Send failed for %s (%s): %s", item["code"], ref, e) await self._safe_edit(io, progress, t("error", lang, ref=ref))