forget-cancellare-vecchi-toot/lib/img_proxy.py

import requests
import threading
import redis
from flask import make_response, abort
import secrets
import hmac
import base64


class ImgProxyCache(object):
    def __init__(self, redis_uri='redis://', timeout=1, expire=60*60*24,
                 prefix='img_proxy', hmac_hash='sha1'):
        self.redis = redis.StrictRedis.from_url(redis_uri)
        self.timeout = timeout
        self.expire = expire
        self.prefix = prefix
        self.redis.client_setname('img_proxy')
        self.hash = hmac_hash

    def key(self, *args):
        return '{prefix}:{args}'.format(
                prefix=self.prefix, args=":".join(args))

    def token(self):
        t = self.redis.get(self.key('hmac_key'))
        if not t:
            t = secrets.token_urlsafe().encode('ascii')
            self.redis.set(self.key('hmac_key'), t)
        return t

    def identifier_for(self, url):
        url_hmac = hmac.new(self.token(), url.encode('UTF-8'), self.hash)
        return base64.urlsafe_b64encode(
                '{}:{}'.format(url_hmac.hexdigest(), url)
                .encode('UTF-8')
                ).strip(b'=')

    def url_for(self, identifier):
        try:
            padding = (4 - len(identifier)) % 4
            identifier += padding * '='
            identifier = base64.urlsafe_b64decode(identifier).decode('UTF-8')
            received_hmac, url = identifier.split(':', 1)
            url_hmac = hmac.new(self.token(), url.encode('UTF-8'), self.hash)
            if not hmac.compare_digest(url_hmac.hexdigest(), received_hmac):
                return None
        except Exception:
            return None
        return url

    def fetch_and_cache(self, url):
        resp = requests.get(url)
        if(resp.status_code != 200):
            return
        mime = resp.headers.get('content-type', 'application/octet-stream')
        self.redis.set(self.key('mime', url),
                       mime, px=self.expire*1000)
        self.redis.set(self.key('body', url),
                       resp.content, px=self.expire*1000)

    def respond(self, identifier):
        url = self.url_for(identifier)
        if not url:
            return abort(403)

        x_imgproxy_cache = 'HIT'
        mime = self.redis.get(self.key('mime', url))
        body = self.redis.get(self.key('body', url))
        if not body or not mime:
            x_imgproxy_cache = 'MISS'
            if self.redis.set(
                    self.key('lock', url), 1, nx=True, ex=10*self.timeout):
                t = threading.Thread(target=self.fetch_and_cache, args=(url,))
                t.start()
                t.join(self.timeout)
                mime = self.redis.get(self.key('mime', url))
                body = self.redis.get(self.key('body', url))
        if not body or not mime:
            return abort(404)

        resp = make_response(body, 200)
        resp.headers.set('content-type', mime)
        resp.headers.set('x-imgproxy-cache', x_imgproxy_cache)
        resp.headers.set('cache-control', 'max-age={}'.format(self.expire))
        return resp