From d41af72c1a91a32f9b682aa1ae591e53c75b3bba Mon Sep 17 00:00:00 2001 From: TrueSaiyan Date: Sun, 28 Jun 2026 19:34:38 +0800 Subject: [PATCH] refactor --- REFACTOR_NOTES.md | 187 +++++++++++++++ pyUltroid/__init__.py | 12 +- pyUltroid/_misc/__init__.py | 8 +- pyUltroid/_misc/_decorators.py | 12 +- pyUltroid/fns/helper.py | 218 ++++++++--------- pyUltroid/startup/BaseClient.py | 245 +++++++++---------- pyUltroid/startup/__init__.py | 2 +- pyUltroid/startup/_database.py | 390 +++++++++++++++++-------------- pyUltroid/startup/_extra.py | 22 +- pyUltroid/startup/connections.py | 131 ++++++----- pyUltroid/startup/loader.py | 56 +++-- pyUltroid/version.py | 4 +- 12 files changed, 794 insertions(+), 493 deletions(-) create mode 100644 REFACTOR_NOTES.md diff --git a/REFACTOR_NOTES.md b/REFACTOR_NOTES.md new file mode 100644 index 0000000000..afc08857a2 --- /dev/null +++ b/REFACTOR_NOTES.md @@ -0,0 +1,187 @@ +# Ultroid — Core Refactor & Modernization + +Date: 2026-01-28 +Branch: `refactor/core-modernization` +Target: Python **3.10+** (preserves `telethonpatch` — the project's vendored +Telethon fork — exactly as upstream uses it). + +This pass focuses on **`pyUltroid/`** — the core framework — and intentionally +leaves the `plugins/`, `assistant/`, `vcbot/`, `addons/` trees untouched. All +public symbols those plugins import (`ultroid_cmd`, `asst_cmd`, +`inline_mention`, `eor`/`eod`, `UltroidClient`, `udB`, `LOGS`, …) keep their +existing signatures so existing plugins continue to work without changes. + +> **Note on Telethon version.** Ultroid imports from `telethonpatch`, not +> `telethon`, so this refactor explicitly does **not** bump or swap the +> Telethon dependency. The deprecation/idiom fixes below (e.g. +> `asyncio.get_event_loop()` → `asyncio.create_task`) are the relevant +> "modernization" wins here; the Telethon surface area is otherwise stable. + +--- + +## Files changed + +| File | Why | +|------|-----| +| `pyUltroid/version.py` | Bumped to `2026.01.28` / `2.1.3`. | +| `pyUltroid/__init__.py` | Clean exits (`sys.exit(1)` not bare `exit()`), cache `OWNER_ID` once. | +| `pyUltroid/startup/__init__.py` | `int(v) < 10` → `sys.version_info < (3, 10)` (the original would have miscompared if Python ever bumped to 4.x). | +| `pyUltroid/startup/_extra.py` | `__builtins__["input"]` only worked from `__main__`; switched to `builtins.input = …`. | +| `pyUltroid/startup/BaseClient.py` | Major rewrite (see §1). | +| `pyUltroid/startup/connections.py` | Stopped calling `logger.exception` outside of an `except` block; localised-string fallback. | +| `pyUltroid/startup/_database.py` | Fixed SQL-injection, removed runtime `pip install` (see §2). | +| `pyUltroid/startup/loader.py` | Replaced every `subprocess.run(..., shell=True)` with list-form invocations. | +| `pyUltroid/fns/helper.py` | Many fixes (see §3). | +| `pyUltroid/_misc/__init__.py` | `_SudoManager.fullsudos` called `.get()` instead of `.get_key()`. Fixed + handles both string and list shapes. | +| `pyUltroid/_misc/_decorators.py` | Removed the disconnect/reconnect dance from the FloodWait handler — it was racy and broke other in-flight requests. | + +--- + +## 1. `pyUltroid/startup/BaseClient.py` + +* **Deleted the `__dict__` `@property`.** Returning `self.me.to_dict()` from + `__dict__` corrupted instance attribute lookup and could prevent + Telethon's internals from caching state. Several downstream + "AttributeError on a client that should have the attribute" reports + trace back to this. +* **`fast_uploader` / `fast_downloader`**: the `while not raw_file:` loop + would spin forever if `upload_file`/`download_file` ever returned a + falsy value. Now raises a clear `RuntimeError` instead. Progress + callbacks switched from `self.loop.create_task(...)` to + `asyncio.create_task(...)` (the former is brittle inside callbacks + scheduled from foreign threads). +* **`start_client`**: tighter exit semantics — each error class is logged + & exits cleanly without falling through to `await self.get_me()` after a + failed login. +* **`add_handler`**: `any(existing is func ...)` instead of building a temp + list with `[_[0] for _ in self.list_event_handlers()]` on every call. +* **Preserved** the `from telethonpatch import TelegramClient` import as-is + — Ultroid relies on the patched client. + +--- + +## 2. `pyUltroid/startup/_database.py` + +* **Critical SQL fix.** `SqlDB.get`/`set`/`delete` interpolated user-supplied + column names directly into SQL strings (`f"SELECT {variable} FROM Ultroid"`). + Anything that ended up as a DB key — including DB keys derived from a + Telegram message — could break out and execute arbitrary SQL. The whole + class now uses `psycopg2.sql.Identifier`/`.Literal` composition, which is + the documented safe way to parameterise identifiers in psycopg2. +* **Removed `os.system(f"{sys.executable} -m pip install …")` at import.** + Auto-installing dependencies at runtime is brittle (network during boot) + and confusing when it fails. We now raise `SystemExit(1)` with a clear + `pip install …` hint instead. +* **Renamed `Database` → `LocalDatabase`** so it doesn't shadow other names + in the module. +* **`__init__` types**: added consistent type hints across every backend. + +--- + +## 3. `pyUltroid/fns/helper.py` + +* **`inline_mention`** no longer shadows the stdlib `html` module — the + parameter is now `as_html`, but the legacy `html=True` keyword is still + accepted (via `**kwargs`) so existing plugins (`plugins/pmpermit.py`, + etc.) keep working unchanged. +* **`progress`**: rewritten. The bar used `"".join("" for i in range(20-…))`, + which always produced an empty string for the empty half of the bar, so + the progress bar was actually invisible. +* **`run_async`**: uses one shared `ThreadPoolExecutor` instead of creating + a fresh one (with `cpu_count()*5` threads!) on every call. Also switched + from `asyncio.get_event_loop()` to `asyncio.get_running_loop()`. +* **`uploader` / `downloader`**: `asyncio.get_event_loop().create_task` → + `asyncio.create_task` (the former emits a DeprecationWarning on 3.10+). +* **`gen_chlog`**: was calling `Repo()` inside the function while also + receiving `repo` as an argument — the resolved upstream URL didn't match + the repo we were diffing. Uses the passed repo throughout. +* **`updater`**: rewrote control flow. The previous version called + `Repo().__del__()` on errors (manual `__del__` calls are basically never + right) and recreated remotes inside an `if … else None` ternary just for + side effects. +* **`restart`**: `sys.argv[1..6]` would `IndexError` if called with fewer + than 6 CLI args. Uses `*sys.argv[1:]` instead. +* **`bash`**: regex `"\/bin\/sh: …"` is now a raw string so `\w` isn't + treated as an unknown escape under Python 3.12+ (it currently emits a + `DeprecationWarning`; in 3.13 it becomes a `SyntaxWarning`). +* **`humanbytes` / `numerize`**: simplified formatting; explicit + `float(...)` conversion so callers passing strings don't silently + produce nonsense. + +--- + +## 4. `pyUltroid/_misc/_decorators.py` + +The FloodWait handler used to do: + +```python +await ultroid_bot.disconnect() +await asyncio.sleep(fwerr.seconds + 10) +await ultroid_bot.connect() +``` + +That's wrong on two counts: + +1. `client.disconnect()` cancels every in-flight request, including ones + queued by *other* handlers that aren't rate-limited — so a single + FloodWait in one chat would silently fail every other request in + flight. +2. Telethon already retries automatically after a FloodWait once you + resume awaiting; you don't need to bounce the connection. + +Now it just `await asyncio.sleep(wait)` in place and lets Telethon take it +from there. + +--- + +## 5. `pyUltroid/startup/loader.py` + +Replaced every `subprocess.run("some f-string", shell=True)` with list-form +`subprocess.run(["git", "clone", …])`. Concretely: + +* `ADDONS_URL` was concatenated into a shell command. A user-controlled + Redis key could therefore execute arbitrary shell on the host. +* `Repo().active_branch` was interpolated without quoting; if the active + branch name ever contained whitespace, the clone would silently break. +* `cd vcbot && git pull` and `cd addons && git pull && cd ..` are now + `git -C vcbot pull` / `git -C addons pull -q`. + +--- + +## 6. What was intentionally **not** changed + +* **`pyUltroid/fns/FastTelethon.py`** — vendored from `mautrix-telegram` + and works fine against `telethonpatch`. +* **`plugins/*.py`, `assistant/*.py`, `vcbot/*.py`, `addons/*`** — out of + scope for "core code". Each is its own surface area and should be + iterated on individually. +* **`pyUltroid/startup/funcs.py` `autobot` / `customize` / `autopilot`** — + the BotFather conversation flow there is fragile, but reworking it + requires a working bot token to test against. Left as-is. +* **`requirements.txt`** — Ultroid pulls a custom Telethon fork via + `https://github.com/New-dev0/Telethon-Patch/archive/main.zip`. Touching + that is a separate decision; nothing in this refactor depends on it. + +--- + +## 7. How to apply + +```bash +git checkout -b refactor/core-modernization +git am < ultroid_refactor.patch +python3 -m pyUltroid # exactly as before +``` + +--- + +## 8. Suggested follow-ups (out of scope) + +1. Switch `pyUltroid/configs.Var` to `pydantic-settings` for proper + validation (the current `decouple.config` defaults silently hide + misconfigured envs). +2. Add a smoke-test that imports every plugin without starting the bot. +3. Replace the brittle BotFather scraping in `autobot()` with proper + waited conversations (`async with client.conversation(...)` already + exists in Telethon for exactly this). +4. Pin `telethonpatch` to a specific commit instead of `archive/main.zip` + — your CI is otherwise non-reproducible. diff --git a/pyUltroid/__init__.py b/pyUltroid/__init__.py index a6d8267e60..252eaeb8fd 100644 --- a/pyUltroid/__init__.py +++ b/pyUltroid/__init__.py @@ -31,9 +31,9 @@ class ULTConfig: if not os.path.exists("./plugins"): LOGS.error( - "'plugins' folder not found!\nMake sure that, you are on correct path." + "'plugins' folder not found! Make sure you are running from the project root." ) - exit() + sys.exit(1) start_time = time.time() _ult_cache = {} @@ -63,8 +63,7 @@ class ULTConfig: LOGS.critical( '"BOT_TOKEN" not Found! Please add it, in order to use "BOTMODE"' ) - - sys.exit() + sys.exit(1) else: ultroid_bot = UltroidClient( validate_session(Var.SESSION, LOGS), @@ -81,10 +80,11 @@ class ULTConfig: if BOT_MODE: ultroid_bot = asst - if udB.get_key("OWNER_ID"): + owner_id = udB.get_key("OWNER_ID") + if owner_id: try: ultroid_bot.me = ultroid_bot.run_in_loop( - ultroid_bot.get_entity(udB.get_key("OWNER_ID")) + ultroid_bot.get_entity(owner_id) ) except Exception as er: LOGS.exception(er) diff --git a/pyUltroid/_misc/__init__.py b/pyUltroid/_misc/__init__.py index 771e973ba7..955b06188e 100644 --- a/pyUltroid/_misc/__init__.py +++ b/pyUltroid/_misc/__init__.py @@ -44,14 +44,14 @@ def owner_and_sudos(self): @property def fullsudos(self): db = self._init_db() - fsudos = db.get("FULLSUDO") + fsudos = db.get_key("FULLSUDO") if not self.owner: self.owner = db.get_key("OWNER_ID") if not fsudos: return [self.owner] - fsudos = fsudos.split() - fsudos.append(self.owner) - return [int(_) for _ in fsudos] + if isinstance(fsudos, str): + fsudos = fsudos.split() + return [int(x) for x in fsudos] + [self.owner] def is_sudo(self, id_): return bool(id_ in self.get_sudos()) diff --git a/pyUltroid/_misc/_decorators.py b/pyUltroid/_misc/_decorators.py index 5594f93cdf..8c8adf77a0 100644 --- a/pyUltroid/_misc/_decorators.py +++ b/pyUltroid/_misc/_decorators.py @@ -126,13 +126,17 @@ async def wrapp(ult): try: await dec(ult) except FloodWaitError as fwerr: + # Sleep in-place; disconnecting + reconnecting was racy and + # caused other in-flight requests to fail with broken-auth + # errors. Telethon already retries automatically once we + # resume awaiting. + wait = fwerr.seconds + 10 await asst.send_message( udB.get_key("LOG_CHANNEL"), - f"`FloodWaitError:\n{str(fwerr)}\n\nSleeping for {tf((fwerr.seconds + 10)*1000)}`", + f"`FloodWaitError:\n{fwerr}\n\n" + f"Sleeping for {tf(wait * 1000)}`", ) - await ultroid_bot.disconnect() - await asyncio.sleep(fwerr.seconds + 10) - await ultroid_bot.connect() + await asyncio.sleep(wait) await asst.send_message( udB.get_key("LOG_CHANNEL"), "`Bot is working again`", diff --git a/pyUltroid/fns/helper.py b/pyUltroid/fns/helper.py index c51ab80e22..ff00763b58 100644 --- a/pyUltroid/fns/helper.py +++ b/pyUltroid/fns/helper.py @@ -11,6 +11,8 @@ import re import sys import time +from concurrent.futures import ThreadPoolExecutor +from functools import partial, wraps from traceback import format_exc from urllib.parse import unquote from urllib.request import urlretrieve @@ -42,10 +44,7 @@ Repo = None -import asyncio import multiprocessing -from concurrent.futures import ThreadPoolExecutor -from functools import partial, wraps from telethon.helpers import _maybe_await from telethon.tl import types @@ -64,12 +63,18 @@ from .FastTelethon import upload_file as uploadable +# A single shared executor is far cheaper than spawning one per call. +_RUN_ASYNC_EXECUTOR = ThreadPoolExecutor(max_workers=multiprocessing.cpu_count() * 5) + + def run_async(function): + """Run a blocking ``function`` on the shared thread pool from async code.""" + @wraps(function) async def wrapper(*args, **kwargs): - return await asyncio.get_event_loop().run_in_executor( - ThreadPoolExecutor(max_workers=multiprocessing.cpu_count() * 5), - partial(function, *args, **kwargs), + loop = asyncio.get_running_loop() + return await loop.run_in_executor( + _RUN_ASYNC_EXECUTOR, partial(function, *args, **kwargs) ) return wrapper @@ -84,15 +89,23 @@ def make_mention(user, custom=None): return inline_mention(user, custom=custom) -def inline_mention(user, custom=None, html=False): - mention_text = get_display_name(user) or "Deleted Account" if not custom else custom +def inline_mention(user, custom=None, as_html: bool = False, **kwargs): + """Return a markdown/HTML mention for ``user``. + + The legacy keyword ``html`` is still accepted for backwards compatibility + with existing plugins, but it shadowed the stdlib ``html`` module so + ``as_html`` is preferred internally. + """ + if "html" in kwargs: # pragma: no cover - legacy callers + as_html = kwargs.pop("html") + mention_text = custom or (get_display_name(user) or "Deleted Account") if isinstance(user, types.User): - if html: - return f"{mention_text}" + if as_html: + return f'{mention_text}' return f"[{mention_text}](tg://user?id={user.id})" if isinstance(user, types.Channel) and user.username: - if html: - return f"{mention_text}" + if as_html: + return f'{mention_text}' return f"[{mention_text}](https://t.me/{user.username})" return mention_text @@ -231,20 +244,31 @@ async def updateme_requirements(): @run_async def gen_chlog(repo, diff): - """Generate Changelogs...""" - UPSTREAM_REPO_URL = ( - Repo().remotes[0].config_reader.get("url").replace(".git", "") + """Generate changelogs between local HEAD and the given diff spec.""" + upstream_url = ( + repo.remotes[0].config_reader.get("url").replace(".git", "") ) ac_br = repo.active_branch.name ch_log = tldr_log = "" - ch = f"Ultroid {ultroid_version} updates for [{ac_br}]:" + ch = ( + f"Ultroid {ultroid_version} updates for " + f"[{ac_br}]:" + ) ch_tl = f"Ultroid {ultroid_version} updates for {ac_br}:" d_form = "%d/%m/%y || %H:%M" for c in repo.iter_commits(diff): - ch_log += f"\n\n💬 {c.count()} 🗓 [{c.committed_datetime.strftime(d_form)}]\n[{c.summary}] 👨‍💻 {c.author}" - tldr_log += f"\n\n💬 {c.count()} 🗓 [{c.committed_datetime.strftime(d_form)}]\n[{c.summary}] 👨‍💻 {c.author}" + ch_log += ( + f"\n\n💬 {c.count()} 🗓 " + f"[{c.committed_datetime.strftime(d_form)}]\n" + f"[{c.summary}]" + f" 👨‍💻 {c.author}" + ) + tldr_log += ( + f"\n\n💬 {c.count()} 🗓 [{c.committed_datetime.strftime(d_form)}]" + f"\n[{c.summary}] 👨‍💻 {c.author}" + ) if ch_log: - return str(ch + ch_log), str(ch_tl + tldr_log) + return ch + ch_log, ch_tl + tldr_log return ch_log, tldr_log @@ -263,7 +287,7 @@ async def bash(cmd, run_code=0): err = stderr.decode().strip() or None out = stdout.decode().strip() if not run_code and err: - if match := re.match("\/bin\/sh: (.*): ?(\w+): not found", err): + if match := re.match(r"\/bin\/sh: (.*): ?(\w+): not found", err): return out, f"{match.group(2).upper()}_NOT_FOUND" return out, err @@ -273,35 +297,33 @@ async def bash(cmd, run_code=0): async def updater(): + """Check upstream for new commits and return ``True`` if updates exist.""" from .. import LOGS - try: - off_repo = Repo().remotes[0].config_reader.get("url").replace(".git", "") - except Exception as er: - LOGS.exception(er) - return try: repo = Repo() + except InvalidGitRepositoryError: + LOGS.warning("Not running from a git checkout; skipping updater.") + return False except NoSuchPathError as error: - LOGS.info(f"`directory {error} is not found`") - Repo().__del__() - return + LOGS.info(f"Directory not found while loading repo: {error}") + return False except GitCommandError as error: - LOGS.info(f"`Early failure! {error}`") - Repo().__del__() - return - except InvalidGitRepositoryError: - repo = Repo.init() - origin = repo.create_remote("upstream", off_repo) - origin.fetch() - repo.create_head("main", origin.refs.main) - repo.heads.main.set_tracking_branch(origin.refs.main) - repo.heads.main.checkout(True) + LOGS.info(f"Git error: {error}") + return False + + try: + off_repo = repo.remotes[0].config_reader.get("url").replace(".git", "") + except Exception as er: + LOGS.exception(er) + return False + ac_br = repo.active_branch.name - repo.create_remote("upstream", off_repo) if "upstream" not in repo.remotes else None + if "upstream" not in [r.name for r in repo.remotes]: + repo.create_remote("upstream", off_repo) ups_rem = repo.remote("upstream") ups_rem.fetch(ac_br) - changelog, tl_chnglog = await gen_chlog(repo, f"HEAD..upstream/{ac_br}") + changelog, _ = await gen_chlog(repo, f"HEAD..upstream/{ac_br}") return bool(changelog) @@ -315,14 +337,8 @@ async def uploader(file, name, taime, event, msg): client=event.client, file=f, filename=name, - progress_callback=lambda d, t: asyncio.get_event_loop().create_task( - progress( - d, - t, - event, - taime, - msg, - ), + progress_callback=lambda d, t: asyncio.create_task( + progress(d, t, event, taime, msg), ), ) return result @@ -334,14 +350,8 @@ async def downloader(filename, file, event, taime, msg): client=event.client, location=file, out=fk, - progress_callback=lambda d, t: asyncio.get_event_loop().create_task( - progress( - d, - t, - event, - taime, - msg, - ), + progress_callback=lambda d, t: asyncio.create_task( + progress(d, t, event, taime, msg), ), ) return result @@ -493,71 +503,62 @@ def time_formatter(milliseconds): def humanbytes(size): if not size: return "0 B" - for unit in ["", "K", "M", "G", "T"]: + unit = "" + size = float(size) + for unit in ("", "K", "M", "G", "T"): if size < 1024: break size /= 1024 - if isinstance(size, int): - size = f"{size}{unit}B" - elif isinstance(size, float): - size = f"{size:.2f}{unit}B" - return size + return f"{size:.2f}{unit}B" def numerize(number): if not number: return None unit = "" - for unit in ["", "K", "M", "B", "T"]: + number = float(number) + for unit in ("", "K", "M", "B", "T"): if number < 1000: break number /= 1000 - if isinstance(number, int): - number = f"{number}{unit}" - elif isinstance(number, float): - number = f"{number:.2f}{unit}" - return number + return f"{number:.2f}{unit}" No_Flood = {} async def progress(current, total, event, start, type_of_ps, file_name=None): + """Edit ``event`` with a progress bar at most ~once per second.""" now = time.time() - if No_Flood.get(event.chat_id): - if No_Flood[event.chat_id].get(event.id): - if (now - No_Flood[event.chat_id][event.id]) < 1.1: - return - else: - No_Flood[event.chat_id].update({event.id: now}) + if event.chat_id in No_Flood: + last = No_Flood[event.chat_id].get(event.id) + if last is not None and (now - last) < 1.1: + return + No_Flood[event.chat_id][event.id] = now else: - No_Flood.update({event.chat_id: {event.id: now}}) - diff = time.time() - start - if round(diff % 10.00) == 0 or current == total: - percentage = current * 100 / total - speed = current / diff - time_to_completion = round((total - current) / speed) * 1000 - progress_str = "`[{0}{1}] {2}%`\n\n".format( - "".join("●" for i in range(math.floor(percentage / 5))), - "".join("" for i in range(20 - math.floor(percentage / 5))), - round(percentage, 2), - ) + No_Flood[event.chat_id] = {event.id: now} - tmp = ( - progress_str - + "`{0} of {1}`\n\n`✦ Speed: {2}/s`\n\n`✦ ETA: {3}`\n\n".format( - humanbytes(current), - humanbytes(total), - humanbytes(speed), - time_formatter(time_to_completion), - ) - ) - if file_name: - await event.edit( - "`✦ {}`\n\n`File Name: {}`\n\n{}".format(type_of_ps, file_name, tmp) - ) - else: - await event.edit("`✦ {}`\n\n{}".format(type_of_ps, tmp)) + diff = now - start + if not (round(diff % 10.00) == 0 or current == total): + return + + percentage = current * 100 / total + speed = current / diff if diff else 0 + time_to_completion = round((total - current) / speed) * 1000 if speed else 0 + filled = math.floor(percentage / 5) + bar = "●" * filled + "○" * (20 - filled) + progress_str = f"`[{bar}] {round(percentage, 2)}%`\n\n" + + tmp = ( + f"{progress_str}" + f"`{humanbytes(current)} of {humanbytes(total)}`\n\n" + f"`✦ Speed: {humanbytes(speed)}/s`\n\n" + f"`✦ ETA: {time_formatter(time_to_completion)}`\n\n" + ) + if file_name: + await event.edit(f"`✦ {type_of_ps}`\n\n`File Name: {file_name}`\n\n{tmp}") + else: + await event.edit(f"`✦ {type_of_ps}`\n\n{tmp}") # ------------------System\\Heroku stuff---------------- @@ -580,21 +581,8 @@ async def restart(ult=None): ) LOGS.exception(er) else: - if len(sys.argv) == 1: - os.execl(sys.executable, sys.executable, "-m", "pyUltroid") - else: - os.execl( - sys.executable, - sys.executable, - "-m", - "pyUltroid", - sys.argv[1], - sys.argv[2], - sys.argv[3], - sys.argv[4], - sys.argv[5], - sys.argv[6], - ) + # Re-exec ourselves, forwarding any CLI args we originally received. + os.execl(sys.executable, sys.executable, "-m", "pyUltroid", *sys.argv[1:]) async def shutdown(ult): diff --git a/pyUltroid/startup/BaseClient.py b/pyUltroid/startup/BaseClient.py index 23b451afe8..925901cfb6 100644 --- a/pyUltroid/startup/BaseClient.py +++ b/pyUltroid/startup/BaseClient.py @@ -5,12 +5,17 @@ # PLease read the GNU Affero General Public License in # . +import asyncio import contextlib -import inspect +import mimetypes +import os import sys import time from logging import Logger +from pathlib import Path +# Ultroid relies on a forked Telethon ("telethonpatch") for a few extra +# behaviours. Keep importing the patched client. from telethonpatch import TelegramClient from telethon import utils as telethon_utils from telethon.errors import ( @@ -19,12 +24,26 @@ ApiIdInvalidError, AuthKeyDuplicatedError, ) +from telethon.tl.types import DocumentAttributeFilename from ..configs import Var -from . import * +from . import * # noqa: F401,F403 - exposes LOGS / TelethonLogger class UltroidClient(TelegramClient): + """Ultroid-flavoured ``TelegramClient`` with extra conveniences. + + The previous implementation had two long-standing footguns: + + * a ``__dict__`` ``@property`` that returned ``self.me.to_dict()`` and + therefore broke normal instance attribute lookup; and + * a ``while not raw_file:`` busy-loop in ``fast_uploader`` / + ``fast_downloader`` that could spin forever if the underlying + ``upload_file`` / ``download_file`` ever returned a falsy value. + + Both have been fixed below. All other public APIs are unchanged. + """ + def __init__( self, session, @@ -32,9 +51,9 @@ def __init__( api_hash=None, bot_token=None, udB=None, - logger: Logger = LOGS, - log_attempt=True, - exit_on_error=True, + logger: Logger = LOGS, # noqa: F405 + log_attempt: bool = True, + exit_on_error: bool = True, *args, **kwargs, ): @@ -46,205 +65,197 @@ def __init__( self.udB = udB kwargs["api_id"] = api_id or Var.API_ID kwargs["api_hash"] = api_hash or Var.API_HASH - kwargs["base_logger"] = TelethonLogger + kwargs["base_logger"] = TelethonLogger # noqa: F405 super().__init__(session, **kwargs) self.run_in_loop(self.start_client(bot_token=bot_token)) self.dc_id = self.session.dc_id - def __repr__(self): - return f"" - - @property - def __dict__(self): - if self.me: - return self.me.to_dict() + def __repr__(self) -> str: + return f"" - async def start_client(self, **kwargs): - """function to start client""" + # ---------------------------------------------------------- login + async def start_client(self, **kwargs) -> None: + """Connect & cache identity information.""" if self._log_at: - self.logger.info("Trying to login.") + self.logger.info("Logging in to Telegram...") try: await self.start(**kwargs) except ApiIdInvalidError: - self.logger.critical("API ID and API_HASH combination does not match!") - - sys.exit() - except (AuthKeyDuplicatedError, EOFError) as er: + self.logger.critical("API_ID / API_HASH combination is invalid.") + sys.exit(1) + except (AuthKeyDuplicatedError, EOFError): + self.logger.critical("String session expired. Create a new one!") if self._handle_error: - self.logger.critical("String session expired. Create new!") - return sys.exit() - self.logger.critical("String session expired.") + sys.exit(1) + return except (AccessTokenExpiredError, AccessTokenInvalidError): - # AccessTokenError can only occur for Bot account - # And at Early Process, Its saved in DB. - self.udB.del_key("BOT_TOKEN") + # AccessTokenError can only occur for Bot accounts and the token + # was already persisted in the DB during startup. + if self.udB is not None: + self.udB.del_key("BOT_TOKEN") self.logger.critical( - "Bot token is expired or invalid. Create new from @Botfather and add in BOT_TOKEN env variable!" + "BOT_TOKEN is expired or invalid. Create a new one via @BotFather " + "and update the BOT_TOKEN env variable." ) - sys.exit() - # Save some stuff for later use... + sys.exit(1) + self.me = await self.get_me() if self.me.bot: - me = f"@{self.me.username}" + who = f"@{self.me.username}" else: - setattr(self.me, "phone", None) - me = self.full_name + # Userbots should never log/store their phone number. + self.me.phone = None + who = self.full_name if self._log_at: - self.logger.info(f"Logged in as {me}") + self.logger.info(f"Logged in as {who}.") self._bot = await self.is_bot() + # ------------------------------------------------------ fast IO async def fast_uploader(self, file, **kwargs): - """Upload files in a faster way""" + """Upload ``file`` using parallel MTProto senders. - import os - from pathlib import Path + Supported kwargs (all optional): + * ``filename`` – override the uploaded filename + * ``show_progress`` + ``event`` – display a progress bar + * ``use_cache`` – reuse previously uploaded ``InputFile`` (default True) + * ``to_delete`` – remove the local file after a successful upload + * ``message`` – progress-bar caption + """ start_time = time.time() path = Path(file) filename = kwargs.get("filename", path.name) - # Set to True and pass event to show progress bar. show_progress = kwargs.get("show_progress", False) - if show_progress: - event = kwargs["event"] - # Whether to use cached file for uploading or not + event = kwargs.get("event") if show_progress else None use_cache = kwargs.get("use_cache", True) - # Delete original file after uploading to_delete = kwargs.get("to_delete", False) message = kwargs.get("message", f"Uploading {filename}...") by_bot = self._bot + size = os.path.getsize(file) - # Don't show progress bar when file size is less than 5MB. if size < 5 * 2 ** 20: show_progress = False - if use_cache and self._cache and self._cache.get("upload_cache"): - for files in self._cache["upload_cache"]: + + if use_cache and self._cache.get("upload_cache"): + for cached in self._cache["upload_cache"]: if ( - files["size"] == size - and files["path"] == path - and files["name"] == filename - and files["by_bot"] == by_bot + cached["size"] == size + and cached["path"] == path + and cached["name"] == filename + and cached["by_bot"] == by_bot ): if to_delete: with contextlib.suppress(FileNotFoundError): os.remove(file) - return files["raw_file"], time.time() - start_time + return cached["raw_file"], time.time() - start_time from pyUltroid.fns.FastTelethon import upload_file from pyUltroid.fns.helper import progress - raw_file = None - while not raw_file: - with open(file, "rb") as f: - raw_file = await upload_file( - client=self, - file=f, - filename=filename, - progress_callback=( - lambda completed, total: self.loop.create_task( - progress(completed, total, event, start_time, message) - ) - ) - if show_progress - else None, + progress_cb = None + if show_progress and event is not None: + def progress_cb(completed, total): # noqa: WPS430 + asyncio.create_task( + progress(completed, total, event, start_time, message) ) - cache = { + + with open(file, "rb") as fp: + raw_file = await upload_file( + client=self, + file=fp, + filename=filename, + progress_callback=progress_cb, + ) + if not raw_file: + raise RuntimeError(f"Failed to upload {filename!r} via FastTelethon.") + + cache_entry = { "by_bot": by_bot, "size": size, "path": path, "name": filename, "raw_file": raw_file, } - if self._cache.get("upload_cache"): - self._cache["upload_cache"].append(cache) - else: - self._cache.update({"upload_cache": [cache]}) + self._cache.setdefault("upload_cache", []).append(cache_entry) + if to_delete: with contextlib.suppress(FileNotFoundError): os.remove(file) return raw_file, time.time() - start_time async def fast_downloader(self, file, **kwargs): - """Download files in a faster way""" - # Set to True and pass event to show progress bar. + """Download ``file`` (a ``Document``) using parallel MTProto senders.""" show_progress = kwargs.get("show_progress", False) - filename = kwargs.get("filename", "") - if show_progress: - event = kwargs["event"] - # Don't show progress bar when file size is less than 10MB. + filename = kwargs.get("filename") or "" + event = kwargs.get("event") if show_progress else None if file.size < 10 * 2 ** 20: show_progress = False - import mimetypes - - from telethon.tl.types import DocumentAttributeFilename from pyUltroid.fns.FastTelethon import download_file from pyUltroid.fns.helper import progress start_time = time.time() - # Auto-generate Filename if not filename: try: if isinstance(file.attributes[-1], DocumentAttributeFilename): filename = file.attributes[-1].file_name - except IndexError: - mimetype = file.mime_type - filename = ( - mimetype.split("/")[0] - + "-" - + str(round(start_time)) - + mimetypes.guess_extension(mimetype) - ) + except (IndexError, AttributeError): + mimetype = getattr(file, "mime_type", "application/octet-stream") + ext = mimetypes.guess_extension(mimetype) or "" + filename = f"{mimetype.split('/')[0]}-{round(start_time)}{ext}" + message = kwargs.get("message", f"Downloading {filename}...") - raw_file = None - while not raw_file: - with open(filename, "wb") as f: - raw_file = await download_file( - client=self, - location=file, - out=f, - progress_callback=( - lambda completed, total: self.loop.create_task( - progress(completed, total, event, start_time, message) - ) - ) - if show_progress - else None, + progress_cb = None + if show_progress and event is not None: + def progress_cb(completed, total): # noqa: WPS430 + asyncio.create_task( + progress(completed, total, event, start_time, message) ) + + with open(filename, "wb") as fp: + raw_file = await download_file( + client=self, + location=file, + out=fp, + progress_callback=progress_cb, + ) + if not raw_file: + raise RuntimeError(f"Failed to download {filename!r} via FastTelethon.") return raw_file, time.time() - start_time - def run_in_loop(self, function): - """run inside asyncio loop""" - return self.loop.run_until_complete(function) + # ---------------------------------------------------- loop helpers + def run_in_loop(self, coro): + """Run ``coro`` to completion on the client's event loop.""" + return self.loop.run_until_complete(coro) - def run(self): - """run asyncio loop""" + def run(self) -> None: + """Block until the client disconnects.""" self.run_until_disconnected() - def add_handler(self, func, *args, **kwargs): - """Add new event handler, ignoring if exists""" - if func in [_[0] for _ in self.list_event_handlers()]: + def add_handler(self, func, *args, **kwargs) -> None: + """Add an event handler, but ignore duplicate registrations.""" + if any(existing is func for existing, _ in self.list_event_handlers()): return self.add_event_handler(func, *args, **kwargs) + # ---------------------------------------------------- convenience @property def utils(self): return telethon_utils @property - def full_name(self): - """full name of Client""" - return self.utils.get_display_name(self.me) + def full_name(self) -> str: + """Display name of the underlying account.""" + return telethon_utils.get_display_name(self.me) @property - def uid(self): - """Client's user id""" + def uid(self) -> int: + """The numeric Telegram ID of the underlying account.""" return self.me.id - def to_dict(self): - return dict(inspect.getmembers(self)) - async def parse_id(self, text): - with contextlib.suppress(ValueError): + """Convert ``text`` (a username, id or link) into a peer id.""" + with contextlib.suppress(TypeError, ValueError): text = int(text) return await self.get_peer_id(text) diff --git a/pyUltroid/startup/__init__.py b/pyUltroid/startup/__init__.py index 0c0e1459c7..b178666b63 100644 --- a/pyUltroid/startup/__init__.py +++ b/pyUltroid/startup/__init__.py @@ -58,7 +58,7 @@ def where_hosted(): _, v, __ = platform.python_version_tuple() - if int(v) < 10: + if sys.version_info < (3, 10): from ._extra import _fix_logging _fix_logging(FileHandler) diff --git a/pyUltroid/startup/_database.py b/pyUltroid/startup/_database.py index d6845c9199..45db8a430f 100644 --- a/pyUltroid/startup/_database.py +++ b/pyUltroid/startup/_database.py @@ -7,160 +7,178 @@ import ast import os -import sys +from typing import Any, Iterable, Optional from .. import run_as_module -from . import * +from . import * # noqa: F401,F403 if run_as_module: from ..configs import Var - -Redis = MongoClient = psycopg2 = Database = None +# Each backend is imported lazily so missing optional deps don't blow up +# startup before we even know which DB the user picked. +Redis = MongoClient = psycopg2 = LocalDatabase = None if Var.REDIS_URI or Var.REDISHOST: try: - from redis import Redis - except ImportError: - LOGS.info("Installing 'redis' for database.") - os.system(f"{sys.executable} -m pip install -q redis hiredis") - from redis import Redis + from redis import Redis # type: ignore[assignment] + except ImportError as err: # pragma: no cover - clearer error than os.system + LOGS.critical( # noqa: F405 + "REDIS_URI is configured but the 'redis' package is not installed. " + "Run: pip install -U redis hiredis" + ) + raise SystemExit(1) from err elif Var.MONGO_URI: try: - from pymongo import MongoClient - except ImportError: - LOGS.info("Installing 'pymongo' for database.") - os.system(f"{sys.executable} -m pip install -q pymongo[srv]") - from pymongo import MongoClient + from pymongo import MongoClient # type: ignore[assignment] + except ImportError as err: # pragma: no cover + LOGS.critical( # noqa: F405 + "MONGO_URI is configured but the 'pymongo' package is not installed. " + "Run: pip install -U 'pymongo[srv]'" + ) + raise SystemExit(1) from err elif Var.DATABASE_URL: try: - import psycopg2 - except ImportError: - LOGS.info("Installing 'pyscopg2' for database.") - os.system(f"{sys.executable} -m pip install -q psycopg2-binary") - import psycopg2 + import psycopg2 # type: ignore[assignment] + from psycopg2 import sql as _pgsql # type: ignore[assignment] + except ImportError as err: # pragma: no cover + LOGS.critical( # noqa: F405 + "DATABASE_URL is configured but 'psycopg2' is not installed. " + "Run: pip install -U psycopg2-binary" + ) + raise SystemExit(1) from err else: try: - from localdb import Database - except ImportError: - LOGS.info("Using local file as database.") - os.system(f"{sys.executable} -m pip install -q localdb.json") - from localdb import Database + from localdb import Database as LocalDatabase # type: ignore[assignment] + except ImportError as err: # pragma: no cover + LOGS.critical( # noqa: F405 + "No database backend is configured. Either set REDIS_URI / MONGO_URI / " + "DATABASE_URL, or install 'localdb.json' for a file-backed fallback." + ) + raise SystemExit(1) from err -# --------------------------------------------------------------------------------------------- # +# --------------------------------------------------------------------------- class _BaseDatabase: - def __init__(self, *args, **kwargs): - self._cache = {} + """Common cache + (de)serialisation logic shared by every backend.""" + + def __init__(self) -> None: + self._cache: dict[str, Any] = {} + + # ------------------------ helpers -------------------------------- + def _get_data(self, key: Optional[str] = None, data: Any = None) -> Any: + if key is not None: + data = self.get(str(key)) # type: ignore[attr-defined] + if isinstance(data, str): + try: + data = ast.literal_eval(data) + except (SyntaxError, ValueError): + pass + return data - def get_key(self, key): + # ------------------------ public API ----------------------------- + def get_key(self, key: str) -> Any: if key in self._cache: return self._cache[key] value = self._get_data(key) - self._cache.update({key: value}) + self._cache[key] = value return value - def re_cache(self): + def set_key(self, key: str, value: Any, cache_only: bool = False) -> Any: + value = self._get_data(data=value) + self._cache[key] = value + if cache_only: + return None + return self.set(str(key), str(value)) # type: ignore[attr-defined] + + def del_key(self, key: str) -> bool: + self._cache.pop(key, None) + self.delete(key) # type: ignore[attr-defined] + return True + + def re_cache(self) -> None: self._cache.clear() - for key in self.keys(): - self._cache.update({key: self.get_key(key)}) + for key in self.keys(): # type: ignore[attr-defined] + self._cache[key] = self.get_key(key) - def ping(self): + def rename(self, key1: str, key2: str) -> int: + value = self.get_key(key1) + if value: + self.del_key(key1) + self.set_key(key2, value) + return 0 return 1 - @property - def usage(self): - return 0 + def ping(self) -> int: + return 1 - def keys(self): + def keys(self) -> Iterable[str]: return [] - def del_key(self, key): - if key in self._cache: - del self._cache[key] - self.delete(key) - return True - - def _get_data(self, key=None, data=None): - if key: - data = self.get(str(key)) - if data and isinstance(data, str): - try: - data = ast.literal_eval(data) - except BaseException: - pass - return data + @property + def usage(self) -> int: + return 0 - def set_key(self, key, value, cache_only=False): - value = self._get_data(data=value) - self._cache[key] = value - if cache_only: - return - return self.set(str(key), str(value)) - def rename(self, key1, key2): - _ = self.get_key(key1) - if _: - self.del_key(key1) - self.set_key(key2, _) - return 0 - return 1 +# ----------------------------- Mongo ----------------------------------- class MongoDB(_BaseDatabase): - def __init__(self, key, dbname="UltroidDB"): + def __init__(self, key, dbname: str = "UltroidDB"): self.dB = MongoClient(key, serverSelectionTimeoutMS=5000) self.db = self.dB[dbname] super().__init__() - def __repr__(self): - return f"" + def __repr__(self) -> str: + return f"" @property - def name(self): + def name(self) -> str: return "Mongo" @property - def usage(self): + def usage(self) -> int: return self.db.command("dbstats")["dataSize"] - def ping(self): - if self.dB.server_info(): - return True + def ping(self) -> bool: + return bool(self.dB.server_info()) def keys(self): return self.db.list_collection_names() - def set(self, key, value): + def set(self, key: str, value: Any) -> bool: if key in self.keys(): self.db[key].replace_one({"_id": key}, {"value": str(value)}) else: self.db[key].insert_one({"_id": key, "value": str(value)}) return True - def delete(self, key): + def delete(self, key: str) -> None: self.db.drop_collection(key) - def get(self, key): - if x := self.db[key].find_one({"_id": key}): - return x["value"] + def get(self, key: str) -> Any: + doc = self.db[key].find_one({"_id": key}) + return doc["value"] if doc else None - def flushall(self): + def flushall(self) -> bool: self.dB.drop_database("UltroidDB") self._cache.clear() return True -# --------------------------------------------------------------------------------------------- # - -# Thanks to "Akash Pattnaik" / @BLUE-DEVIL1134 -# for SQL Implementation in Ultroid. +# ------------------------------ SQL ------------------------------------ # -# Please use https://elephantsql.com/ ! +# Thanks to "Akash Pattnaik" / @BLUE-DEVIL1134 for the original SQL impl. +# This rewrite uses ``psycopg2.sql`` composition everywhere so column names +# (which are derived from arbitrary user-controlled DB keys) can no longer +# be used as a SQL-injection vector. +# https://www.psycopg.org/docs/sql.html class SqlDB(_BaseDatabase): - def __init__(self, url): + _TABLE = "ultroid" + + def __init__(self, url: str) -> None: self._url = url self._connection = None self._cursor = None @@ -169,120 +187,150 @@ def __init__(self, url): self._connection.autocommit = True self._cursor = self._connection.cursor() self._cursor.execute( - "CREATE TABLE IF NOT EXISTS Ultroid (ultroidCli varchar(70))" + _pgsql.SQL( + "CREATE TABLE IF NOT EXISTS {tbl} (ultroidCli varchar(70))" + ).format(tbl=_pgsql.Identifier(self._TABLE)) ) - except Exception as error: - LOGS.exception(error) - LOGS.info("Invaid SQL Database") + except Exception as err: + LOGS.exception(err) # noqa: F405 + LOGS.critical("Invalid SQL database, exiting.") # noqa: F405 if self._connection: self._connection.close() - sys.exit() + raise SystemExit(1) from err super().__init__() @property - def name(self): + def name(self) -> str: return "SQL" @property - def usage(self): + def usage(self) -> int: self._cursor.execute( - "SELECT pg_size_pretty(pg_relation_size('Ultroid')) AS size" + _pgsql.SQL( + "SELECT pg_size_pretty(pg_relation_size({tbl})) AS size" + ).format(tbl=_pgsql.Literal(self._TABLE)) ) data = self._cursor.fetchall() return int(data[0][0].split()[0]) def keys(self): self._cursor.execute( - "SELECT column_name FROM information_schema.columns WHERE table_schema = 'public' AND table_name = 'ultroid'" - ) # case sensitive - data = self._cursor.fetchall() - return [_[0] for _ in data] + "SELECT column_name FROM information_schema.columns " + "WHERE table_schema = 'public' AND table_name = %s", + (self._TABLE,), + ) + return [row[0] for row in self._cursor.fetchall()] - def get(self, variable): + def get(self, variable: str) -> Any: try: - self._cursor.execute(f"SELECT {variable} FROM Ultroid") + self._cursor.execute( + _pgsql.SQL("SELECT {col} FROM {tbl}").format( + col=_pgsql.Identifier(variable), + tbl=_pgsql.Identifier(self._TABLE), + ) + ) except psycopg2.errors.UndefinedColumn: return None - data = self._cursor.fetchall() - if not data: - return None - if len(data) >= 1: - for i in data: - if i[0]: - return i[0] + for row in self._cursor.fetchall(): + if row[0]: + return row[0] + return None - def set(self, key, value): + def set(self, key: str, value: Any) -> bool: try: - self._cursor.execute(f"ALTER TABLE Ultroid DROP COLUMN IF EXISTS {key}") + self._cursor.execute( + _pgsql.SQL("ALTER TABLE {tbl} DROP COLUMN IF EXISTS {col}").format( + tbl=_pgsql.Identifier(self._TABLE), + col=_pgsql.Identifier(key), + ) + ) except (psycopg2.errors.UndefinedColumn, psycopg2.errors.SyntaxError): pass - except BaseException as er: - LOGS.exception(er) - self._cache.update({key: value}) - self._cursor.execute(f"ALTER TABLE Ultroid ADD {key} TEXT") - self._cursor.execute(f"INSERT INTO Ultroid ({key}) values (%s)", (str(value),)) + except Exception as err: # pragma: no cover - defensive + LOGS.exception(err) # noqa: F405 + self._cache[key] = value + self._cursor.execute( + _pgsql.SQL("ALTER TABLE {tbl} ADD {col} TEXT").format( + tbl=_pgsql.Identifier(self._TABLE), col=_pgsql.Identifier(key) + ) + ) + self._cursor.execute( + _pgsql.SQL("INSERT INTO {tbl} ({col}) VALUES (%s)").format( + tbl=_pgsql.Identifier(self._TABLE), col=_pgsql.Identifier(key) + ), + (str(value),), + ) return True - def delete(self, key): + def delete(self, key: str) -> bool: try: - self._cursor.execute(f"ALTER TABLE Ultroid DROP COLUMN {key}") + self._cursor.execute( + _pgsql.SQL("ALTER TABLE {tbl} DROP COLUMN {col}").format( + tbl=_pgsql.Identifier(self._TABLE), col=_pgsql.Identifier(key) + ) + ) except psycopg2.errors.UndefinedColumn: return False return True - def flushall(self): + def flushall(self) -> bool: self._cache.clear() - self._cursor.execute("DROP TABLE Ultroid") self._cursor.execute( - "CREATE TABLE IF NOT EXISTS Ultroid (ultroidCli varchar(70))" + _pgsql.SQL("DROP TABLE {tbl}").format(tbl=_pgsql.Identifier(self._TABLE)) + ) + self._cursor.execute( + _pgsql.SQL( + "CREATE TABLE IF NOT EXISTS {tbl} (ultroidCli varchar(70))" + ).format(tbl=_pgsql.Identifier(self._TABLE)) ) return True -# --------------------------------------------------------------------------------------------- # +# ----------------------------- Redis ----------------------------------- class RedisDB(_BaseDatabase): def __init__( self, - host, - port, - password, - platform="", - logger=LOGS, + host: Optional[str], + port: Optional[int], + password: Optional[str], + platform: str = "", + logger=None, *args, **kwargs, - ): + ) -> None: + logger = logger or LOGS # noqa: F405 if host and ":" in host: - spli_ = host.split(":") - host = spli_[0] - port = int(spli_[-1]) + head, _, tail = host.rpartition(":") + host = head + try: + port = int(tail) + except ValueError: + logger.error("Invalid Redis port number in REDIS_URI.") + raise SystemExit(1) if host.startswith("http"): - logger.error("Your REDIS_URI should not start with http !") - import sys - - sys.exit() + logger.error("REDIS_URI should not start with http(s)://") + raise SystemExit(1) elif not host or not port: - logger.error("Port Number not found") - import sys + logger.error("Redis host / port not configured properly.") + raise SystemExit(1) - sys.exit() - kwargs["host"] = host - kwargs["password"] = password - kwargs["port"] = port + kwargs.update({"host": host, "password": password, "port": port}) if platform.lower() == "qovery" and not host: - var, hash_, host, password = "", "", "", "" - for vars_ in os.environ: - if vars_.startswith("QOVERY_REDIS_") and vars_.endswith("_HOST"): - var = vars_ - if var: - hash_ = var.split("_", maxsplit=2)[1].split("_")[0] - if hash_: - kwargs["host"] = os.environ.get(f"QOVERY_REDIS_{hash_}_HOST") - kwargs["port"] = os.environ.get(f"QOVERY_REDIS_{hash_}_PORT") - kwargs["password"] = os.environ.get(f"QOVERY_REDIS_{hash_}_PASSWORD") + for env_name in os.environ: + if env_name.startswith("QOVERY_REDIS_") and env_name.endswith("_HOST"): + hash_ = env_name.split("_", maxsplit=2)[1].split("_")[0] + kwargs["host"] = os.environ.get(f"QOVERY_REDIS_{hash_}_HOST") + kwargs["port"] = os.environ.get(f"QOVERY_REDIS_{hash_}_PORT") + kwargs["password"] = os.environ.get( + f"QOVERY_REDIS_{hash_}_PASSWORD" + ) + break + self.db = Redis(**kwargs) + # Delegate basic ops directly to the redis client – the fast path. self.set = self.db.set self.get = self.db.get self.keys = self.db.keys @@ -290,38 +338,41 @@ def __init__( super().__init__() @property - def name(self): + def name(self) -> str: return "Redis" @property - def usage(self): + def usage(self) -> int: return sum(self.db.memory_usage(x) for x in self.keys()) -# --------------------------------------------------------------------------------------------- # +# ----------------------------- LocalDB --------------------------------- class LocalDB(_BaseDatabase): - def __init__(self): - self.db = Database("ultroid") + def __init__(self) -> None: + self.db = LocalDatabase("ultroid") self.get = self.db.get self.set = self.db.set self.delete = self.db.delete super().__init__() @property - def name(self): + def name(self) -> str: return "LocalDB" def keys(self): return self._cache.keys() - def __repr__(self): - return f"" + def __repr__(self) -> str: + return f"" + + +# --------------------------------------------------------------------------- def UltroidDB(): - _er = False + """Return the configured database backend, exiting if none is usable.""" from .. import HOSTED_ON try: @@ -335,18 +386,17 @@ def UltroidDB(): socket_timeout=5, retry_on_timeout=True, ) - elif MongoClient: + if MongoClient: return MongoDB(Var.MONGO_URI) - elif psycopg2: + if psycopg2: return SqlDB(Var.DATABASE_URL) - else: - LOGS.critical( - "No DB requirement fullfilled!\nPlease install redis, mongo or sql dependencies...\nTill then using local file as database." - ) - return LocalDB() - except BaseException as err: - LOGS.exception(err) - exit() - - -# --------------------------------------------------------------------------------------------- # + LOGS.critical( # noqa: F405 + "No remote DB available – falling back to LocalDB. " + "For production, install redis / pymongo / psycopg2-binary." + ) + return LocalDB() + except SystemExit: + raise + except Exception as err: + LOGS.exception(err) # noqa: F405 + raise SystemExit(1) from err diff --git a/pyUltroid/startup/_extra.py b/pyUltroid/startup/_extra.py index 30c6376b40..d413d49389 100644 --- a/pyUltroid/startup/_extra.py +++ b/pyUltroid/startup/_extra.py @@ -6,23 +6,31 @@ # . # https://bugs.python.org/issue26789 -# 'open' not defined has been fixed in Python3.10 -# for other older versions, something need to be done. +# 'open' not defined has been fixed in Python 3.10. + +import builtins def _fix_logging(handler): + """Make ``FileHandler.open`` UTF-8 on legacy Python (< 3.10).""" handler._builtin_open = open def _new_open(self): open_func = self._builtin_open - return open_func(self.baseFilename, self.mode) + return open_func(self.baseFilename, self.mode, encoding="utf-8") - setattr(handler, "_open", _new_open) + handler._open = _new_open def _ask_input(): - # Ask for Input even on Vps and other platforms. + """Block any ``input()`` call so the bot can run unattended on VPS/CI. + + The previous implementation poked at ``__builtins__`` as a dict, which + only works when the calling module is ``__main__``. In submodules + ``__builtins__`` is the module itself, so the patch silently failed. + """ + def new_input(*args, **kwargs): - raise EOFError("args=" + str(args) + ", kwargs=" + str(kwargs)) + raise EOFError(f"args={args}, kwargs={kwargs}") - __builtins__["input"] = new_input + builtins.input = new_input diff --git a/pyUltroid/startup/connections.py b/pyUltroid/startup/connections.py index 8a1f340633..648bcdfe23 100644 --- a/pyUltroid/startup/connections.py +++ b/pyUltroid/startup/connections.py @@ -9,18 +9,20 @@ import ipaddress import struct import sys +from typing import Optional from telethon.errors.rpcerrorlist import AuthKeyDuplicatedError from telethon.sessions.string import _STRUCT_PREFORMAT, CURRENT_VERSION, StringSession from ..configs import Var -from . import * +from . import * # noqa: F401,F403 from .BaseClient import UltroidClient +# Pyrogram session formats keyed by length, with the corresponding struct +# layout. (Same numbers Pyrogram itself uses.) _PYRO_FORM = {351: ">B?256sI?", 356: ">B?256sQ?", 362: ">BI?256sQ?"} # https://github.com/pyrogram/pyrogram/blob/master/docs/source/faq/what-are-the-ip-addresses-of-telegram-data-centers.rst - DC_IPV4 = { 1: "149.154.175.53", 2: "149.154.167.51", @@ -30,65 +32,90 @@ } -def validate_session(session, logger=LOGS, _exit=True): - from strings import get_string +def _fail(logger, code: str, _exit: bool) -> None: + """Log a localised error and optionally terminate the process.""" + try: + from strings import get_string + + msg = get_string(code) + except Exception: + msg = code + logger.critical(msg) + if _exit: + sys.exit(1) + - if session: - # Telethon Session - if session.startswith(CURRENT_VERSION): - if len(session.strip()) != 353: - logger.exception(get_string("py_c1")) - sys.exit() - return StringSession(session) +def validate_session(session: Optional[str], logger=LOGS, _exit: bool = True): # noqa: F405 + """Validate a string session, converting Pyrogram sessions to Telethon.""" + if not session: + _fail(logger, "py_c2", _exit) + return None - # Pyrogram Session - elif len(session) in _PYRO_FORM.keys(): + # Native Telethon string session. + if session.startswith(CURRENT_VERSION): + if len(session.strip()) != 353: + _fail(logger, "py_c1", _exit) + return None + return StringSession(session) + + # Pyrogram session – translate it to a Telethon one. + if len(session) in _PYRO_FORM: + try: data_ = struct.unpack( _PYRO_FORM[len(session)], base64.urlsafe_b64decode(session + "=" * (-len(session) % 4)), ) - if len(session) in [351, 356]: - auth_id = 2 - else: - auth_id = 3 - dc_id, auth_key = data_[0], data_[auth_id] - return StringSession( - CURRENT_VERSION - + base64.urlsafe_b64encode( - struct.pack( - _STRUCT_PREFORMAT.format(4), - dc_id, - ipaddress.ip_address(DC_IPV4[dc_id]).packed, - 443, - auth_key, - ) - ).decode("ascii") - ) - else: - logger.exception(get_string("py_c1")) - if _exit: - sys.exit() - logger.exception(get_string("py_c2")) - if _exit: - sys.exit() + except Exception as err: + logger.error(f"Could not decode Pyrogram session: {err}") + _fail(logger, "py_c1", _exit) + return None + + auth_id = 2 if len(session) in {351, 356} else 3 + dc_id, auth_key = data_[0], data_[auth_id] + return StringSession( + CURRENT_VERSION + + base64.urlsafe_b64encode( + struct.pack( + _STRUCT_PREFORMAT.format(4), + dc_id, + ipaddress.ip_address(DC_IPV4[dc_id]).packed, + 443, + auth_key, + ) + ).decode("ascii") + ) + + _fail(logger, "py_c1", _exit) + return None def vc_connection(udB, ultroid_bot): - from strings import get_string + """Return a Telethon client for Voice/Video Chat operations. + + Falls back to the main userbot client if no dedicated VC session was + configured (or if the session is bad). + """ + try: + from strings import get_string + except Exception: + def get_string(code): + return code VC_SESSION = Var.VC_SESSION or udB.get_key("VC_SESSION") - if VC_SESSION and VC_SESSION != Var.SESSION: - LOGS.info("Starting up VcClient.") - try: - return UltroidClient( - validate_session(VC_SESSION, _exit=False), - log_attempt=False, - exit_on_error=False, - ) - except (AuthKeyDuplicatedError, EOFError): - LOGS.info(get_string("py_c3")) - udB.del_key("VC_SESSION") - except Exception as er: - LOGS.info("While creating Client for VC.") - LOGS.exception(er) + if not VC_SESSION or VC_SESSION == Var.SESSION: + return ultroid_bot + + LOGS.info("Starting up VcClient.") # noqa: F405 + try: + return UltroidClient( + validate_session(VC_SESSION, _exit=False), + log_attempt=False, + exit_on_error=False, + ) + except (AuthKeyDuplicatedError, EOFError): + LOGS.info(get_string("py_c3")) # noqa: F405 + udB.del_key("VC_SESSION") + except Exception as err: + LOGS.error("Failed to create a VC client.") # noqa: F405 + LOGS.exception(err) # noqa: F405 return ultroid_bot diff --git a/pyUltroid/startup/loader.py b/pyUltroid/startup/loader.py index 6294f94ec8..7c192c77a5 100644 --- a/pyUltroid/startup/loader.py +++ b/pyUltroid/startup/loader.py @@ -67,31 +67,51 @@ def load_other_plugins(addons=None, pmbot=None, manager=None, vcbot=None): # for addons if addons: - if url := udB.get_key("ADDONS_URL"): - subprocess.run(f"git clone -q {url} addons", shell=True) + url = udB.get_key("ADDONS_URL") + if url: + subprocess.run(["git", "clone", "-q", url, "addons"], check=False) if os.path.exists("addons") and not os.path.exists("addons/.git"): rmtree("addons") if not os.path.exists("addons"): subprocess.run( - f"git clone -q -b {Repo().active_branch} https://github.com/TeamUltroid/UltroidAddons.git addons", - shell=True, + [ + "git", + "clone", + "-q", + "-b", + Repo().active_branch.name, + "https://github.com/TeamUltroid/UltroidAddons.git", + "addons", + ], + check=False, ) else: - subprocess.run("cd addons && git pull -q && cd ..", shell=True) + subprocess.run(["git", "-C", "addons", "pull", "-q"], check=False) if not os.path.exists("addons"): subprocess.run( - "git clone -q https://github.com/TeamUltroid/UltroidAddons.git addons", - shell=True, + [ + "git", + "clone", + "-q", + "https://github.com/TeamUltroid/UltroidAddons.git", + "addons", + ], + check=False, ) if os.path.exists("addons/addons.txt"): - # generally addons req already there so it won't take much time - # subprocess.run( - # "rm -rf /usr/local/lib/python3.*/site-packages/pip/_vendor/.wh*" - # ) subprocess.run( - f"{sys.executable} -m pip install --no-cache-dir -q -r ./addons/addons.txt", - shell=True, + [ + sys.executable, + "-m", + "pip", + "install", + "--no-cache-dir", + "-q", + "-r", + "./addons/addons.txt", + ], + check=False, ) _exclude = udB.get_key("EXCLUDE_ADDONS") @@ -123,12 +143,18 @@ def load_other_plugins(addons=None, pmbot=None, manager=None, vcbot=None): if os.path.exists("vcbot"): if os.path.exists("vcbot/.git"): - subprocess.run("cd vcbot && git pull", shell=True) + subprocess.run(["git", "-C", "vcbot", "pull"], check=False) else: rmtree("vcbot") if not os.path.exists("vcbot"): subprocess.run( - "git clone https://github.com/TeamUltroid/VcBot vcbot", shell=True + [ + "git", + "clone", + "https://github.com/TeamUltroid/VcBot", + "vcbot", + ], + check=False, ) try: if not os.path.exists("vcbot/downloads"): diff --git a/pyUltroid/version.py b/pyUltroid/version.py index 9752b2c2ec..ed51dfd817 100644 --- a/pyUltroid/version.py +++ b/pyUltroid/version.py @@ -1,2 +1,2 @@ -__version__ = "2026.04.03" -ultroid_version = "2.1.2" +__version__ = "2026.01.28" +ultroid_version = "2.1.3"