Added support for avatar url validity checking, header filesize checking, header content type checking

This commit is contained in:
Hannah 2019-05-14 18:02:11 +00:00
parent 2f8bd44687
commit 6063ce97e9

19
src/pluralkit/utils.py Normal file → Executable file
View file

@ -6,9 +6,13 @@ import string
from datetime import datetime, timezone, timedelta
from typing import List, Tuple, Union
from urllib.parse import urlparse
from urllib.request import urlopen
from urllib.error import HTTPError
import requests
from pluralkit import db
from pluralkit.errors import InvalidAvatarURLError
from pluralkit.errors import InvalidAvatarURLError, AvatarHTTPError, InvalidAvatarContentTypeError, AvatarFileSizeTooLargeError
def display_relative(time: Union[datetime, timedelta]) -> str:
@ -69,5 +73,18 @@ def validate_avatar_url_or_raise(url):
u = urlparse(url)
if not (u.scheme in ["http", "https"] and u.netloc and u.path):
raise InvalidAvatarURLError()
response = ''
response = requests.head(url) # Requests won't output a ton of garbage to console when there's a 404, just one line.
if (response.status_code() != 200):
raise AvatarHTTPError(error)
u = urlopen(url) # get header info
u.close() # we don't need to read the file
ContentType = u.info()['content-type']
ContentType = str.lower(ContentType) # HTTP header feilds are case insensitive so we may get capital letters from sillier web servers
ContentLength = int(u.info()['content-length'])
if (ContentType != 'image/jpeg') and (ContentType != 'image/png') and (ContentType != 'image/gif'): # check for valid avatar filetype
raise InvalidAvatarContentTypeError()
elif (ContentLength > 1000000):
raise AvatarFileSizeTooLargeError()
# TODO: check file type and size of image