| """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 <number> 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} <number>") |
| 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 <number> free|premium|admin · /block <number> · /unblock <number>") |
|
|
| 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)) |
|
|
