From 0783e7b459be4576e6165f9f229c9aac3fbccb2d Mon Sep 17 00:00:00 2001
From: Audionut <audionut11@gmail.com>
Date: Tue, 17 Dec 2024 21:45:11 +1000
Subject: [PATCH] Clean MTV host handling

---
 src/prep.py         |  2 +-
 src/trackers/MTV.py | 41 +++++++++++++++++++++++++++++++++++++----
 2 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/src/prep.py b/src/prep.py
index 437deb0f..a13d6e3b 100644
--- a/src/prep.py
+++ b/src/prep.py
@@ -88,7 +88,7 @@ async def prompt_user_for_confirmation(self, message: str) -> bool:
             sys.exit(1)
 
     async def check_images_concurrently(self, imagelist, meta):
-        approved_image_hosts = ['ptpimg', 'imgbox']
+        approved_image_hosts = ['ptpimg', 'imgbox', 'imgbb']
         invalid_host_found = False  # Track if any image is on a non-approved host
 
         # Ensure meta['image_sizes'] exists
diff --git a/src/trackers/MTV.py b/src/trackers/MTV.py
index 92b8011b..9bfa86eb 100644
--- a/src/trackers/MTV.py
+++ b/src/trackers/MTV.py
@@ -14,6 +14,7 @@
 from datetime import datetime
 import glob
 import multiprocessing
+from urllib.parse import urlparse
 
 
 class MTV():
@@ -46,9 +47,29 @@ async def upload(self, meta, disctype):
 
     async def upload_with_retry(self, meta, cookiefile, common, img_host_index=1):
         approved_image_hosts = ['ptpimg', 'imgbox', 'imgbb']
+        url_host_mapping = {
+            "i.ibb.co": "imgbb",
+            "ptpimg.me": "ptpimg",
+            "images2.imgbox.com": "imgbox",
+        }
         images_reuploaded = False
 
-        if all(any(host in image['raw_url'] for host in approved_image_hosts) for image in meta['image_list']):
+        normalized_approved_hosts = set(approved_image_hosts + list(url_host_mapping.keys()))  # noqa F841
+        for image in meta['image_list']:
+            raw_url = image['raw_url']
+            parsed_url = urlparse(raw_url)
+            hostname = parsed_url.netloc
+            mapped_host = url_host_mapping.get(hostname, hostname)
+            if meta['debug']:
+                if mapped_host in approved_image_hosts:
+                    console.print(f"[green]URL '{raw_url}' is correctly matched to approved host '{mapped_host}'.")
+                else:
+                    console.print(f"[red]URL '{raw_url}' is not recognized as part of an approved host.")
+
+        if all(
+            url_host_mapping.get(urlparse(image['raw_url']).netloc, urlparse(image['raw_url']).netloc) in approved_image_hosts
+            for image in meta['image_list']
+        ):
             console.print("[green]Images are already hosted on an approved image host. Skipping re-upload.")
             image_list = meta['image_list']
 
@@ -181,6 +202,12 @@ async def handle_image_upload(self, meta, img_host_index=1, approved_image_hosts
         if approved_image_hosts is None:
             approved_image_hosts = ['ptpimg', 'imgbox', 'imgbb']
 
+        url_host_mapping = {
+            "i.ibb.co": "imgbb",
+            "ptpimg.me": "ptpimg",
+            "images2.imgbox.com": "imgbox",
+        }
+
         retry_mode = False
         images_reuploaded = False
         new_images_key = 'mtv_images_key'
@@ -280,9 +307,15 @@ async def handle_image_upload(self, meta, img_host_index=1, approved_image_hosts
             for image in uploaded_images:
                 console.print(f"[debug] Response in upload_image_task: {image['img_url']}, {image['raw_url']}, {image['web_url']}")
 
-        if not all(any(x in image['raw_url'] for x in approved_image_hosts) for image in meta.get(new_images_key, [])):
-            console.print("[red]Unsupported image host detected, please use one of the approved image hosts")
-            return meta[new_images_key], True, images_reuploaded  # Trigger retry_mode if switching hosts
+        for image in meta.get(new_images_key, []):
+            raw_url = image['raw_url']
+            parsed_url = urlparse(raw_url)
+            hostname = parsed_url.netloc
+            mapped_host = url_host_mapping.get(hostname, hostname)
+
+            if mapped_host not in approved_image_hosts:
+                console.print(f"[red]Unsupported image host detected in URL '{raw_url}'. Please use one of the approved image hosts.")
+                return meta[new_images_key], True, images_reuploaded  # Trigger retry_mode if switching hosts
 
         return meta[new_images_key], False, images_reuploaded