From b4a55e2a535d0b153b9e8883a3b12f63778295a9 Mon Sep 17 00:00:00 2001
From: Dylan <dylanpdx@gmail.com>
Date: Thu, 8 Feb 2024 17:00:04 +0000
Subject: [PATCH] Updated image combine logic for checking hosts

---
 combineImg/__init__.py | 10 ++++++----
 twitfix.py             |  4 +++-
 2 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/combineImg/__init__.py b/combineImg/__init__.py
index ba35d22..27c8be7 100644
--- a/combineImg/__init__.py
+++ b/combineImg/__init__.py
@@ -4,6 +4,7 @@ from io import BytesIO
 import base64
 import concurrent.futures
 from time import time as timer
+from urllib.parse import urlparse 
 
 # find the highest res image in an array of images
 def findImageWithMostPixels(imageArray):
@@ -139,10 +140,11 @@ def lambda_handler(event, context):
             "body": "Invalid request."
         }
     images = event["queryStringParameters"].get("imgs","").split(",")
-    for img in images:
-        if not img.startswith("https://pbs.twimg.com"):
-            return {'statusCode':400,'body':'Invalid image URL'}
-    combined = genImageFromURL(images)
+    for img in imgs:
+        result = urlparse(img)
+        if result.hostname != "pbs.twimg.com" or result.scheme != "https":
+            abort(400)
+        combined = genImageFromURL(images)
     if (combined == None):
         return {'statusCode':200,'body':get500ImgBase64(),'isBase64Encoded':True,'headers':{"Content-Type": "image/jpeg","Cache-Control": "public, max-age=86400"}}
     buffered = BytesIO()
diff --git a/twitfix.py b/twitfix.py
index c3620f6..c678c29 100644
--- a/twitfix.py
+++ b/twitfix.py
@@ -17,6 +17,7 @@ from yt_dlp.utils import ExtractorError
 import vxlogging as log
 import zipfile
 import html
+from urllib.parse import urlparse 
 app = Flask(__name__)
 CORS(app)
 user_agent=""
@@ -360,7 +361,8 @@ def rendercombined():
         abort(400)
     #check that each image starts with "https://pbs.twimg.com"
     for img in imgs:
-        if not img.startswith("https://pbs.twimg.com"):
+        result = urlparse(img)
+        if result.hostname != "pbs.twimg.com" or result.scheme != "https":
             abort(400)
     finalImg= combineImg.genImageFromURL(imgs)
     imgIo = BytesIO()