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()