From 4f2a445302056a4de201927060a1da7bfbd13f11 Mon Sep 17 00:00:00 2001
From: Anshul Gupta <ansg191@anshulg.com>
Date: Sat, 16 Nov 2024 19:26:01 -0800
Subject: Add extension check to images

---
 core/image.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

(limited to 'core/image.py')

diff --git a/core/image.py b/core/image.py
index 7e3c68d..e699d79 100644
--- a/core/image.py
+++ b/core/image.py
@@ -19,5 +19,7 @@ def contains_image(html: str) -> Optional[str]:
     soup = BeautifulSoup(html, "html.parser")
     image = soup.find("img", recursive=True)
     if image:
-        return image.get("src")
+        src = image.get("src")
+        if src and src.lower().split('.')[-1] in ['png', 'jpeg', 'gif', 'webp']:
+            return src
     return None
-- 
cgit v1.2.3