From 72da699ac7ff02ac526ad7af8e49e2a0dc7e1d69 Mon Sep 17 00:00:00 2001
From: Norielle Cruz <nogats07@gmail.com>
Date: Sat, 5 Aug 2023 06:41:18 +0800
Subject: [PATCH 1/3] fix: screenshot not showing properly if the website has
 slow loading time

---
 bot/plugins/webshot.py | 39 ++++++++++++++++++++++++++++++---------
 1 file changed, 30 insertions(+), 9 deletions(-)

diff --git a/bot/plugins/webshot.py b/bot/plugins/webshot.py
index 5d489205..45e93ad9 100644
--- a/bot/plugins/webshot.py
+++ b/bot/plugins/webshot.py
@@ -1,3 +1,4 @@
+import os, requests, random, string
 from typing import Dict
 from .plugin import Plugin
 
@@ -11,21 +12,41 @@ def get_source_name(self) -> str:
     def get_spec(self) -> [Dict]:
         return [{
             "name": "screenshot_website",
-            "description": "Show screenshot/image of a website from a given url or domain name",
+            "description": "Show screenshot/image of a website from a given url or domain name.",
             "parameters": {
                 "type": "object",
                 "properties": {
-                    "url": {"type": "string", "description": "Website url or domain name"}
+                    "url": {"type": "string", "description": "Website url or domain name. Correctly formatted url is required. Example: https://www.google.com"}
                 },
                 "required": ["url"],
             },
         }]
+    
+    def generate_random_string(self, length):
+        characters = string.ascii_letters + string.digits
+        return ''.join(random.choice(characters) for _ in range(length))
 
     async def execute(self, function_name, **kwargs) -> Dict:
-        return {
-            'direct_result': {
-                'kind': 'photo',
-                'format': 'url',
-                'value': f'https://image.thum.io/get/maxAge/12/width/720/{kwargs["url"]}'
-            }
-        }
+        try:
+            image_url = f'https://image.thum.io/get/maxAge/12/width/720/{kwargs["url"]}'
+            response = requests.get(image_url, timeout=30)
+
+            if response.status_code == 200:
+                if not os.path.exists("uploads/webshot"):
+                    os.makedirs("uploads/webshot")
+
+                image_file_path = os.path.join("uploads/webshot", f"{self.generate_random_string(15)}.png")
+                with open(image_file_path, "wb") as f:
+                    f.write(response.content)
+
+                return {
+                    'direct_result': {
+                        'kind': 'photo',
+                        'format': 'path',
+                        'value': image_file_path
+                    }
+                }
+            else:
+                return {'result': 'Unable to screenshot website'}
+        except:
+            return {'result': 'Unable to screenshot website'}

From 9165039abfe4899ae9b42f5f4e89db1f0b9ab599 Mon Sep 17 00:00:00 2001
From: Norielle Cruz <nogats07@gmail.com>
Date: Sat, 5 Aug 2023 06:57:31 +0800
Subject: [PATCH 2/3] refactor: add url preloading before actual download

---
 bot/plugins/webshot.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/bot/plugins/webshot.py b/bot/plugins/webshot.py
index 45e93ad9..4684c3da 100644
--- a/bot/plugins/webshot.py
+++ b/bot/plugins/webshot.py
@@ -29,6 +29,11 @@ def generate_random_string(self, length):
     async def execute(self, function_name, **kwargs) -> Dict:
         try:
             image_url = f'https://image.thum.io/get/maxAge/12/width/720/{kwargs["url"]}'
+            
+            # preload url first
+            requests.get(image_url)
+
+            # download the actual image
             response = requests.get(image_url, timeout=30)
 
             if response.status_code == 200:

From 890341666eb4719d6ee8c2d1120eaeaefea56a20 Mon Sep 17 00:00:00 2001
From: Norielle Cruz <nogats07@gmail.com>
Date: Sat, 5 Aug 2023 07:05:09 +0800
Subject: [PATCH 3/3] refactor: delete file if an error occurs

---
 bot/plugins/webshot.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/bot/plugins/webshot.py b/bot/plugins/webshot.py
index 4684c3da..4eb0e03b 100644
--- a/bot/plugins/webshot.py
+++ b/bot/plugins/webshot.py
@@ -54,4 +54,7 @@ async def execute(self, function_name, **kwargs) -> Dict:
             else:
                 return {'result': 'Unable to screenshot website'}
         except:
+            if 'image_file_path' in locals():
+                os.remove(image_file_path)
+                
             return {'result': 'Unable to screenshot website'}