TheHamkerCat vor 5 Jahren
Ursprung
Commit
9f22ecf499
3 geänderte Dateien mit 72 neuen und 30 gelöschten Zeilen
  1. 8 7
      wbb/__main__.py
  2. 7 1
      wbb/modules/arq_storage.py
  3. 57 22
      wbb/modules/reverse.py

+ 8 - 7
wbb/__main__.py

@@ -87,18 +87,19 @@ async def start_bot():
     print(f"[INFO]: BOT STARTED AS {BOT_NAME}!")
     print(f"[INFO]: USERBOT STARTED AS {USERBOT_NAME}!")
     restart_data = await clean_restart_stage()
-    if restart_data:
-        print("[INFO]: SENDING RESTART STATUS")
-        try:
+    try:
+        print("[INFO]: SENDING ONLINE STATUS")
+        if restart_data:
             await app.edit_message_text(
                 restart_data["chat_id"],
                 restart_data["message_id"],
                 "**Restarted Successfully**",
             )
-        except Exception:
-            pass
-    else:
-        await app.send_message(LOG_GROUP_ID, "Bot started!")
+
+        else:
+            await app.send_message(LOG_GROUP_ID, "Bot started!")
+    except Exception:
+        pass
     await idle()
     print("[INFO]: STOPPING BOT AND CLOSING AIOHTTP SESSION")
     await aiohttpsession.close()

+ 7 - 1
wbb/modules/arq_storage.py

@@ -29,10 +29,16 @@ async def upload(m: Message, file: str = None, url: str = None):
         if file:
             remove(file)
         return
+
+    if file:
+        remove(file)
+
     if not resp:
         return await m.edit(err)
+
     if not resp.ok:
         return await m.edit(err)
+
     await m.edit(
         f"**Download Link:** {resp.result}",
         disable_web_page_preview=True,
@@ -51,7 +57,7 @@ async def arq_upload(_, message):
 
             file_id = get_file_id_from_message(r, 100000000, None)
             if not file_id:
-                return await message.reply("Media not supported.")
+                return await message.reply("Unsupported media.")
 
             m = await message.reply("Downloading...")
             file = await app.download_media(file_id)

+ 57 - 22
wbb/modules/reverse.py

@@ -1,5 +1,7 @@
 import os
-from asyncio import get_running_loop
+from asyncio import gather, get_running_loop
+from base64 import b64decode
+from io import BytesIO
 from random import randint
 
 import aiofiles
@@ -7,15 +9,22 @@ import aiohttp
 import requests
 from bs4 import BeautifulSoup
 from pyrogram import filters
+from pyrogram.types import InputMediaPhoto, Message
 
-from wbb import app
+from wbb import MESSAGE_DUMP_CHAT, app
 from wbb.core.decorators.errors import capture_err
 from wbb.utils.functions import get_file_id_from_message
+from wbb.utils.http import get
+
+
+async def get_soup(url: str, headers):
+    html = await get(url, headers=headers)
+    return BeautifulSoup(html, "html.parser")
 
 
 @app.on_message(filters.command("reverse"))
 @capture_err
-async def reverse_image_search(_, message):
+async def reverse_image_search(_, message: Message):
     if not message.reply_to_message:
         return await message.reply_text(
             "Reply to a message to reverse search it."
@@ -31,14 +40,13 @@ async def reverse_image_search(_, message):
         return await message.reply_text(
             "Reply to an image/document/sticker/animation to reverse search it."
         )
-    m = await message.reply_text("Downloading")
+    m = await message.reply_text("Searching...")
     file_id = get_file_id_from_message(reply)
     if not file_id:
         return await m.edit("Can't reverse that")
     image = await app.download_media(
         file_id, f"{randint(1000, 10000)}.jpg"
     )
-    await m.edit("Uploading to google's server")
     async with aiofiles.open(image, "rb") as f:
         if image:
             search_url = "http://www.google.com/searchbyimage/upload"
@@ -63,21 +71,48 @@ async def reverse_image_search(_, message):
     headers = {
         "User-Agent": "Mozilla/5.0 (X11; Linux x86_64; rv:58.0) Gecko/20100101 Firefox/58.0"
     }
-    await m.edit("Scraping query")
-    async with aiohttp.ClientSession() as session:
-        async with session.get(location, headers=headers) as resp:
-            soup = BeautifulSoup(await resp.text(), "html.parser")
+
+    soup = await get_soup(location, headers=headers)
     div = soup.find_all("div", {"class": "r5a77d"})[0]
-    anchor_element = div.find("a")
-    text = anchor_element.text
-    try:
-        await m.edit("Trying to send a photo")
-        await app.send_photo(
-            message.chat.id,
-            photo=f"https://webshot.amanoteam.com/print?q={location}",
-            caption=f"**Result**: [{text}]({location})",
-        )
-        await m.delete()
-    except Exception:
-        text = f"**Result**: [Link]({location})"
-        await m.edit(text)
+    text = div.find("a").text
+
+    url = "https://google.com" + soup.find_all(
+        "a", {"class": "ekf0x hSQtef"}
+    )[0].get("href")
+
+    soup = await get_soup(url, headers=headers)
+
+    text = f"**Result**: [{text}]({location})"
+
+    media = []
+    for img in soup.find_all("img"):
+        if len(media) == 2:
+            break
+
+        if img.get("src"):
+            img = img.get("src")
+            if "image/gif" in img:
+                continue
+
+            img = BytesIO(b64decode(img))
+            img.name = "img.png"
+            media.append(img)
+        elif img.get("data-src"):
+            img = img.get("data-src")
+            media.append(img)
+
+    # Cache images, so we can use file_ids
+    tasks = [app.send_photo(MESSAGE_DUMP_CHAT, img) for img in media]
+    messages = await gather(*tasks)
+
+    await message.reply_media_group(
+        [
+            InputMediaPhoto(
+                i.photo.file_id,
+                caption=text,
+            )
+            for i in messages
+        ]
+    )
+
+    await m.edit(text)