smuggler.py

import hashlib
import hmac
import json
import os

import klein
import requests

from twisted.internet.threads import deferToThread
from twisted.web.http import stringToDatetime


__version__ = "1.0"


SECRET_KEY = os.environ["SMUGGLER_KEY"]

USER_AGENT = os.environ.get(
    "SMUGGLER_USER_AGENT",
    "Smuggler/{}".format(__version__),
)

TIMEOUT = int(os.environ.get("SMUGGLER_TIMEOUT", 10))

MAX_IMAGE_SIZE = int(os.environ.get("SMUGGLER_SIZE_LIMIT", 5242880))

with open(
    os.environ.get(
        "SMUGGLER_MIMETYPES",
        os.path.join(os.path.dirname(__file__), "mimetypes.json"),
        )) as fp:
    MIMETYPES = json.load(fp)


def processResponse(response, request):
    # Ensure that we know what size the image is
    if "Content-Length" not in response.headers:
        request.setResponseCode(413)
        return "cannot determine size of source image"

    # Ensure that the image isn't too large for us
    if int(response.headers["Content-Length"]) > MAX_IMAGE_SIZE:
        request.setResponseCode(413)
        return "source image exceeds max allowed size"

    contentType = response.headers["Content-Type"]

    # Verify Content Type
    contentTypePrefix = contentType.split(";")[0].strip()
    if contentTypePrefix not in MIMETYPES:
        request.setResponseCode(415)
        return "non-Image content-type returned '%s'" % (contentTypePrefix,)

    # Shuttle over various values from the source
    request.setHeader("Content-Type", contentType)

    if "ETag" in response.headers:
        request.setETag(response.headers["ETag"])

    if "Last-Modified" in response.headers:
        request.setLastModified(
            stringToDatetime(response.headers["Last-Modified"]),
        )

    if "Expires" in response.headers:
        request.setHeader("Expires", response.headers["Expires"])

    if "Cache-Control" in response.headers:
        request.setHeader("Cache-Control", response.headers["Cache-Control"])

    if "Content-Encoding" in response.headers:
        request.setHeader(
            "Content-Encoding",
            response.headers["Content-Encoding"],
        )

    return response.raw.read()


@klein.route("/<hmacRequest>/<path:url>", methods=["HEAD", "GET"])
def proxyImage(request, hmacRequest, url):
    hmacDigest = hmac.new(
        SECRET_KEY, url,
        digestmod=hashlib.sha224,
    ).hexdigest()

    # if hmacDigest != hmacRequest:  # TODO: Constant Time
    #     request.setResponseCode(404)
    #     return "checksum mismatch"

    # Using a Thread to do HTTP because nothing I've tried in Twisted actually
    # makes HTTP requests without losing headers like Content-Length.
    d = deferToThread(
        requests.get,
        url,
        headers={
            "Accept": "images/*",
            "Accept-Encoding": request.getHeader("Accept-Encoding"),
            "User-Agent": USER_AGENT,
        },
        timeout=TIMEOUT,
        stream=True,
        verify=False,
    )
    d.addCallback(processResponse, request)

    return d


def wrap_resource(fn):
    def _resource():
        r = fn()

        real_render = r.render

        def _render(request):
            request.setHeader("Server", USER_AGENT)
            request.setHeader("X-Content-Type-Options", "nosniff")
            request.setHeader("X-Frame-Options", "deny")
            return real_render(request)

        r.render = _render

        return r
    return _resource


resource = wrap_resource(klein.resource)