From a0830f88f7a18b9547c492de1c336737a9159821 Mon Sep 17 00:00:00 2001
From: Andrew Godwin <andrew@aeracode.org>
Date: Wed, 2 Sep 2015 00:09:53 -0700
Subject: [PATCH] Initial django session and auth decorators

---
 channels/decorators.py | 70 ++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 70 insertions(+)

diff --git a/channels/decorators.py b/channels/decorators.py
index 6012938..f19db2c 100644
--- a/channels/decorators.py
+++ b/channels/decorators.py
@@ -4,6 +4,7 @@ from importlib import import_module
 
 from django.conf import settings
 from django.utils import six
+from django.contrib import auth
 
 from channels import channel_backends, DEFAULT_CHANNEL_BACKEND
 
@@ -28,6 +29,75 @@ def consumer(*channels, **kwargs):
 
 # TODO: Sessions, auth
 
+def http_session(func):
+    """
+    Wraps a HTTP or WebSocket consumer (or any consumer of messages
+    that provides a "COOKIES" or "GET" attribute) to provide a "session"
+    attribute that behaves like request.session; that is, it's hung off of
+    a per-user session key that is saved in a cookie or passed as the
+    "session_key" GET parameter.
+
+    It won't automatically create and set a session cookie for users who
+    don't have one - that's what SessionMiddleware is for, this is a simpler
+    read-only version for more low-level code.
+
+    If a user does not have a session we can inflate, the "session" attribute will
+    be None, rather than an empty session you can write to.
+    """
+    @functools.wraps(func)
+    def inner(*args, **kwargs):
+        if "COOKIES" not in kwargs and "GET" not in kwargs:
+            print kwargs
+            raise ValueError("No COOKIES or GET sent to consumer; this decorator can only be used on messages containing at least one.")
+        # Make sure there's a session key
+        session_key = None
+        if "GET" in kwargs:
+            session_key = kwargs['GET'].get("session_key")
+        if "COOKIES" in kwargs and session_key is None:
+            session_key = kwargs['COOKIES'].get(settings.SESSION_COOKIE_NAME)
+        # Make a session storage
+        if session_key:
+            session_engine = import_module(settings.SESSION_ENGINE)
+            session = session_engine.SessionStore(session_key=session_key)
+        else:
+            session = None
+        kwargs['session'] = session
+        # Run the consumer
+        result = func(*args, **kwargs)
+        # Persist session if needed (won't be saved if error happens)
+        if session is not None and session.modified:
+            session.save()
+        return result
+    return inner
+
+
+def http_django_auth(func):
+    """
+    Wraps a HTTP or WebSocket consumer (or any consumer of messages
+    that provides a "COOKIES" attribute) to provide both a "session"
+    attribute and a "user" attibute, like AuthMiddleware does.
+
+    This runs http_session() to get a session to hook auth off of.
+    If the user does not have a session cookie set, both "session"
+    and "user" will be None.
+    """
+    @http_session
+    @functools.wraps(func)
+    def inner(*args, **kwargs):
+        # If we didn't get a session, then we don't get a user
+        if kwargs['session'] is None:
+            kwargs['user'] = None
+        # Otherwise, be a bit naughty and make a fake Request with just
+        # a "session" attribute (later on, perhaps refactor contrib.auth to
+        # pass around session rather than request)
+        else:
+            fake_request = type("FakeRequest", (object, ), {"session": kwargs['session']})
+            kwargs['user'] = auth.get_user(fake_request)
+        # Run the consumer
+        return func(*args, **kwargs)
+    return inner
+
+
 def send_channel_session(func):
     """
     Provides a session-like object called "channel_session" to consumers