Wiizard/wiizard/requestHandler.py

########################################################################
#  Wiizard - A Wii games manager
#  Copyright (C) 2023  CYBERDEViL
#
#  This file is part of Wiizard.
#
#  Wiizard is free software: you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software Foundation, either version 3 of the License, or
#  (at your option) any later version.
#
#  Wiizard is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
#  You should have received a copy of the GNU General Public License
#  along with this program.  If not, see <https://www.gnu.org/licenses/>.
#
########################################################################


import requests
from PIL import Image, UnidentifiedImageError
import io


class ImageDownloadException(Exception):
  pass


class RequestHandler:
  USERAGENT    = "Mozilla/5.0 (X11; Linux x86_64; rv:109.0) Gecko/20100101 iceweasel/109.0"
  MAX_IMG_SIZE = 40000000  # 4MB

  def __init__(self):
    self.__headers = {"User-Agent": self.USERAGENT}

  def get(self, url):
    """ This may raise requests.RequestException
    """
    headers  = self.__headers
    response = requests.get(url, headers=headers)
    return response

  def head(self, url):
    """ This may raise requests.RequestException
    """
    headers  = self.__headers
    response = requests.head(url, headers=headers)
    return response

  def downloadImage(self, url, path):
    """ This may raise a ImageDownloadException
    """
    # First do a HEAD request
    try:
      response = self.head(url)
    except requests.RequestException as err:
      raise ImageDownloadException(err)

    headers  = response.headers

    contentType = headers.get("Content-Type", "")
    if contentType != "image/png":
      raise ImageDownloadException("No a img", url)

    contentLength = headers.get("Content-Length", "")
    if not contentLength.isnumeric():
      raise ImageDownloadException("No valid content length set " + url)

    contentLength = int(contentLength)
    if contentLength > __class__.MAX_IMG_SIZE:
      raise ImageDownloadException("What kind of img is this? It's to large.. " + url)

    # Get the image data
    response = requests.get(url, stream=True)

    if response.status_code != 200:
      raise ImageDownloadException("Wrong status for fetching img " + url)

    response.raw.decode_content = True
    outputSize                  = 0
    data                        = b''
    for chunk in response:
      outputSize += len(chunk)
      if outputSize > __class__.MAX_IMG_SIZE:
        raise ImageDownloadException("Img is to large! HEAD was lying! " + url)
      data += chunk

    # Verify image data
    # https://pillow.readthedocs.io/en/stable/PIL.html
    try:
      img = Image.open(io.BytesIO(data))
      # https://pillow.readthedocs.io/en/stable/reference/Image.html#PIL.Image.Image.verify
      # TODO what "suitable exceptions" are referenced in the docs?
      img.verify()
    except UnidentifiedImageError:
      raise ImageDownloadException("Invalid image data for " + url)

    # Write image data
    try:
      with open(path, "wb") as f:
        f.write(data)
    except OSError as err:
      raise ImageDownloadException(err)