Chore: fix interactive verification not work in container environment

BingImageGen.py

@@ -3,27 +3,26 @@ Code derived from:
 https://github.com/acheong08/EdgeGPT/blob/f940cecd24a4818015a8b42a2443dd97c3c2a8f4/src/ImageGen.py
 """
 from log import getlogger
+
+from typing import Union
 from uuid import uuid4
 import os
-import urllib
+import contextlib
+import aiohttp
+import asyncio
+import random
 import time
 import requests
 import regex
 
-BING_URL = "https://www.bing.com"
 logger = getlogger()
 
-
+BING_URL = "https://www.bing.com"
-class ImageGen:
+# Generate random IP between range 13.104.0.0/14
-    """
+FORWARDED_IP = (
-    Image generation by Microsoft Bing
+    f"13.{random.randint(104, 107)}.{random.randint(0, 255)}.{random.randint(0, 255)}"
-    Parameters:
+)
-        auth_cookie: str
+HEADERS = {
-    """
-
-    def __init__(self, auth_cookie: str) -> None:
-        self.session: requests.Session = requests.Session()
-        self.session.headers = {
     "accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7",
     "accept-language": "en-US,en;q=0.9",
     "cache-control": "max-age=0",
@@ -31,8 +30,50 @@ class ImageGen:
     "referrer": "https://www.bing.com/images/create/",
     "origin": "https://www.bing.com",
     "user-agent": "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/110.0.0.0 Safari/537.36 Edg/110.0.1587.63",
+    "x-forwarded-for": FORWARDED_IP,
 }
+
+# Error messages
+error_timeout = "Your request has timed out."
+error_redirect = "Redirect failed"
+error_blocked_prompt = (
+    "Your prompt has been blocked by Bing. Try to change any bad words and try again."
+)
+error_noresults = "Could not get results"
+error_unsupported_lang = "\nthis language is currently not supported by bing"
+error_bad_images = "Bad images"
+error_no_images = "No images"
+#
+sending_message = "Sending request..."
+wait_message = "Waiting for results..."
+download_message = "\nDownloading images..."
+
+
+
+def debug(debug_file, text_var):
+    """helper function for debug"""
+    with open(f"{debug_file}", "a") as f:
+        f.write(str(text_var))
+
+
+class ImageGen:
+    """
+    Image generation by Microsoft Bing
+    Parameters:3
+        auth_cookie: str
+    """
+
+    def __init__(
+        self, auth_cookie: str, debug_file: Union[str, None] = None, quiet: bool = False
+    ) -> None:
+        self.session: requests.Session = requests.Session()
+        self.session.headers = HEADERS
         self.session.cookies.set("_U", auth_cookie)
+        self.quiet = quiet
+        self.debug_file = debug_file
+        if self.debug_file:
+            self.debug = partial(debug, self.debug_file)
+
 
     def get_images(self, prompt: str) -> list:
         """
@@ -40,18 +81,37 @@ class ImageGen:
         Parameters:
             prompt: str
         """
-        print("Sending request...")
+        if not self.quiet:
-        url_encoded_prompt = urllib.parse.quote(prompt)
+            print(sending_message)
-        # https://www.bing.com/images/create?q=<PROMPT>&rt=4&FORM=GENCRE
+        if self.debug_file:
+            self.debug(sending_message)
+        url_encoded_prompt = requests.utils.quote(prompt)
+        # https://www.bing.com/images/create?q=<PROMPT>&rt=3&FORM=GENCRE
         url = f"{BING_URL}/images/create?q={url_encoded_prompt}&rt=4&FORM=GENCRE"
         response = self.session.post(url, allow_redirects=False)
+        # check for content waring message
+        if "this prompt has been blocked" in response.text.lower():
+            if self.debug_file:
+                self.debug(f"ERROR: {error_blocked_prompt}")
+            raise Exception(
+                error_blocked_prompt,
+            )
+        if (
+            "we're working hard to offer image creator in more languages"
+            in response.text.lower()
+        ):
+            if self.debug_file:
+                self.debug(f"ERROR: {error_unsupported_lang}")
+            raise Exception(error_unsupported_lang)
         if response.status_code != 302:
             # if rt4 fails, try rt3
             url = f"{BING_URL}/images/create?q={url_encoded_prompt}&rt=3&FORM=GENCRE"
-            response3 = self.session.post(url, allow_redirects=False)
+            response3 = self.session.post(url, allow_redirects=False, timeout=200)
             if response3.status_code != 302:
-                logger.error(f"ERROR: {response.text}")
+                if self.debug_file:
-                return []
+                    self.debug(f"ERROR: {error_redirect}")
+                print(f"ERROR: {response3.text}")
+                raise Exception(error_redirect)
             response = response3
         # Get redirect URL
         redirect_url = response.headers["Location"].replace("&nfy=1", "")
@@ -60,50 +120,196 @@ class ImageGen:
         # https://www.bing.com/images/create/async/results/{ID}?q={PROMPT}
         polling_url = f"{BING_URL}/images/create/async/results/{request_id}?q={url_encoded_prompt}"
         # Poll for results
+        if self.debug_file:
+            self.debug("Polling and waiting for result")
+        if not self.quiet:
             print("Waiting for results...")
+        start_wait = time.time()
         while True:
+            if int(time.time() - start_wait) > 200:
+                if self.debug_file:
+                    self.debug(f"ERROR: {error_timeout}")
+                raise Exception(error_timeout)
+            if not self.quiet:
                 print(".", end="", flush=True)
             response = self.session.get(polling_url)
             if response.status_code != 200:
-                logger.error("Could not get results", exc_info=True)
+                if self.debug_file:
-                return []
+                    self.debug(f"ERROR: {error_noresults}")
-            if response.text == "":
+                raise Exception(error_noresults)
+            if not response.text or response.text.find("errorMessage") != -1:
                 time.sleep(1)
                 continue
             else:
                 break
-
         # Use regex to search for src=""
         image_links = regex.findall(r'src="([^"]+)"', response.text)
         # Remove size limit
         normal_image_links = [link.split("?w=")[0] for link in image_links]
         # Remove duplicates
-        return list(set(normal_image_links))
+        normal_image_links = list(set(normal_image_links))
+
+        # bad_images = [
+        #     "https://r.bing.com/rp/in-2zU3AJUdkgFe7ZKv19yPBHVs.png",
+        #     "https://r.bing.com/rp/TX9QuO3WzcCJz1uaaSwQAz39Kb0.jpg",
+        # ]
+        # for img in normal_image_links:
+        #     if img in bad_images:
+        #         raise Exception("Bad images")
+        # No images
+        if not normal_image_links:
+            raise Exception(error_no_images)
+        return normal_image_links
 
     def save_images(self, links: list, output_dir: str) -> str:
         """
         Saves images to output directory
         """
-        print("\nDownloading images...")
+        
-        try:
-            os.mkdir(output_dir)
-        except FileExistsError:
-            pass
         # image name
         image_name = str(uuid4())
         # since matrix only support one media attachment per message, we just need one link
         if links:
             link = links.pop()
-        else:
-            logger.error("Get Image URL failed")
-            # return "" if there is no link
-            return ""
 
+        image_path = os.path.join(output_dir, f"{image_name}.jpeg")
+
+        with contextlib.suppress(FileExistsError):
+            os.mkdir(output_dir)
+        try:
             with self.session.get(link, stream=True) as response:
                 # save response to file
                 response.raise_for_status()
-            with open(f"{output_dir}/{image_name}.jpeg", "wb") as output_file:
+                with open(
+                    os.path.join(output_dir, image_path), "wb"
+                ) as output_file:
                     for chunk in response.iter_content(chunk_size=8192):
                         output_file.write(chunk)
+            return image_path
+        except requests.exceptions.MissingSchema as url_exception:
+            raise Exception(
+                "Inappropriate contents found in the generated images. Please try again or try another prompt.",
+            ) from url_exception
 
+
+class ImageGenAsync:
+    """
+    Image generation by Microsoft Bing
+    Parameters:
+        auth_cookie: str
+    """
+
+    def __init__(self, auth_cookie: str, quiet: bool = True) -> None:
+        self.session = aiohttp.ClientSession(
+            headers=HEADERS,
+            cookies={"_U": auth_cookie},
+        )
+        self.quiet = quiet
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, *excinfo) -> None:
+        await self.session.close()
+
+    async def get_images(self, prompt: str) -> list:
+        """
+        Fetches image links from Bing
+        Parameters:
+            prompt: str
+        """
+        if not self.quiet:
+            print("Sending request...")
+        url_encoded_prompt = requests.utils.quote(prompt)
+        # https://www.bing.com/images/create?q=<PROMPT>&rt=3&FORM=GENCRE
+        url = f"{BING_URL}/images/create?q={url_encoded_prompt}&rt=4&FORM=GENCRE"
+        async with self.session.post(url, allow_redirects=False) as response:
+            content = await response.text()
+            if "this prompt has been blocked" in content.lower():
+                raise Exception(
+                    "Your prompt has been blocked by Bing. Try to change any bad words and try again.",
+                )
+            if response.status != 302:
+                # if rt4 fails, try rt3
+                url = (
+                    f"{BING_URL}/images/create?q={url_encoded_prompt}&rt=3&FORM=GENCRE"
+                )
+                async with self.session.post(
+                    url,
+                    allow_redirects=False,
+                    timeout=200,
+                ) as response3:
+                    if response3.status != 302:
+                        print(f"ERROR: {response3.text}")
+                        raise Exception("Redirect failed")
+                    response = response3
+        # Get redirect URL
+        redirect_url = response.headers["Location"].replace("&nfy=1", "")
+        request_id = redirect_url.split("id=")[-1]
+        await self.session.get(f"{BING_URL}{redirect_url}")
+        # https://www.bing.com/images/create/async/results/{ID}?q={PROMPT}
+        polling_url = f"{BING_URL}/images/create/async/results/{request_id}?q={url_encoded_prompt}"
+        # Poll for results
+        if not self.quiet:
+            print("Waiting for results...")
+        while True:
+            if not self.quiet:
+                print(".", end="", flush=True)
+            # By default, timeout is 300s, change as needed
+            response = await self.session.get(polling_url)
+            if response.status != 200:
+                raise Exception("Could not get results")
+            content = await response.text()
+            if content and content.find("errorMessage") == -1:
+                break
+
+            await asyncio.sleep(1)
+            continue
+        # Use regex to search for src=""
+        image_links = regex.findall(r'src="([^"]+)"', content)
+        # Remove size limit
+        normal_image_links = [link.split("?w=")[0] for link in image_links]
+        # Remove duplicates
+        normal_image_links = list(set(normal_image_links))
+
+        # Bad images
+        bad_images = [
+            "https://r.bing.com/rp/in-2zU3AJUdkgFe7ZKv19yPBHVs.png",
+            "https://r.bing.com/rp/TX9QuO3WzcCJz1uaaSwQAz39Kb0.jpg",
+        ]
+        for im in normal_image_links:
+            if im in bad_images:
+                raise Exception("Bad images")
+        # No images
+        if not normal_image_links:
+            raise Exception("No images")
+        return normal_image_links
+
+    async def save_images(self, links: list, output_dir: str) -> str:
+        """
+        Saves images to output directory
+        """
+        if not self.quiet:
+            print("\nDownloading images...")
+        with contextlib.suppress(FileExistsError):
+            os.mkdir(output_dir)
+
+        # image name
+        image_name = str(uuid4())
+        # since matrix only support one media attachment per message, we just need one link
+        if links:
+            link = links.pop()
+
+        image_path = os.path.join(output_dir, f"{image_name}.jpeg")
+        try:
+            async with self.session.get(link, raise_for_status=True) as response:
+                # save response to file
+                with open(image_path, "wb") as output_file:
+                    async for chunk in response.content.iter_chunked(8192):
+                        output_file.write(chunk)
             return f"{output_dir}/{image_name}.jpeg"
+            
+        except aiohttp.client_exceptions.InvalidURL as url_exception:
+            raise Exception(
+                "Inappropriate contents found in the generated images. Please try again or try another prompt.",
+            ) from url_exception

bot.py

@@ -2,6 +2,7 @@ import sys
 import asyncio
 import re
 import os
+from functools import partial
 import traceback
 from typing import Optional, Union
 from nio import (
@@ -27,7 +28,7 @@ from send_message import send_room_message
 from v3 import Chatbot
 from log import getlogger
 from bing import BingBot
-from BingImageGen import ImageGen
+from BingImageGen import ImageGenAsync
 from send_image import send_room_image
 
 logger = getlogger()
@@ -39,7 +40,8 @@ class Bot:
         homeserver: str,
         user_id: str,
         device_id: str,
-        chatgpt_api_endpoint: str = os.environ.get("CHATGPT_API_ENDPOINT") or "https://api.openai.com/v1/chat/completions",
+        chatgpt_api_endpoint: str = os.environ.get(
+            "CHATGPT_API_ENDPOINT") or "https://api.openai.com/v1/chat/completions",
         api_key: Optional[str] = os.environ.get("OPENAI_API_KEY") or "",
         room_id: Union[str, None] = None,
         bing_api_endpoint: Optional[str] = '',
@@ -48,7 +50,7 @@ class Bot:
         jailbreakEnabled: Optional[bool] = True,
         bing_auth_cookie: Optional[str] = '',
     ):
-        if (homeserver is None or user_id is None \
+        if (homeserver is None or user_id is None
                 or device_id is None):
             logger.warning("homeserver && user_id && device_id is required")
             sys.exit(1)
@@ -82,10 +84,14 @@ class Bot:
             self.client.access_token = self.access_token
 
         # setup event callbacks
-        self.client.add_event_callback(self.message_callback, (RoomMessageText, ))
+        self.client.add_event_callback(
-        self.client.add_event_callback(self.decryption_failure, (MegolmEvent, ))
+            self.message_callback, (RoomMessageText, ))
-        self.client.add_event_callback(self.invite_callback, (InviteMemberEvent, ))
+        self.client.add_event_callback(
-        self.client.add_to_device_callback(self.to_device_callback, (KeyVerificationEvent, ))
+            self.decryption_failure, (MegolmEvent, ))
+        self.client.add_event_callback(
+            self.invite_callback, (InviteMemberEvent, ))
+        self.client.add_to_device_callback(
+            self.to_device_callback, (KeyVerificationEvent, ))
 
         # regular expression to match keyword [!gpt {prompt}] [!chat {prompt}]
         self.gpt_prog = re.compile(r"^\s*!gpt\s*(.+)$")
@@ -115,11 +121,15 @@ class Bot:
 
         # initialize bingbot
         if self.bing_api_endpoint != '':
-            self.bingbot = BingBot(bing_api_endpoint, jailbreakEnabled=self.jailbreakEnabled)
+            self.bingbot = BingBot(
+                bing_api_endpoint, jailbreakEnabled=self.jailbreakEnabled)
 
-        # initialize BingImageGen
+        # initialize BingImageGenAsync
         if self.bing_auth_cookie != '':
-            self.imageGen = ImageGen(self.bing_auth_cookie)
+            self.imageGen = ImageGenAsync(self.bing_auth_cookie, quiet=True)
+
+        # get current event loop
+        self.loop = asyncio.get_running_loop()
 
     # message_callback RoomMessageText event
     async def message_callback(self, room: MatrixRoom, event: RoomMessageText) -> None:
@@ -157,7 +167,13 @@ class Bot:
                 prompt = n.group(1)
                 if self.api_key != '':
                     try:
-                        await self.chat(room_id, reply_to_event_id, prompt, sender_id, raw_user_message)
+                        await self.chat(room_id,
+                                        reply_to_event_id,
+                                        prompt,
+                                        sender_id,
+                                        raw_user_message
+                                        )
+
                     except Exception as e:
                         logger.error(e)
                         await send_room_message(self.client, room_id, reply_message=str(e))
@@ -169,7 +185,12 @@ class Bot:
             if m:
                 prompt = m.group(1)
                 try:
-                    await self.gpt(room_id, reply_to_event_id, prompt, sender_id, raw_user_message)
+                    await self.gpt(
+                        room_id,
+                        reply_to_event_id,
+                        prompt, sender_id,
+                        raw_user_message
+                    )
                 except Exception as e:
                     logger.error(e)
                     await send_room_message(self.client, room_id, reply_message=str(e))
@@ -181,7 +202,14 @@ class Bot:
                     prompt = b.group(1)
                     # raw_content_body used for construct formatted_body
                     try:
-                        await self.bing(room_id, reply_to_event_id, prompt, sender_id, raw_user_message)
+                        await self.bing(
+                            room_id,
+                            reply_to_event_id,
+                            prompt,
+                            sender_id,
+                            raw_user_message
+                        )
+
                     except Exception as e:
                         await send_room_message(self.client, room_id, reply_message=str(e))
 
@@ -206,7 +234,7 @@ class Bot:
             return
 
         logger.error(
-            f"Failed to decrypt message: {event.event_id} from {event.sender} in {room.room_id}\n" + \
+            f"Failed to decrypt message: {event.event_id} from {event.sender} in {room.room_id}\n" +
             "Please make sure the bot current session is verified"
         )
 
@@ -346,7 +374,9 @@ class Bot:
                 # keyboard so that user can accept/reject via keyboard.
                 # For emoji verification bot must not run as service or
                 # in background.
-                yn = input("Do the emojis match? (Y/N) (C for Cancel) ")
+                # yn = input("Do the emojis match? (Y/N) (C for Cancel) ")
+                # automatic match, so we use y
+                yn = "y"
                 if yn.lower() == "y":
                     estr = ("Match! The verification for this "
                             "device will be accepted.")
@@ -478,7 +508,7 @@ class Bot:
             logger.error(f"Error: {e}", exc_info=True)
 
     # !bing command
-    async def bing(self, room_id, reply_to_event_id, prompt, sender_id, raw_content_body):
+    async def bing(self, room_id, reply_to_event_id, prompt, sender_id, raw_user_message):
         try:
             # sending typing state
             await self.client.room_typing(room_id, timeout=180000)
@@ -490,7 +520,7 @@ class Bot:
         text = text.strip()
         try:
             await send_room_message(self.client, room_id, reply_message=text,
-                                    reply_to_event_id=reply_to_event_id, sender=sender_id, raw_content_body=raw_content_body)
+                                    reply_to_event_id=reply_to_event_id, sender_id=sender_id, user_message=raw_user_message)
         except Exception as e:
             logger.error(f"Error: {e}", exc_info=True)
 
@@ -499,14 +529,20 @@ class Bot:
         try:
             await self.client.room_typing(room_id, timeout=180000)
             # generate image
-            generated_image_path = self.imageGen.save_images(
+            try:
-                self.imageGen.get_images(prompt),
+
-                "images",
+                links = await self.imageGen.get_images(prompt)
-            )
+                image_path = await self.imageGen.save_images(links, "images")
+            except Exception as e:
+                logger.error(f"Image Generation error: {e}", exc_info=True)
+
             # send image
-            if generated_image_path != "":
+            try:
-                await send_room_image(self.client, room_id, generated_image_path)
+                await send_room_image(self.client, room_id, image_path)
-                await self.client.room_typing(room_id, bool=False)
+                await self.client.room_typing(room_id, typing_state=False)
+            except Exception as e:
+                logger.error(e, exc_info=True)
+
         except Exception as e:
             logger.error(f"Error: {e}", exc_info=True)
 
@@ -562,4 +598,3 @@ class Bot:
                         f"device {device_id} and session key "
                         f"{olm_device.keys['ed25519']}.")
             self.client.verify_device(olm_device)
-