From 3705d11c1953d833c548ecd1bdd94c395a58fb4d Mon Sep 17 00:00:00 2001
From: Myned <alphasea14@gmail.com>
Date: Mon, 20 Nov 2017 02:12:56 -0500
Subject: [PATCH] Added check for filesize before requesting from iqdb

---
 src/main/misc/exceptions.py | 3 +++
 src/main/utils/scraper.py   | 9 +++++++++
 src/main/utils/utils.py     | 6 ++++--
 3 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/src/main/misc/exceptions.py b/src/main/misc/exceptions.py
index dcc8c84..7959441 100644
--- a/src/main/misc/exceptions.py
+++ b/src/main/misc/exceptions.py
@@ -10,6 +10,9 @@ async def send_error(ctx, error):
 # class NSFW(errext.CheckFailure):
 #     pass
 
+class SizeError(Exception):
+    pass
+
 
 class Wrong(Exception):
     pass
diff --git a/src/main/utils/scraper.py b/src/main/utils/scraper.py
index 0482e86..ccb60fe 100644
--- a/src/main/utils/scraper.py
+++ b/src/main/utils/scraper.py
@@ -3,12 +3,21 @@ import re
 
 from bs4 import BeautifulSoup
 from lxml import html
+from hurry.filesize import size, alternative
 
 from misc import exceptions as exc
 from utils import utils as u
 
 
 async def get_post(url):
+    try:
+        image = await u.fetch(url, response=True)
+        filesize = int(image.headers['Content-Length'])
+        if filesize > 8192 * 1024:
+            raise exc.SizeError(size(filesize, system=alternative))
+    except ValueError:
+        raise exc.MissingArgument
+
     await asyncio.sleep(u.RATE_LIMIT)
 
     content = await u.fetch('http://iqdb.harry.lu', params={'url': url})
diff --git a/src/main/utils/utils.py b/src/main/utils/utils.py
index d8c7c91..c255aea 100644
--- a/src/main/utils/utils.py
+++ b/src/main/utils/utils.py
@@ -107,9 +107,11 @@ def close(loop):
     print('Finished cancelling tasks.')
 
 
-async def fetch(url, *, params={}, json=False):
+async def fetch(url, *, params={}, json=False, response=False):
     async with session.get(url, params=params, headers={'User-Agent': 'Myned/Modumind/dev'}) as r:
-        if json:
+        if response:
+            return r
+        elif json:
             return await r.json()
         return await r.read()