1# -*- coding: utf-8 -*-
2"""
3    webapp2_extras.sessions
4    =======================
5
6    Lightweight but flexible session support for webapp2.
7
8    :copyright: 2011 by tipfy.org.
9    :license: Apache Sotware License, see LICENSE for details.
10"""
11import re
12
13import webapp2
14
15from webapp2_extras import securecookie
16from webapp2_extras import security
17
18#: Default configuration values for this module. Keys are:
19#:
20#: secret_key
21#:     Secret key to generate session cookies. Set this to something random
22#:     and unguessable. This is the only required configuration key:
23#:     an exception is raised if it is not defined.
24#:
25#: cookie_name
26#:     Name of the cookie to save a session or session id. Default is
27#:     `session`.
28#:
29#: session_max_age:
30#:     Default session expiration time in seconds. Limits the duration of the
31#:     contents of a cookie, even if a session cookie exists. If None, the
32#:     contents lasts as long as the cookie is valid. Default is None.
33#:
34#: cookie_args
35#:     Default keyword arguments used to set a cookie. Keys are:
36#:
37#:     - max_age: Cookie max age in seconds. Limits the duration
38#:       of a session cookie. If None, the cookie lasts until the client
39#:       is closed. Default is None.
40#:
41#:     - domain: Domain of the cookie. To work accross subdomains the
42#:       domain must be set to the main domain with a preceding dot, e.g.,
43#:       cookies set for `.mydomain.org` will work in `foo.mydomain.org` and
44#:       `bar.mydomain.org`. Default is None, which means that cookies will
45#:       only work for the current subdomain.
46#:
47#:     - path: Path in which the authentication cookie is valid.
48#:       Default is `/`.
49#:
50#:     - secure: Make the cookie only available via HTTPS.
51#:
52#:     - httponly: Disallow JavaScript to access the cookie.
53#:
54#: backends
55#:     A dictionary of available session backend classes used by
56#:     :meth:`SessionStore.get_session`.
57default_config = {
58    'secret_key':      None,
59    'cookie_name':     'session',
60    'session_max_age': None,
61    'cookie_args': {
62        'max_age':     None,
63        'domain':      None,
64        'path':        '/',
65        'secure':      None,
66        'httponly':    False,
67    },
68    'backends': {
69        'securecookie': 'webapp2_extras.sessions.SecureCookieSessionFactory',
70        'datastore':    'webapp2_extras.appengine.sessions_ndb.' \
71                        'DatastoreSessionFactory',
72        'memcache':     'webapp2_extras.appengine.sessions_memcache.' \
73                        'MemcacheSessionFactory',
74    },
75}
76
77_default_value = object()
78
79
80class _UpdateDictMixin(object):
81    """Makes dicts call `self.on_update` on modifications.
82
83    From werkzeug.datastructures.
84    """
85
86    on_update = None
87
88    def calls_update(name):
89        def oncall(self, *args, **kw):
90            rv = getattr(super(_UpdateDictMixin, self), name)(*args, **kw)
91            if self.on_update is not None:
92                self.on_update()
93            return rv
94        oncall.__name__ = name
95        return oncall
96
97    __setitem__ = calls_update('__setitem__')
98    __delitem__ = calls_update('__delitem__')
99    clear = calls_update('clear')
100    pop = calls_update('pop')
101    popitem = calls_update('popitem')
102    setdefault = calls_update('setdefault')
103    update = calls_update('update')
104    del calls_update
105
106
107class SessionDict(_UpdateDictMixin, dict):
108    """A dictionary for session data."""
109
110    __slots__ = ('container', 'new', 'modified')
111
112    def __init__(self, container, data=None, new=False):
113        self.container = container
114        self.new = new
115        self.modified = False
116        dict.update(self, data or ())
117
118    def pop(self, key, *args):
119        # Only pop if key doesn't exist, do not alter the dictionary.
120        if key in self:
121            return super(SessionDict, self).pop(key, *args)
122        if args:
123            return args[0]
124        raise KeyError(key)
125
126    def on_update(self):
127        self.modified = True
128
129    def get_flashes(self, key='_flash'):
130        """Returns a flash message. Flash messages are deleted when first read.
131
132        :param key:
133            Name of the flash key stored in the session. Default is '_flash'.
134        :returns:
135            The data stored in the flash, or an empty list.
136        """
137        return self.pop(key, [])
138
139    def add_flash(self, value, level=None, key='_flash'):
140        """Adds a flash message. Flash messages are deleted when first read.
141
142        :param value:
143            Value to be saved in the flash message.
144        :param level:
145            An optional level to set with the message. Default is `None`.
146        :param key:
147            Name of the flash key stored in the session. Default is '_flash'.
148        """
149        self.setdefault(key, []).append((value, level))
150
151
152class BaseSessionFactory(object):
153    """Base class for all session factories."""
154
155    #: Name of the session.
156    name = None
157    #: A reference to :class:`SessionStore`.
158    session_store = None
159    #: Keyword arguments to save the session.
160    session_args = None
161    #: The session data, a :class:`SessionDict` instance.
162    session = None
163
164    def __init__(self, name, session_store):
165        self.name = name
166        self.session_store = session_store
167        self.session_args = session_store.config['cookie_args'].copy()
168        self.session = None
169
170    def get_session(self, max_age=_default_value):
171        raise NotImplementedError()
172
173    def save_session(self, response):
174        raise NotImplementedError()
175
176
177class SecureCookieSessionFactory(BaseSessionFactory):
178    """A session factory that stores data serialized in a signed cookie.
179
180    Signed cookies can't be forged because the HMAC signature won't match.
181
182    This is the default factory passed as the `factory` keyword to
183    :meth:`SessionStore.get_session`.
184
185    .. warning::
186       The values stored in a signed cookie will be visible in the cookie,
187       so do not use secure cookie sessions if you need to store data that
188       can't be visible to users. For this, use datastore or memcache sessions.
189    """
190
191    def get_session(self, max_age=_default_value):
192        if self.session is None:
193            data = self.session_store.get_secure_cookie(self.name,
194                                                        max_age=max_age)
195            new = data is None
196            self.session = SessionDict(self, data=data, new=new)
197
198        return self.session
199
200    def save_session(self, response):
201        if self.session is None or not self.session.modified:
202            return
203
204        self.session_store.save_secure_cookie(
205            response, self.name, dict(self.session), **self.session_args)
206
207
208class CustomBackendSessionFactory(BaseSessionFactory):
209    """Base class for sessions that use custom backends, e.g., memcache."""
210
211    #: The session unique id.
212    sid = None
213
214    #: Used to validate session ids.
215    _sid_re = re.compile(r'^\w{22}$')
216
217    def get_session(self, max_age=_default_value):
218        if self.session is None:
219            data = self.session_store.get_secure_cookie(self.name,
220                                                        max_age=max_age)
221            sid = data.get('_sid') if data else None
222            self.session = self._get_by_sid(sid)
223
224        return self.session
225
226    def _get_by_sid(self, sid):
227        raise NotImplementedError()
228
229    def _is_valid_sid(self, sid):
230        """Check if a session id has the correct format."""
231        return sid and self._sid_re.match(sid) is not None
232
233    def _get_new_sid(self):
234        return security.generate_random_string(entropy=128)
235
236
237class SessionStore(object):
238    """A session provider for a single request.
239
240    The session store can provide multiple sessions using different keys,
241    even using different backends in the same request, through the method
242    :meth:`get_session`. By default it returns a session using the default key.
243
244    To use, define a base handler that extends the dispatch() method to start
245    the session store and save all sessions at the end of a request::
246
247        import webapp2
248
249        from webapp2_extras import sessions
250
251        class BaseHandler(webapp2.RequestHandler):
252            def dispatch(self):
253                # Get a session store for this request.
254                self.session_store = sessions.get_store(request=self.request)
255
256                try:
257                    # Dispatch the request.
258                    webapp2.RequestHandler.dispatch(self)
259                finally:
260                    # Save all sessions.
261                    self.session_store.save_sessions(self.response)
262
263            @webapp2.cached_property
264            def session(self):
265                # Returns a session using the default cookie key.
266                return self.session_store.get_session()
267
268    Then just use the session as a dictionary inside a handler::
269
270        # To set a value:
271        self.session['foo'] = 'bar'
272
273        # To get a value:
274        foo = self.session.get('foo')
275
276    A configuration dict can be passed to :meth:`__init__`, or the application
277    must be initialized with the ``secret_key`` configuration defined. The
278    configuration is a simple dictionary::
279
280        config = {}
281        config['webapp2_extras.sessions'] = {
282            'secret_key': 'my-super-secret-key',
283        }
284
285        app = webapp2.WSGIApplication([
286            ('/', HomeHandler),
287        ], config=config)
288
289    Other configuration keys are optional.
290    """
291
292    #: Configuration key.
293    config_key = __name__
294
295    def __init__(self, request, config=None):
296        """Initializes the session store.
297
298        :param request:
299            A :class:`webapp2.Request` instance.
300        :param config:
301            A dictionary of configuration values to be overridden. See
302            the available keys in :data:`default_config`.
303        """
304        self.request = request
305        # Base configuration.
306        self.config = request.app.config.load_config(self.config_key,
307            default_values=default_config, user_values=config,
308            required_keys=('secret_key',))
309        # Tracked sessions.
310        self.sessions = {}
311
312    @webapp2.cached_property
313    def serializer(self):
314        # Serializer and deserializer for signed cookies.
315        return securecookie.SecureCookieSerializer(self.config['secret_key'])
316
317    def get_backend(self, name):
318        """Returns a configured session backend, importing it if needed.
319
320        :param name:
321            The backend keyword.
322        :returns:
323            A :class:`BaseSessionFactory` subclass.
324        """
325        backends = self.config['backends']
326        backend = backends[name]
327        if isinstance(backend, basestring):
328            backend = backends[name] = webapp2.import_string(backend)
329
330        return backend
331
332    # Backend based sessions --------------------------------------------------
333
334    def _get_session_container(self, name, factory):
335        if name not in self.sessions:
336            self.sessions[name] = factory(name, self)
337
338        return self.sessions[name]
339
340    def get_session(self, name=None, max_age=_default_value, factory=None,
341                    backend='securecookie'):
342        """Returns a session for a given name. If the session doesn't exist, a
343        new session is returned.
344
345        :param name:
346            Cookie name. If not provided, uses the ``cookie_name``
347            value configured for this module.
348        :param max_age:
349            A maximum age in seconds for the session to be valid. Sessions
350            store a timestamp to invalidate them if needed. If `max_age` is
351            None, the timestamp won't be checked.
352        :param factory:
353            A session factory that creates the session using the preferred
354            backend. For convenience, use the `backend` argument instead,
355            which defines a backend keyword based on the configured ones.
356        :param backend:
357            A configured backend keyword. Available ones are:
358
359            - ``securecookie``: uses secure cookies. This is the default
360              backend.
361            - ``datastore``: uses App Engine's datastore.
362            - ``memcache``:  uses App Engine's memcache.
363        :returns:
364            A dictionary-like session object.
365        """
366        factory = factory or self.get_backend(backend)
367        name = name or self.config['cookie_name']
368
369        if max_age is _default_value:
370            max_age = self.config['session_max_age']
371
372        container = self._get_session_container(name, factory)
373        return container.get_session(max_age=max_age)
374
375    # Signed cookies ----------------------------------------------------------
376
377    def get_secure_cookie(self, name, max_age=_default_value):
378        """Returns a deserialized secure cookie value.
379
380        :param name:
381            Cookie name.
382        :param max_age:
383            Maximum age in seconds for a valid cookie. If the cookie is older
384            than this, returns None.
385        :returns:
386            A secure cookie value or None if it is not set.
387        """
388        if max_age is _default_value:
389            max_age = self.config['session_max_age']
390
391        value = self.request.cookies.get(name)
392        if value:
393            return self.serializer.deserialize(name, value, max_age=max_age)
394
395    def set_secure_cookie(self, name, value, **kwargs):
396        """Sets a secure cookie to be saved.
397
398        :param name:
399            Cookie name.
400        :param value:
401            Cookie value. Must be a dictionary.
402        :param kwargs:
403            Options to save the cookie. See :meth:`get_session`.
404        """
405        assert isinstance(value, dict), 'Secure cookie values must be a dict.'
406        container = self._get_session_container(name,
407                                                SecureCookieSessionFactory)
408        container.get_session().update(value)
409        container.session_args.update(kwargs)
410
411    # Saving to a response object ---------------------------------------------
412
413    def save_sessions(self, response):
414        """Saves all sessions in a response object.
415
416        :param response:
417            A :class:`webapp.Response` object.
418        """
419        for session in self.sessions.values():
420            session.save_session(response)
421
422    def save_secure_cookie(self, response, name, value, **kwargs):
423        value = self.serializer.serialize(name, value)
424        response.set_cookie(name, value, **kwargs)
425
426
427# Factories -------------------------------------------------------------------
428
429
430#: Key used to store :class:`SessionStore` in the request registry.
431_registry_key = 'webapp2_extras.sessions.SessionStore'
432
433
434def get_store(factory=SessionStore, key=_registry_key, request=None):
435    """Returns an instance of :class:`SessionStore` from the request registry.
436
437    It'll try to get it from the current request registry, and if it is not
438    registered it'll be instantiated and registered. A second call to this
439    function will return the same instance.
440
441    :param factory:
442        The callable used to build and register the instance if it is not yet
443        registered. The default is the class :class:`SessionStore` itself.
444    :param key:
445        The key used to store the instance in the registry. A default is used
446        if it is not set.
447    :param request:
448        A :class:`webapp2.Request` instance used to store the instance. The
449        active request is used if it is not set.
450    """
451    request = request or webapp2.get_request()
452    store = request.registry.get(key)
453    if not store:
454        store = request.registry[key] = factory(request)
455
456    return store
457
458
459def set_store(store, key=_registry_key, request=None):
460    """Sets an instance of :class:`SessionStore` in the request registry.
461
462    :param store:
463        An instance of :class:`SessionStore`.
464    :param key:
465        The key used to retrieve the instance from the registry. A default
466        is used if it is not set.
467    :param request:
468        A :class:`webapp2.Request` instance used to retrieve the instance. The
469        active request is used if it is not set.
470    """
471    request = request or webapp2.get_request()
472    request.registry[key] = store
473
474
475# Don't need to import it. :)
476default_config['backends']['securecookie'] = SecureCookieSessionFactory
477