From d11ef986a46409c8197383dbc4f87db08c8a5eaf Mon Sep 17 00:00:00 2001
From: MinRK <benjaminrk@gmail.com>
Date: Tue, 30 Aug 2011 16:27:12 -0700
Subject: [PATCH] authenticate Websockets with the session cookie

Now all Notebook connections are authenticated.
---
 IPython/frontend/html/notebook/handlers.py    | 68 ++++++++++++++++---
 .../html/notebook/static/js/kernel.js         |  6 ++
 2 files changed, 65 insertions(+), 9 deletions(-)

diff --git a/IPython/frontend/html/notebook/handlers.py b/IPython/frontend/html/notebook/handlers.py
index c4e588b4f..6f953fc6c 100644
--- a/IPython/frontend/html/notebook/handlers.py
+++ b/IPython/frontend/html/notebook/handlers.py
@@ -16,6 +16,9 @@ Authors:
 # Imports
 #-----------------------------------------------------------------------------
 
+import logging
+import Cookie
+
 from tornado import web
 from tornado import websocket
 
@@ -165,20 +168,58 @@ class ZMQStreamHandler(websocket.WebSocketHandler):
         else:
             self.write_message(msg)
 
+class AuthenticatedZMQStreamHandler(ZMQStreamHandler):
+    def open(self, kernel_id):
+        self.kernel_id = kernel_id
+        self.session = Session()
+        self.save_on_message = self.on_message
+        self.on_message = self.on_first_message
+    
+    def get_current_user(self):
+        password = self.get_secure_cookie("password")
+        if password is None:
+            # clear cookies, to prevent future Invalid cookie signature warnings
+            self._cookies = Cookie.SimpleCookie()
+        if self.application.password and self.application.password != password:
+            return None
+        return self.get_secure_cookie("user") or 'anonymous'
+    
+    def _inject_cookie_message(self, msg):
+        """Inject the first message, which is the document cookie,
+        for authentication."""
+        if isinstance(msg, unicode):
+            # Cookie can't constructor doesn't accept unicode strings for some reason
+            msg = msg.encode('utf8', 'replace')
+        try:
+            self._cookies = Cookie.SimpleCookie(msg)
+        except:
+            logging.warn("couldn't parse cookie string: %s",msg, exc_info=True)
+    
+    def on_first_message(self, msg):
+        self._inject_cookie_message(msg)
+        if self.get_current_user() is None:
+            logging.warn("Couldn't authenticate WebSocket connection")
+            raise web.HTTPError(403)
+        self.on_message = self.save_on_message
+                
 
-class IOPubHandler(ZMQStreamHandler):
+class IOPubHandler(AuthenticatedZMQStreamHandler):
 
     def initialize(self, *args, **kwargs):
         self._kernel_alive = True
         self._beating = False
         self.iopub_stream = None
         self.hb_stream = None
-
-    def open(self, kernel_id):
+    
+    def on_first_message(self, msg):
+        try:
+            super(IOPubHandler, self).on_first_message(msg)
+        except web.HTTPError:
+            self.close()
+            return
         km = self.application.kernel_manager
-        self.kernel_id = kernel_id
-        self.session = Session()
         self.time_to_dead = km.time_to_dead
+        kernel_id = self.kernel_id
         try:
             self.iopub_stream = km.create_iopub_stream(kernel_id)
             self.hb_stream = km.create_hb_stream(kernel_id)
@@ -187,9 +228,13 @@ class IOPubHandler(ZMQStreamHandler):
             # close the connection.
             if not self.stream.closed():
                 self.stream.close()
+            self.close()
         else:
             self.iopub_stream.on_recv(self._on_zmq_reply)
             self.start_hb(self.kernel_died)
+    
+    def on_message(self, msg):
+        pass
 
     def on_close(self):
         # This method can be called twice, once by self.kernel_died and once 
@@ -245,15 +290,20 @@ class IOPubHandler(ZMQStreamHandler):
         self.on_close()
 
 
-class ShellHandler(ZMQStreamHandler):
+class ShellHandler(AuthenticatedZMQStreamHandler):
 
     def initialize(self, *args, **kwargs):
         self.shell_stream = None
 
-    def open(self, kernel_id):
+    def on_first_message(self, msg):
+        try:
+            super(ShellHandler, self).on_first_message(msg)
+        except web.HTTPError:
+            self.close()
+            return
         km = self.application.kernel_manager
         self.max_msg_size = km.max_msg_size
-        self.kernel_id = kernel_id
+        kernel_id = self.kernel_id
         try:
             self.shell_stream = km.create_shell_stream(kernel_id)
         except web.HTTPError:
@@ -261,8 +311,8 @@ class ShellHandler(ZMQStreamHandler):
             # close the connection.
             if not self.stream.closed():
                 self.stream.close()
+            self.close()
         else:
-            self.session = Session()
             self.shell_stream.on_recv(self._on_zmq_reply)
 
     def on_message(self, msg):
diff --git a/IPython/frontend/html/notebook/static/js/kernel.js b/IPython/frontend/html/notebook/static/js/kernel.js
index 708b93323..46304b9e8 100644
--- a/IPython/frontend/html/notebook/static/js/kernel.js
+++ b/IPython/frontend/html/notebook/static/js/kernel.js
@@ -91,6 +91,12 @@ var IPython = (function (IPython) {
         console.log("Starting WS:", ws_url);
         this.shell_channel = new this.WebSocket(ws_url + "/shell");
         this.iopub_channel = new this.WebSocket(ws_url + "/iopub");
+        send_cookie = function(){
+            this.send(document.cookie);
+            console.log(this);
+        }
+        this.shell_channel.onopen = send_cookie;
+        this.iopub_channel.onopen = send_cookie;
     };