codekingpro/portable-devtools

1from hashlib import sha2562import sys3 4from flask import session, request, make_response, url_for, current_app, \5    redirect6from werkzeug.exceptions import Unauthorized7 8 9class Paranoid(object):10    def __init__(self, app=None):11        self.invalid_session_handler = self._default_invalid_session_handler12        if app:13            self.init_app(app)14 15    def init_app(self, app):16        @app.before_request17        def before_request():18            token = self.create_token()19            existing_token = self.get_token_from_session()20            if existing_token is None:21                # this is a new session, so we write our id in it22                self.write_token_to_session(token)23            elif existing_token != token:24                # this session is invalid, so we get rid of it25                if callable(self.invalid_session_handler):26                    response = make_response(self.invalid_session_handler())27                else:28                    if self.invalid_session_handler.startswith(29                            ('http://', 'https://', '/')):30                        url = self.invalid_session_handler31                    else:32                        url = url_for(self.invalid_session_handler)33                    response = redirect(url)34                self.clear_session(response)35                return response36 37    def on_invalid_session(self, f):38        self.invalid_session_handler = f39        return f40 41    def _default_invalid_session_handler(self):42        try:43            raise Unauthorized()44        except Exception as e:45            response = current_app.handle_user_exception(e)46        return response47 48    @property49    def redirect_view(self):50        return self.invalid_session_handler51 52    @redirect_view.setter53    def redirect_view(self, view):54        self.invalid_session_handler = view55 56    def _get_remote_addr(self):57        address = request.headers.get('X-Forwarded-For', request.remote_addr)58        if address is None:  # pragma: no cover59            address = 'x.x.x.x'60        address = address.encode('utf-8').split(b',')[0].strip()61        return address62 63    def create_token(self):64        """Create a session protection token for this client.65 66        This method generates a session protection token for the cilent, which67        consists in a hash of the user agent and the IP address. This method68        can be overriden by subclasses to implement different token generation69        algorithms.70        """71        user_agent = request.headers.get('User-Agent')72        if user_agent is None:  # pragma: no cover73            user_agent = 'no user agent'74        user_agent = user_agent.encode('utf-8')75        base = self._get_remote_addr() + b'|' + user_agent76        h = sha256()77        h.update(base)78        return h.hexdigest()79 80    def get_token_from_session(self):81        """Return the session protection token stored from the client session.82 83        This method retrieves the stored session protection token, or None if84        this is a brand new session that doesn't have a token in it. This85        default implementation finds the token in the user session. Subclasses86        can override this method and implement other storage methods.87        """88        return session.get('_paranoid_token')89 90    def write_token_to_session(self, token):91        """Write a session protection token to the client session.92 93        This methods writes the session protection token. This default94        implementation writes the token to the user session. Subclasses can95        override this method to implement other storage methods.96        """97        session['_paranoid_token'] = token98 99    def clear_session(self, response):100        """Clear the session.101 102        This method is invoked when the session is found to be invalid.103        Subclasses can override this method to implement a custom session104        reset.105        """106        session.clear()107 108        # if flask-login is installed, we try to clear the109        # "remember me" cookie, just in case it is set110        if 'flask_login' in sys.modules:111            remember_cookie = current_app.config.get('REMEMBER_COOKIE',112                                                     'remember_token')113            response.set_cookie(remember_cookie, '', expires=0, max_age=0)114