diff --git a/meme_otron/meme_otron.py b/meme_otron/meme_otron.py index ac2111f..17549a9 100644 --- a/meme_otron/meme_otron.py +++ b/meme_otron/meme_otron.py @@ -72,12 +72,10 @@ def compute_part(*args: str, input_data: Optional[bytes] = None, if input_data is None or len(input_data) == 0: if len(args) <= 1: return None, 'Image: received no input data nor URL' - elif not utils.validate_url(args[1]): - return None, 'Image: invalid URL format' else: - input_data = utils.read_web(args[1]) + input_data, err = utils.read_web(args[1]) if input_data is None: - return None, 'Image: could not reach URL' + return None, 'Image: ' + err img = img_factory.build_image_only(input_data) if img is None: return None, 'Image: invalid image format' diff --git a/meme_otron/utils.py b/meme_otron/utils.py index 3711b10..7153e50 100644 --- a/meme_otron/utils.py +++ b/meme_otron/utils.py @@ -235,12 +235,19 @@ def read_stream(stream: BinaryIO) -> bytes: return output_data -def read_web(url: str, timeout: int = 5) -> Optional[bytes]: +def read_web(url: str, *, timeout: int = 5, + max_file_size: Optional[int] = None) -> Tuple[Optional[bytes], Optional[str]]: + if not validate_url(url): + return None, 'Invalid URL' try: with urlopen(url, None, timeout) as web_file: - return web_file.read() + if web_file.getcode() != 200: + return None, 'Invalid return code' + if max_file_size is not None and int(web_file.info()['Content-Length']) > max_file_size: + return None, 'File too big' + return web_file.read(), None except URLError: - return None + return None, 'Could not connect' # endregion