1
0
Fork 0
mirror of https://github.com/myned/modufur.git synced 2024-12-25 14:47:29 +00:00

Added check for filesize before requesting from iqdb

This commit is contained in:
Myned 2017-11-20 02:12:56 -05:00
parent c129f00ff3
commit 3705d11c19
3 changed files with 16 additions and 2 deletions

View file

@ -10,6 +10,9 @@ async def send_error(ctx, error):
# class NSFW(errext.CheckFailure): # class NSFW(errext.CheckFailure):
# pass # pass
class SizeError(Exception):
pass
class Wrong(Exception): class Wrong(Exception):
pass pass

View file

@ -3,12 +3,21 @@ import re
from bs4 import BeautifulSoup from bs4 import BeautifulSoup
from lxml import html from lxml import html
from hurry.filesize import size, alternative
from misc import exceptions as exc from misc import exceptions as exc
from utils import utils as u from utils import utils as u
async def get_post(url): 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) await asyncio.sleep(u.RATE_LIMIT)
content = await u.fetch('http://iqdb.harry.lu', params={'url': url}) content = await u.fetch('http://iqdb.harry.lu', params={'url': url})

View file

@ -107,9 +107,11 @@ def close(loop):
print('Finished cancelling tasks.') 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: 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.json()
return await r.read() return await r.read()