fix: multiple issues in urls + catastrophic backtracking in Regex, freezes the bot

Author: TheHamkerCat

Dockerfile

@@ -10,6 +10,7 @@ ENV PYTHONUNBUFFERED=1
 RUN apt-get update -y && apt-get install -y --no-install-recommends \
     curl ca-certificates \
     git gcc build-essential \
+    iputils-ping \
     && rm -rf /var/lib/apt/lists/*
 
 # install uv

wbb/modules/misc.py

@@ -21,6 +21,7 @@
 OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 SOFTWARE.
 """
+
 import re
 import secrets
 import string
@@ -34,8 +35,6 @@
 from wbb.core.decorators.errors import capture_err
 from wbb.utils import random_line
 from wbb.utils.http import get
-from wbb.utils.json_prettify import json_prettify
-from wbb.utils.pastebin import paste
 
 __MODULE__ = "Misc"
 __HELP__ = """
@@ -61,9 +60,6 @@
     Translate A Message
     Ex: /tr en
 
-/json [URL]
-    Get parsed JSON response from a rest API.
-
 /arq
     Statistics Of ARQ API.
 
@@ -166,9 +162,7 @@ async def rtfm(_, message):
     await message.delete()
     if not message.reply_to_message:
         return await message.reply_text("Reply To A Message lol")
-    await message.reply_to_message.reply_text(
-        "Are You Lost? READ THE FUCKING DOCS!"
-    )
+    await message.reply_to_message.reply_text("Are You Lost? READ THE FUCKING DOCS!")
 
 
 @app.on_message(filters.command("runs"))
@@ -222,16 +216,12 @@ async def getid(client, message):
 @capture_err
 async def random(_, message):
     if len(message.command) != 2:
-        return await message.reply_text(
-            '"/random" Needs An Argurment.' " Ex: `/random 5`"
-        )
+        return await message.reply_text('"/random" Needs An Argurment. Ex: `/random 5`')
     length = message.text.split(None, 1)[1]
     try:
         if 1 < int(length) < 1000:
             alphabet = string.ascii_letters + string.digits
-            password = "".join(
-                secrets.choice(alphabet) for i in range(int(length))
-            )
+            password = "".join(secrets.choice(alphabet) for i in range(int(length)))
             await message.reply_text(f"`{password}`")
         else:
             await message.reply_text("Specify A Length Between 1-1000")
@@ -265,30 +255,8 @@ async def tr(_, message):
     await message.reply_text(result.result.translatedText)
 
 
-@app.on_message(filters.command("json"))
-@capture_err
-async def json_fetch(_, message):
-    if len(message.command) != 2:
-        return await message.reply_text("/json [URL]")
-    url = message.text.split(None, 1)[1]
-    m = await message.reply_text("Fetching")
-    try:
-        data = await get(url)
-        data = await json_prettify(data)
-        if len(data) < 4090:
-            await m.edit(data)
-        else:
-            link = await paste(data)
-            await m.edit(
-                f"[OUTPUT_TOO_LONG]({link})",
-                disable_web_page_preview=True,
-            )
-    except Exception as e:
-        await m.edit(str(e))
-
-
 @app.on_message(filters.command(["kickme", "banme"]))
 async def kickbanme(_, message):
     await message.reply_text(
-        "Haha, it doesn't work that way, You're stuck with everyone here."
+        "Haha, it doesn't work that way, You're stuck here with everyone."
     )

wbb/modules/regex.py

@@ -1,6 +1,7 @@
 # https://github.com/PaulSonOfLars/tgbot/blob/master/tg_bot/modules/sed.py
+import asyncio
+import multiprocessing as mp
 import re
-import sre_constants
 
 from pyrogram import filters
 
@@ -11,6 +12,62 @@
 __HELP__ = "**Usage:**\ns/foo/bar"
 
 DELIMITERS = ("/", ":", "|", "_")
+REGEX_TIMEOUT_SECONDS = 5
+
+
+def _regex_sub_worker(
+    pattern: str,
+    replacement: str,
+    source_text: str,
+    ignore_case: bool,
+    replace_all: bool,
+    result_queue,
+):
+    flags = re.I if ignore_case else 0
+    count = 0 if replace_all else 1
+
+    try:
+        result = re.sub(pattern, replacement, source_text, count=count, flags=flags)
+        result_queue.put(("ok", result))
+    except re.error:
+        result_queue.put(("regex_error", ""))
+
+
+def run_regex_with_timeout(
+    pattern: str,
+    replacement: str,
+    source_text: str,
+    ignore_case: bool,
+    replace_all: bool,
+) -> str:
+    result_queue = mp.Queue(maxsize=1)
+    process = mp.Process(
+        target=_regex_sub_worker,
+        args=(
+            pattern,
+            replacement,
+            source_text,
+            ignore_case,
+            replace_all,
+            result_queue,
+        ),
+    )
+    process.start()
+    process.join(REGEX_TIMEOUT_SECONDS)
+
+    if process.is_alive():
+        process.terminate()
+        process.join()
+        raise asyncio.TimeoutError
+
+    if result_queue.empty():
+        return ""
+
+    status, result = result_queue.get()
+    if status == "regex_error":
+        raise re.error("invalid regex")
+
+    return result
 
 
 @app.on_message(
@@ -31,31 +88,31 @@ async def sed(_, message):
             to_fix = message.reply_to_message.caption
         else:
             return
-        try:
-            repl, repl_with, flags = sed_result
-        except Exception:
+        if not sed_result:
             return
+        repl, repl_with, flags = sed_result
 
         if not repl:
             return await message.reply_text(
-                "You're trying to replace... " "nothing with something?"
+                "You're trying to replace... nothing with something?"
             )
 
         try:
             if infinite_checker(repl):
                 return await message.reply_text("Nice try -_-")
 
-            if "i" in flags and "g" in flags:
-                text = re.sub(repl, repl_with, to_fix, flags=re.I).strip()
-            elif "i" in flags:
-                text = re.sub(
-                    repl, repl_with, to_fix, count=1, flags=re.I
-                ).strip()
-            elif "g" in flags:
-                text = re.sub(repl, repl_with, to_fix).strip()
-            else:
-                text = re.sub(repl, repl_with, to_fix, count=1).strip()
-        except sre_constants.error:
+            text = await asyncio.to_thread(
+                run_regex_with_timeout,
+                repl,
+                repl_with,
+                to_fix,
+                "i" in flags,
+                "g" in flags,
+            )
+            text = text.strip()
+        except asyncio.TimeoutError:
+            return await message.reply_text("Regex took too long to compute.")
+        except re.error:
             return
 
         # empty string errors -_-
@@ -75,8 +132,9 @@ def infinite_checker(repl):
         r"\(.{1,}\)\{.{1,}(,)?\}\(.*\)(\+|\* |\{.*\})",
     ]
     for match in regex:
-        status = re.search(match, repl)
-        return bool(status)
+        if re.search(match, repl):
+            return True
+    return False
 
 
 def separate_sed(sed_string):

wbb/modules/rss.py

@@ -19,7 +19,11 @@
     remove_rss_feed,
     update_rss_feed,
 )
-from wbb.utils.functions import get_http_status_code, get_urls_from_text
+from wbb.utils.functions import (
+    get_http_status_code,
+    get_urls_from_text,
+    is_safe_url,
+)
 from wbb.utils.rss import Feed
 
 __MODULE__ = "RSS"
@@ -34,6 +38,13 @@
 """
 
 
+def get_parsed_feed_url(parsed, fallback: str) -> str:
+    href = parsed.get("href")
+    if isinstance(href, str) and href:
+        return href
+    return fallback
+
+
 async def rss_worker():
     log.info("RSS Worker started")
     while not await sleep(RSS_DELAY):
@@ -47,9 +58,20 @@ async def rss_worker():
             chat = _feed["chat_id"]
             try:
                 url = _feed["url"]
+                if not is_safe_url(url):
+                    await remove_rss_feed(chat)
+                    log.info(f"Removed RSS Feed from {chat} (Unsafe URL)")
+                    continue
+
                 last_title = _feed.get("last_title")
 
                 parsed = await loop.run_in_executor(None, parse, url)
+                final_url = get_parsed_feed_url(parsed, url)
+                if not is_safe_url(final_url):
+                    await remove_rss_feed(chat)
+                    log.info(f"Removed RSS Feed from {chat} (Unsafe redirect)")
+                    continue
+
                 feed = Feed(parsed)
 
                 if feed.title == last_title:
@@ -91,6 +113,9 @@ async def add_feed_func(_, m: Message):
         return await m.reply("[ERROR]: Invalid URL")
 
     url = urls[0]
+    if not is_safe_url(url):
+        return await m.reply("[ERROR]: URL is not allowed (SSRF protection).")
+
     status = await get_http_status_code(url)
     if status != 200:
         return await m.reply("[ERROR]: Invalid Url")
@@ -99,6 +124,10 @@ async def add_feed_func(_, m: Message):
     try:
         loop = get_event_loop()
         parsed = await loop.run_in_executor(None, parse, url)
+        final_url = get_parsed_feed_url(parsed, url)
+        if not is_safe_url(final_url):
+            return await m.reply("[ERROR]: URL is not allowed (SSRF protection).")
+
         feed = Feed(parsed)
     except Exception:
         return await m.reply(ns)
@@ -112,7 +141,7 @@ async def add_feed_func(_, m: Message):
         await m.reply(feed.parsed(), disable_web_page_preview=True)
     except Exception:
         return await m.reply(ns)
-    await add_rss_feed(chat_id, parsed.url, feed.title)
+    await add_rss_feed(chat_id, final_url, feed.title)
 
 
 @app.on_message(filters.command("rm_feed"))