microblog.pub/utils/media.py

import base64
import mimetypes
from enum import Enum
from gzip import GzipFile
from io import BytesIO
from typing import Any

import gridfs
import piexif
import requests
from PIL import Image


def load(url, user_agent):
    """Initializes a `PIL.Image` from the URL."""
    with requests.get(url, stream=True, headers={"User-Agent": user_agent}) as resp:
        resp.raise_for_status()
        resp.raw.decode_content = True
        return Image.open(BytesIO(resp.raw.read()))


def to_data_uri(img):
    out = BytesIO()
    img.save(out, format=img.format)
    out.seek(0)
    data = base64.b64encode(out.read()).decode("utf-8")
    return f"data:{img.get_format_mimetype()};base64,{data}"


class Kind(Enum):
    ATTACHMENT = "attachment"
    ACTOR_ICON = "actor_icon"
    UPLOAD = "upload"


class MediaCache(object):
    def __init__(self, gridfs_db: str, user_agent: str) -> None:
        self.fs = gridfs.GridFS(gridfs_db)
        self.user_agent = user_agent

    def cache_attachment(self, url: str) -> None:
        if self.fs.find_one({"url": url, "kind": Kind.ATTACHMENT.value}):
            return
        if (
            url.endswith(".png")
            or url.endswith(".jpg")
            or url.endswith(".jpeg")
            or url.endswith(".gif")
        ):
            i = load(url, self.user_agent)
            # Save the original attachment (gzipped)
            with BytesIO() as buf:
                f1 = GzipFile(mode="wb", fileobj=buf)
                i.save(f1, format=i.format)
                f1.close()
                buf.seek(0)
                self.fs.put(
                    buf,
                    url=url,
                    size=None,
                    content_type=i.get_format_mimetype(),
                    kind=Kind.ATTACHMENT.value,
                )
            # Save a thumbnail (gzipped)
            i.thumbnail((720, 720))
            with BytesIO() as buf:
                with GzipFile(mode="wb", fileobj=buf) as f1:
                    i.save(f1, format=i.format)
                buf.seek(0)
                self.fs.put(
                    buf,
                    url=url,
                    size=720,
                    content_type=i.get_format_mimetype(),
                    kind=Kind.ATTACHMENT.value,
                )
            return

        # The attachment is not an image, download and save it anyway
        with requests.get(
            url, stream=True, headers={"User-Agent": self.user_agent}
        ) as resp:
            resp.raise_for_status()
            with BytesIO() as buf:
                with GzipFile(mode="wb", fileobj=buf) as f1:
                    for chunk in resp.iter_content():
                        if chunk:
                            f1.write(chunk)
                buf.seek(0)
                self.fs.put(
                    buf,
                    url=url,
                    size=None,
                    content_type=mimetypes.guess_type(url)[0],
                    kind=Kind.ATTACHMENT.value,
                )

    def cache_actor_icon(self, url: str) -> None:
        if self.fs.find_one({"url": url, "kind": Kind.ACTOR_ICON.value}):
            return
        i = load(url, self.user_agent)
        for size in [50, 80]:
            t1 = i.copy()
            t1.thumbnail((size, size))
            with BytesIO() as buf:
                with GzipFile(mode="wb", fileobj=buf) as f1:
                    t1.save(f1, format=i.format)
                buf.seek(0)
                self.fs.put(
                    buf,
                    url=url,
                    size=size,
                    content_type=i.get_format_mimetype(),
                    kind=Kind.ACTOR_ICON.value,
                )

    def save_upload(self, obuf: BytesIO, filename: str) -> str:
        # Remove EXIF metadata
        if filename.lower().endswith(".jpg") or filename.lower().endswith(".jpeg"):
            obuf.seek(0)
            with BytesIO() as buf2:
                piexif.remove(obuf.getvalue(), buf2)
                obuf.truncate(0)
                obuf.write(buf2.getvalue())

        obuf.seek(0)
        mtype = mimetypes.guess_type(filename)[0]
        with BytesIO() as gbuf:
            with GzipFile(mode="wb", fileobj=gbuf) as gzipfile:
                gzipfile.write(obuf.getvalue())

            gbuf.seek(0)
            oid = self.fs.put(
                gbuf,
                content_type=mtype,
                upload_filename=filename,
                kind=Kind.UPLOAD.value,
            )
            return str(oid)

    def cache(self, url: str, kind: Kind) -> None:
        if kind == Kind.ACTOR_ICON:
            self.cache_actor_icon(url)
        else:
            self.cache_attachment(url)

    def get_actor_icon(self, url: str, size: int) -> Any:
        return self.get_file(url, size, Kind.ACTOR_ICON)

    def get_attachment(self, url: str, size: int) -> Any:
        return self.get_file(url, size, Kind.ATTACHMENT)

    def get_file(self, url: str, size: int, kind: Kind) -> Any:
        return self.fs.find_one({"url": url, "size": size, "kind": kind.value})