增加websocket认证

connect.py

@@ -269,8 +269,6 @@ class Tty(object):
         today_connect_log_dir = os.path.join(tty_log_dir, date_start)
         log_file_path = os.path.join(today_connect_log_dir, '%s_%s_%s' % (self.username, self.asset_name, time_start))
 
-
-
         try:
             is_dir(today_connect_log_dir, mode=0777)
         except OSError:
@@ -308,7 +306,7 @@ class Tty(object):
         # 2. get 映射用户
         # 3. get 映射用户的账号，密码或者key
         # self.connect_info = {'user': '', 'asset': '', 'ip': '', 'port': 0, 'role_name': '', 'role_pass': '', 'role_key': ''}
-        self.connect_info = {'user': 'a', 'asset': 'b', 'ip': '127.0.0.1', 'port': 22, 'role_name': 'root', 'role_pass': 'redhat', 'role_key': ''}
+        self.connect_info = {'user': 'a', 'asset': 'b', 'ip': '127.0.0.1', 'port': 22, 'role_name': 'root', 'role_pass': '', 'role_key': '/root/.ssh/id_rsa.bak'}
         return self.connect_info
 
     def get_connection(self):

jlog/views.py

@@ -104,6 +104,10 @@ def log_record(request):
 
 
 def web_terminal(request):
-    web_terminal_uri = 'ws://%s/terminal' % web_socket_host
+    #username = get_session.get('username', '')
+    token = request.COOKIES.get('sessionid')
+    username = request.user.username
+    asset_name = '127.0.0.1'
+    web_terminal_uri = 'ws://%s/terminal?username=%s&asset_name=%s&token=%s' % (web_socket_host, username, asset_name, token)
     return render_to_response('jlog/web_terminal.html', locals())
 

run_websocket.py

@@ -7,7 +7,7 @@ import os
 import sys
 import os.path
 import threading
-import uuid
+import urllib
 
 import tornado.ioloop
 import tornado.options
@@ -15,6 +15,7 @@ import tornado.web
 import tornado.websocket
 import tornado.httpserver
 import tornado.gen
+import tornado.httpclient
 from tornado.websocket import WebSocketClosedError
 
 from tornado.options import define, options
@@ -40,6 +41,20 @@ define("port", default=3000, help="run on the given port", type=int)
 define("host", default='0.0.0.0', help="run port on", type=str)
 
 
+def require_auth(func):
+    def _deco(request, *args, **kwargs):
+        username = request.get_argument('username', '')
+        asset_name = request.get_argument('asset_name', '')
+        token = request.get_argument('token', '')
+        print username, asset_name, token
+        client = tornado.httpclient.HTTPClient()
+        # response = client.fetch('http://some/url') + urllib.urlencode({'username': username,
+        #                                                                'asset_name': asset_name, 'token': token})
+        return request.close()
+        # return func(request, *args, **kwargs)
+    return _deco
+
+
 class MyThread(threading.Thread):
     def __init__(self, *args, **kwargs):
         super(MyThread, self).__init__(*args, **kwargs)
@@ -121,6 +136,7 @@ class MonitorHandler(tornado.websocket.WebSocketHandler):
     def check_origin(self, origin):
         return True
 
+    @require_auth
     def open(self):
         # 获取监控的path
         self.file_path = self.get_argument('file_path', '')
@@ -174,6 +190,8 @@ class WebTerminalKillHandler(tornado.web.RequestHandler):
             print ws.id
             if ws.id == int(ws_id):
                 print "killed"
+                ws.log.is_finished = True
+                ws.log.save()
                 ws.close()
         print len(WebTerminalHandler.clients)
 
@@ -194,7 +212,12 @@ class WebTerminalHandler(tornado.websocket.WebSocketHandler):
     def check_origin(self, origin):
         return True
 
+    @require_auth
     def open(self):
+        asset_name = self.get_argument('asset_name', '')
+        username = self.get_argument('username', '')
+        token = self.get_argument('token', '')
+        print asset_name, username, token
         self.term = WebTty('a', 'b')
         self.term.get_connection()
         self.channel = self.term.ssh.invoke_shell(term='xterm')
@@ -223,11 +246,14 @@ class WebTerminalHandler(tornado.websocket.WebSocketHandler):
         print 'On_close'
         if self in WebTerminalHandler.clients:
             WebTerminalHandler.clients.remove(self)
-        self.log_file_f.write('End time is %s' % datetime.datetime.now())
-        self.log.is_finished = True
-        self.log.end_time = datetime.datetime.now()
-        self.log.save()
-        self.close()
+        try:
+            self.log_file_f.write('End time is %s' % datetime.datetime.now())
+            self.log.is_finished = True
+            self.log.end_time = datetime.datetime.now()
+            self.log.save()
+            self.close()
+        except AttributeError:
+            pass
 
     def forward_outbound(self):
         self.log_file_f, self.log_time_f, self.log = self.term.get_log_file()