diff --git a/plugins/calculator.py b/plugins/calculator.py index 5e318e6a3d..9e12039027 100644 --- a/plugins/calculator.py +++ b/plugins/calculator.py @@ -10,10 +10,48 @@ __doc__ = get_help("help_calculator") +import ast +import operator import re from . import Button, asst, callback, get_string, in_pattern, udB, ultroid_cmd +# Safe math operations for calculator - no arbitrary code execution +_SAFE_MATH_OPS = { + ast.Add: operator.add, + ast.Sub: operator.sub, + ast.Mult: operator.mul, + ast.Div: operator.truediv, + ast.Mod: operator.mod, + ast.Pow: operator.pow, + ast.USub: operator.neg, + ast.UAdd: operator.pos, +} + + +def _safe_eval_math(expr): + """Safely evaluate a mathematical expression without using eval().""" + try: + tree = ast.parse(expr, mode="eval") + except SyntaxError: + raise ValueError(f"Invalid expression: {expr}") + + def _eval_node(node): + if isinstance(node, ast.Expression): + return _eval_node(node.body) + elif isinstance(node, ast.Constant) and isinstance(node.value, (int, float)): + return node.value + elif isinstance(node, ast.BinOp) and type(node.op) in _SAFE_MATH_OPS: + left = _eval_node(node.left) + right = _eval_node(node.right) + return _SAFE_MATH_OPS[type(node.op)](left, right) + elif isinstance(node, ast.UnaryOp) and type(node.op) in _SAFE_MATH_OPS: + return _SAFE_MATH_OPS[type(node.op)](_eval_node(node.operand)) + else: + raise ValueError(f"Unsupported expression") + + return _eval_node(tree) + CALC = {} m = [ @@ -105,7 +143,7 @@ async def _(e): if get: if get.endswith(("*", ".", "/", "-", "+")): get = get[:-1] - out = eval(get) + out = _safe_eval_math(get) try: num = float(out) await e.answer(f"Answer : {num}", cache_time=0, alert=True) diff --git a/plugins/nightmode.py b/plugins/nightmode.py index 1d2792d22c..d2eb4bea92 100644 --- a/plugins/nightmode.py +++ b/plugins/nightmode.py @@ -29,6 +29,8 @@ Ex- `nmtime 01 00 06 30` """ +import ast + from . import LOGS try: @@ -125,7 +127,7 @@ async def open_grp(): async def close_grp(): __, _, h2, m2 = 0, 0, 7, 0 if udB.get_key("NIGHT_TIME"): - _, __, h2, m2 = eval(udB.get_key("NIGHT_TIME")) + _, __, h2, m2 = ast.literal_eval(udB.get_key("NIGHT_TIME")) for chat in keym.get(): try: await ultroid_bot( @@ -148,7 +150,7 @@ async def close_grp(): try: h1, m1, h2, m2 = 0, 0, 7, 0 if udB.get_key("NIGHT_TIME"): - h1, m1, h2, m2 = eval(udB.get_key("NIGHT_TIME")) + h1, m1, h2, m2 = ast.literal_eval(udB.get_key("NIGHT_TIME")) sch = AsyncIOScheduler() sch.add_job(close_grp, trigger="cron", hour=h1, minute=m1) sch.add_job(open_grp, trigger="cron", hour=h2, minute=m2) diff --git a/pyUltroid/fns/executor.py b/pyUltroid/fns/executor.py index 06226842b5..63a3020d16 100644 --- a/pyUltroid/fns/executor.py +++ b/pyUltroid/fns/executor.py @@ -41,8 +41,8 @@ async def run(self, *args) -> int: def terminate(self, pid: int) -> bool: try: - self._processes.pop(pid) - self._processes[pid].kill() + process = self._processes.pop(pid) + process.kill() return True except KeyError: return False @@ -65,12 +65,11 @@ async def error(self, pid: int) -> str: error.append(err) return "\n".join(error) - @property - def _auto_remove_processes(self) -> None: - while self._processes: - for proc in self._processes.keys(): - if proc.returncode is not None: # process is still running - try: - self._processes.pop(proc) - except KeyError: - pass + def cleanup_finished_processes(self) -> None: + """Remove finished processes from the tracking dict.""" + finished = [ + pid for pid, proc in self._processes.items() + if proc.returncode is not None + ] + for pid in finished: + self._processes.pop(pid, None) diff --git a/pyUltroid/fns/helper.py b/pyUltroid/fns/helper.py index c51ab80e22..ce518c56ba 100644 --- a/pyUltroid/fns/helper.py +++ b/pyUltroid/fns/helper.py @@ -64,11 +64,14 @@ from .FastTelethon import upload_file as uploadable +_shared_executor = ThreadPoolExecutor(max_workers=multiprocessing.cpu_count() * 5) + + def run_async(function): @wraps(function) async def wrapper(*args, **kwargs): return await asyncio.get_event_loop().run_in_executor( - ThreadPoolExecutor(max_workers=multiprocessing.cpu_count() * 5), + _shared_executor, partial(function, *args, **kwargs), ) @@ -453,7 +456,8 @@ def mediainfo(media): i = str(media.document.attributes[0]) if "supports_streaming=True" in i: m = "video" - m = "video as doc" + else: + m = "video as doc" else: m = "video" elif "audio" in mim: diff --git a/pyUltroid/fns/tools.py b/pyUltroid/fns/tools.py index b1f75831bb..b296fb5bd8 100644 --- a/pyUltroid/fns/tools.py +++ b/pyUltroid/fns/tools.py @@ -5,6 +5,7 @@ # PLease read the GNU Affero General Public License in # . +import ast import json import math import os @@ -94,7 +95,10 @@ def json_parser(data, indent=None, ascii=False): if indent: parsed = json.dumps(data, indent=indent, ensure_ascii=ascii) except JSONDecodeError: - parsed = eval(data) + try: + parsed = ast.literal_eval(data) + except (ValueError, SyntaxError): + parsed = data return parsed @@ -1049,11 +1053,8 @@ def recycle_type(exte): def _get_value(stri): try: - value = eval(stri.strip()) - except Exception as er: - from .. import LOGS - - LOGS.debug(er) + value = ast.literal_eval(stri.strip()) + except (ValueError, SyntaxError): value = stri.strip() return value diff --git a/pyUltroid/startup/BaseClient.py b/pyUltroid/startup/BaseClient.py index 23b451afe8..9b6c328280 100644 --- a/pyUltroid/startup/BaseClient.py +++ b/pyUltroid/startup/BaseClient.py @@ -54,10 +54,11 @@ def __init__( def __repr__(self): return f"" - @property - def __dict__(self): + def as_dict(self): + """Return the client's user info as a dict.""" if self.me: return self.me.to_dict() + return {} async def start_client(self, **kwargs): """function to start client""" @@ -132,8 +133,9 @@ async def fast_uploader(self, file, **kwargs): from pyUltroid.fns.FastTelethon import upload_file from pyUltroid.fns.helper import progress + max_retries = 3 raw_file = None - while not raw_file: + for attempt in range(max_retries): with open(file, "rb") as f: raw_file = await upload_file( client=self, @@ -147,6 +149,11 @@ async def fast_uploader(self, file, **kwargs): if show_progress else None, ) + if raw_file: + break + self.logger.warning(f"Upload attempt {attempt + 1}/{max_retries} failed for {filename}") + if not raw_file: + raise RuntimeError(f"Failed to upload {filename} after {max_retries} attempts") cache = { "by_bot": by_bot, "size": size, @@ -196,8 +203,9 @@ async def fast_downloader(self, file, **kwargs): ) message = kwargs.get("message", f"Downloading {filename}...") + max_retries = 3 raw_file = None - while not raw_file: + for attempt in range(max_retries): with open(filename, "wb") as f: raw_file = await download_file( client=self, @@ -211,6 +219,11 @@ async def fast_downloader(self, file, **kwargs): if show_progress else None, ) + if raw_file: + break + self.logger.warning(f"Download attempt {attempt + 1}/{max_retries} failed for {filename}") + if not raw_file: + raise RuntimeError(f"Failed to download {filename} after {max_retries} attempts") return raw_file, time.time() - start_time def run_in_loop(self, function): diff --git a/pyUltroid/startup/_database.py b/pyUltroid/startup/_database.py index 2611a5c32d..d27d7685d4 100644 --- a/pyUltroid/startup/_database.py +++ b/pyUltroid/startup/_database.py @@ -7,11 +7,23 @@ import ast import os +import re +import subprocess import sys from .. import run_as_module from . import * +# Regex to validate SQL identifier names (column names, etc.) +_VALID_SQL_IDENTIFIER = re.compile(r"^[a-zA-Z_][a-zA-Z0-9_]*$") + + +def _sanitize_sql_key(key): + """Validate that a key is a safe SQL identifier to prevent SQL injection.""" + if not _VALID_SQL_IDENTIFIER.match(key): + raise ValueError(f"Invalid SQL identifier: {key!r}") + return key + if run_as_module: from ..configs import Var @@ -22,28 +34,28 @@ from redis import Redis except ImportError: LOGS.info("Installing 'redis' for database.") - os.system(f"{sys.executable} -m pip install -q redis hiredis") + subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "redis", "hiredis"]) from redis import Redis 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]") + subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "pymongo[srv]"]) from pymongo import MongoClient 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") + LOGS.info("Installing 'psycopg2' for database.") + subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "psycopg2-binary"]) import psycopg2 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") + subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "localdb.json"]) from localdb import Database # --------------------------------------------------------------------------------------------- # @@ -199,6 +211,7 @@ def keys(self): return [_[0] for _ in data] def get(self, variable): + variable = _sanitize_sql_key(variable) try: self._cursor.execute(f"SELECT {variable} FROM Ultroid") except psycopg2.errors.UndefinedColumn: @@ -212,6 +225,7 @@ def get(self, variable): return i[0] def set(self, key, value): + key = _sanitize_sql_key(key) try: self._cursor.execute(f"ALTER TABLE Ultroid DROP COLUMN IF EXISTS {key}") except (psycopg2.errors.UndefinedColumn, psycopg2.errors.SyntaxError): @@ -224,6 +238,7 @@ def set(self, key, value): return True def delete(self, key): + key = _sanitize_sql_key(key) try: self._cursor.execute(f"ALTER TABLE Ultroid DROP COLUMN {key}") except psycopg2.errors.UndefinedColumn: diff --git a/pyUltroid/startup/loader.py b/pyUltroid/startup/loader.py index 6294f94ec8..12742928d0 100644 --- a/pyUltroid/startup/loader.py +++ b/pyUltroid/startup/loader.py @@ -68,21 +68,20 @@ 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) + subprocess.run(["git", "clone", "-q", url, "addons"]) 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", str(Repo().active_branch), + "https://github.com/TeamUltroid/UltroidAddons.git", "addons"], ) else: - subprocess.run("cd addons && git pull -q && cd ..", shell=True) + subprocess.run(["git", "-C", "addons", "pull", "-q"]) 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"], ) if os.path.exists("addons/addons.txt"): # generally addons req already there so it won't take much time @@ -90,8 +89,7 @@ def load_other_plugins(addons=None, pmbot=None, manager=None, vcbot=None): # "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"], ) _exclude = udB.get_key("EXCLUDE_ADDONS") @@ -123,12 +121,12 @@ 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"]) 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"] ) try: if not os.path.exists("vcbot/downloads"): diff --git a/requirements.txt b/requirements.txt index acbc790d4b..75e387a839 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,12 +1,12 @@ # Important Requirements here. -telethon -gitpython +telethon>=1.42.0,<2.0 +gitpython>=3.1.40,<4.0 https://github.com/New-dev0/Telethon-Patch/archive/main.zip -python-decouple -python-dotenv -telegraph +python-decouple>=3.8,<4.0 +python-dotenv>=1.0.0,<2.0 +telegraph>=2.2.0,<3.0 enhancer -requests -aiohttp +requests>=2.31.0,<3.0 +aiohttp>=3.9.0,<4.0 catbox-uploader -cloudscraper \ No newline at end of file +cloudscraper>=1.2.71,<2.0 \ No newline at end of file diff --git a/resources/session/ssgen.py b/resources/session/ssgen.py index a4cbae34a0..311a326825 100644 --- a/resources/session/ssgen.py +++ b/resources/session/ssgen.py @@ -93,7 +93,7 @@ def telethon_session(): return except UserIsBotError: print("You are trying to Generate Session for your Bot's Account?") - print("Here is That!\n{ultroid.session.save()}\n\n") + print(f"Here is That!\n{ultroid.session.save()}\n\n") print("NOTE: You can't use that as User Session..") except ApiIdInvalidError: print(