1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65
| import base64 import hashlib import hmac import json from flask import make_response, request
def cookie_encode(data, key): msg = base64.b64encode(json.dumps(data).encode('utf-8')) sig = base64.b64encode(hmac.new(key.encode('utf-8'), msg, digestmod=hashlib.md5).digest()) return '!' + sig.decode('utf-8') + '?' + msg.decode('utf-8')
def cookie_decode(data, key): if cookie_is_encoded(data): sig, msg = data.split('?', 1) sig = sig[1:] expected_sig = base64.b64encode(hmac.new(key.encode('utf-8'), msg.encode('utf-8'), digestmod=hashlib.md5).digest()).decode('utf-8') if _lscmp(sig, expected_sig): return json.loads(base64.b64decode(msg).decode('utf-8')) return None
def waf(data): blacklist = [b'R', b'secret', b'eval', b'file', b'compile', b'open', b'os.popen'] valid = False for word in blacklist: if word in data.encode('utf-8'): valid = True break return valid
def cookie_check(key, secret=None): data = request.cookies.get(key) if data and cookie_is_encoded(data): decoded_data = cookie_decode(data, secret) if decoded_data and waf(json.dumps(decoded_data)): return True return False
def get_cookie(key, default=None, secret=None): value = request.cookies.get(key) if secret and value: dec = cookie_decode(value, secret) return dec if dec else default return value or default
def cookie_is_encoded(data): return data.startswith('!') and '?' in data
def _lscmp(a, b): """Compare two strings in a way that is resistant to timing attacks.""" return not sum(0 if x == y else 1 for x, y in zip(a, b)) and len(a) == len(b)
def set_cookie(name, value, secret=None, **options): if secret: value = cookie_encode({name: value}, secret) elif not isinstance(value, str): raise TypeError('Secret key missing for non-string Cookie.') if len(value) > 4096: raise ValueError('Cookie value too long.') resp = make_response("success") resp.set_cookie(name, value, max_age=options.get('max_age', 3600)) return resp
|