From 34f2a6ab9541d52b23a1942417faaa47571dccd0 Mon Sep 17 00:00:00 2001
From: Gordon Shumway <39967334+oddluck@users.noreply.github.com>
Date: Tue, 26 Mar 2019 12:26:20 -0400
Subject: [PATCH] get header and verify image type

---
 ASCII/plugin.py | 16 ++++++++++++++--
 1 file changed, 14 insertions(+), 2 deletions(-)

diff --git a/ASCII/plugin.py b/ASCII/plugin.py
index bcbbc0d..dc33140 100644
--- a/ASCII/plugin.py
+++ b/ASCII/plugin.py
@@ -230,7 +230,13 @@ class ASCII(callbacks.Plugin):
         filename = "{0}/{1}".format(filepath, url.split('/')[-1])
         ua = UserAgent()
         header = {'User-Agent':str(ua.random)}
-        response = requests.get(url, headers=header)
+        image_formats = ("image/png", "image/jpeg", "image/jpg", "image/gif")
+        r = requests.head(url, headers=header)
+        if r.headers["content-type"] in image_formats:
+            response = requests.get(url, headers=header)
+        else:
+            irc.reply("Invalid file type.")
+            return
         if response.status_code == 200:
             with open("{0}".format(filename), 'wb') as f:
                 f.write(response.content)
@@ -322,7 +328,13 @@ class ASCII(callbacks.Plugin):
         filename = "{0}/{1}".format(filepath, url.split('/')[-1])
         ua = UserAgent()
         header = {'User-Agent':str(ua.random)}
-        response = requests.get(url, headers=header)
+        image_formats = ("image/png", "image/jpeg", "image/jpg", "image/gif")
+        r = requests.head(url, headers=header)
+        if r.headers["content-type"] in image_formats:
+            response = requests.get(url, headers=header)
+        else:
+            irc.reply("Invalid file type.")
+            return
         if response.status_code == 200:
             with open("{0}".format(filename), 'wb') as f:
                 f.write(response.content)