[Feature] 完成app注册sdk对接

coco/app.py

@@ -5,10 +5,9 @@ import logging
 
 from .config import Config
 from .sshd import SSHServer
-from .ws import WSServer
+from .httpd import HttpServer
 from .logging import create_logger
 from .sdk import AppService
-from .auth import AppAccessKey
 
 
 __version__ = '0.4.0'
@@ -27,6 +26,7 @@ class Coco:
         'SSHD_PORT': 2222,
         'WS_PORT': 5000,
         'ACCESS_KEY': '',
+        'ACCESS_KEY_ENV': 'COCO_ACCESS_KEY',
         'ACCESS_KEY_FILE': os.path.join(BASE_DIR, 'keys', '.access_key'),
         'SECRET_KEY': None,
         'LOG_LEVEL': 'INFO',
@@ -43,7 +43,6 @@ class Coco:
         self.config = self.config_class(BASE_DIR, defaults=self.default_config)
         self.sessions = []
         self.clients = []
-        self.ws = None
         self.root_path = root_path
         self.name = name
         self.lock = threading.Lock()
@@ -55,7 +54,7 @@ class Coco:
         if root_path is None:
             self.root_path = BASE_DIR
 
-        self.make_logger()
+        self.httpd = None
         self.sshd = None
         self.running = True
 
@@ -63,8 +62,9 @@ class Coco:
         create_logger(self)
 
     def prepare(self):
+        self.make_logger()
         self.sshd = SSHServer(self)
-        self.ws = WSServer(self)
+        self.httpd = HttpServer(self)
         self.initial_service()
 
     def heartbeat(self):
@@ -94,7 +94,7 @@ class Coco:
         thread.start()
 
     def run_ws(self):
-        thread = threading.Thread(target=self.ws.run, args=())
+        thread = threading.Thread(target=self.httpd.run, args=())
         thread.daemon = True
         thread.start()
 
@@ -123,6 +123,7 @@ class Coco:
 
     def initial_service(self):
         self.service = AppService(self)
+        self.service.initial()
 
     def monitor_session(self):
         pass

coco/auth.py

@@ -10,19 +10,6 @@ from . import utils
 from .exception import LoadAccessKeyError
 
 
-def make_signature(access_key_secret, date=None):
-    if isinstance(date, bytes):
-        date = date.decode("utf-8")
-    if isinstance(date, int):
-        date_gmt = utils.http_date(date)
-    elif date is None:
-        date_gmt = utils.http_date(int(time.time()))
-    else:
-        date_gmt = date
-    data = str(access_key_secret) + "\n" + date_gmt
-    return utils.content_md5(data)
-
-
 class AccessKeyAuth(object):
     def __init__(self, access_key):
         self.id = access_key.id
@@ -34,49 +21,6 @@ class AccessKeyAuth(object):
         req.headers['Authorization'] = "Sign {0}:{1}".format(self.id, signature)
         return req
 
-#
-# class AccessTokenAuth(object):
-#     def __init__(self, token):
-#         self.token = token
-#
-#     def sign_request(self, req):
-#         req.headers['Authorization'] = 'Bearer {0}'.format(self.token)
-#         return req
-#
-#
-# class SessionAuth(object):
-#     def __init__(self, session_id, csrf_token):
-#         self.session_id = session_id
-#         self.csrf_token = csrf_token
-#
-#     def sign_request(self, req):
-#         cookie = [v for v in req.headers.get('Cookie', '').split(';')
-#                   if v.strip()]
-#         cookie.extend(['sessionid='+self.session_id,
-#                        'csrftoken='+self.csrf_token])
-#         req.headers['Cookie'] = ';'.join(cookie)
-#         req.headers['X-CSRFTOKEN'] = self.csrf_token
-#         return req
-
-
-# class Auth(object):
-#     def __init__(self, token=None, access_key_id=None,
-#                  access_key_secret=None,
-#                  session_id=None, csrf_token=None):
-#
-#         if token is not None:
-#             self.instance = AccessTokenAuth(token)
-#         elif access_key_id and access_key_secret:
-#             self.instance = AccessKeyAuth(access_key_id, access_key_secret)
-#         elif session_id and csrf_token:
-#             self.instance = SessionAuth(session_id, csrf_token)
-#         else:
-#             raise SyntaxError('Need token or access_key_id, access_key_secret '
-#                               'or session_id, csrf_token')
-#
-#     def sign_request(self, req):
-#         return self.instance.sign_request(req)
-
 
 class AccessKey(object):
     def __init__(self, id=None, secret=None):
@@ -97,13 +41,13 @@ class AccessKey(object):
     @classmethod
     def load_from_val(cls, val, **kwargs):
         id, secret = cls.clean(val, **kwargs)
-        return cls(id, secret)
+        return cls(id=id, secret=secret)
 
     @classmethod
     def load_from_env(cls, env, **kwargs):
         value = os.environ.get(env)
         id, secret = cls.clean(value, **kwargs)
-        return cls(id, secret)
+        return cls(id=id, secret=secret)
 
     @classmethod
     def load_from_f(cls, f, **kwargs):
@@ -117,7 +61,7 @@ class AccessKey(object):
                     break
             f.close()
         id, secret = cls.clean(value, **kwargs)
-        return cls(id, secret)
+        return cls(id=id, secret=secret)
 
     def save_to_f(self, f, silent=False):
         if isinstance(f, str):
@@ -131,44 +75,57 @@ class AccessKey(object):
         finally:
             f.close()
 
-    def __nonzero__(self):
+    def __bool__(self):
         return bool(self.id and self.secret)
-    __bool__ = __nonzero__
 
     def __str__(self):
         return '{0}:{1}'.format(self.id, self.secret)
-    __repr__ = __str__
+
+    def __repr__(self):
+        return '{0}:{1}'.format(self.id, self.secret)
 
 
 class AppAccessKey(AccessKey):
     """使用Access key来认证"""
 
-    def __init__(self, app, id=None, secret=None):
+    def __init__(self, id=None, secret=None):
         super().__init__(id=id, secret=secret)
+        self.app = None
+
+    def set_app(self, app):
         self.app = app
-        self._key_store = app.config['ACCESS_KEY_STORE']
-        self._key_env = app.config['ACCESS_KEY_ENV']
-        self._key_val = app.config['ACCESS_KEY']
+
+    @property
+    def _key_env(self):
+        return self.app.config['ACCESS_KEY_ENV']
+
+    @property
+    def _key_val(self):
+        return self.app.config['ACCESS_KEY']
+
+    @property
+    def _key_file(self):
+        return self.app.config['ACCESS_KEY_FILE']
 
     def load_from_conf_env(self, sep=':', silent=False):
         return super().load_from_env(self._key_env, sep=sep, silent=silent)
 
-    def load_from_conf(self, sep=':', silent=False):
+    def load_from_conf_val(self, sep=':', silent=False):
         return super().load_from_val(self._key_val, sep=sep, silent=silent)
 
-    def load_from_key_store(self, sep=':', silent=False):
-        return super().load_from_f(self._key_store, sep=sep, silent=silent)
+    def load_from_conf_file(self, sep=':', silent=False):
+        return super().load_from_f(self._key_file, sep=sep, silent=silent)
 
     def load(self, **kwargs):
         """Should return access_key_id, access_key_secret"""
-        for method in [self.load_from_env,
-                       self.load_from_conf,
-                       self.load_from_key_store]:
+        for method in [self.load_from_conf_env,
+                       self.load_from_conf_val,
+                       self.load_from_conf_file]:
             try:
                 return method(**kwargs)
             except LoadAccessKeyError:
                 continue
         return None
 
-    def save_to_key_store(self):
-        return super().save_to_f(self._key_store)
\ No newline at end of file
+    def save_to_file(self):
+        return super().save_to_f(self._key_file)
\ No newline at end of file

coco/ws.py → coco/httpd.py

@@ -72,7 +72,7 @@ class MonitorWehSocketHandler(BaseWehSocketHandler):
     pass
 
 
-class WSServer:
+class HttpServer:
     routers = [
         (r'/ws/interactive/', InteractiveWehSocketHandler),
         (r'/ws/proxy/(?P<asset_id>[0-9]+)/(?P<system_user_id>[0-9]+)/', ProxyWehSocketHandler),

coco/sdk.py

@@ -80,7 +80,7 @@ class AppRequests(object):
     @staticmethod
     def clean_result(resp):
         if resp.status_code >= 500:
-            raise ResponseError("Response code is {0.code}: {0.text}".format(resp))
+            raise ResponseError("Response code is {0.status_code}: {0.text}".format(resp))
 
         try:
             _ = resp.json()
@@ -138,82 +138,85 @@ class AppService:
     def __init__(self, app):
         self.app = app
         self.access_key = None
-        self.requests = AppRequests(app.config['JUMPSERVER_ENDPOINT'])
-        self.prepare()
+        self.requests = AppRequests(app.config['CORE_HOST'])
 
-    def prepare(self):
+    def initial(self):
         self.load_access_key()
         self.set_auth()
         self.valid_auth()
 
     def load_access_key(self):
         # Must be get access key if not register it
-        self.access_key = self.access_key_class(self.app).load()
+        self.access_key = self.access_key_class()
+        self.access_key.set_app(self.app)
+        self.access_key = self.access_key.load()
         if self.access_key is None:
+            logger.info("No access key found, register it")
             self.register_and_save()
 
     def set_auth(self):
         self.requests.auth = AccessKeyAuth(self.access_key)
 
     def valid_auth(self):
-        while True:
-            delay = 1
+        delay = 1
+        while delay < 300:
             if self.heatbeat() is None:
                 msg = "Access key is not valid or need admin " \
                       "accepted, waiting %d s" % delay
                 logger.info(msg)
-                delay += 1
-                time.sleep(1)
+                delay += 3
+                time.sleep(3)
             else:
                 break
+        if delay >= 300:
+            logger.info("Start timeout")
+            sys.exit()
 
     def register_and_save(self):
         self.register()
         self.save_access_key()
 
     def save_access_key(self):
-        self.access_key.save_to_key_store()
+        self.access_key.save_to_file()
 
     def register(self):
         try:
-            resp = self.requests.post('terminal-register',
-                                      data={'name': self.app.name},
-                                      use_auth=False)
+            resp = self.requests.post(
+                'terminal-register', data={'name': self.app.name}, use_auth=False
+            )
         except (RequestError, ResponseError) as e:
             logger.error(e)
-            sys.exit()
+            return
 
         if resp.status_code == 201:
             access_key = resp.json()['access_key']
             access_key_id = access_key['id']
             access_key_secret = access_key['secret']
             self.access_key = self.access_key_class(
-                app=self.app, id=access_key_id, secret=access_key_secret
+                id=access_key_id, secret=access_key_secret
             )
+            self.access_key.set_app(self.app)
             logger.info('Register app success: %s' % access_key_id,)
-        elif resp.status_code == 400:
+        elif resp.status_code == 409:
             msg = '{} exist already, register failed'.format(self.app.name)
             logging.error(msg)
             sys.exit()
         else:
-            logging.error('Register terminal {} failed unknown'.format(self.app.name))
+            logging.error('Register terminal {} failed unknown: {}'.format(self.app.name, resp.json()))
             sys.exit()
 
-    def wait_util_accept(self):
-        while True:
-            if self.heatbeat() is None:
-                time.sleep(1)
-            else:
-                break
-
     def heatbeat(self):
         """和Jumpserver维持心跳, 当Terminal断线后,jumpserver可以知晓
 
         Todo: Jumpserver发送的任务也随heatbeat返回, 并执行,如 断开某用户
         """
-        r, content = self.requests.post('terminal-heatbeat', use_auth=True)
-        if r.status_code == 201:
-            return content
+        try:
+            resp = self.requests.post('terminal-heatbeat', use_auth=True)
+        except (ResponseError, RequestError):
+            return None
+
+        if resp.status_code == 201:
+            return True
         else:
             return None
 
@@ -224,9 +227,9 @@ class AppService:
             'asset_id': asset_id,
             'system_user_id': system_user_id,
         }
-        r, content = self.get('validate-user-asset-permission',
-                              use_auth=True,
-                              params=params)
+        r, content = self.requests.get(
+            'validate-user-asset-permission', use_auth=True, params=params
+        )
         if r.status_code == 200:
             return True
         else:
@@ -234,7 +237,7 @@ class AppService:
 
     def get_system_user_auth_info(self, system_user):
         """获取系统用户的认证信息: 密码, ssh私钥"""
-        r, content = self.get('system-user-auth-info', pk=system_user['id'])
+        r, content = self.requests.get('system-user-auth-info', pk=system_user['id'])
         if r.status_code == 200:
             password = content['password'] or ''
             private_key_string = content['private_key'] or ''
@@ -275,7 +278,7 @@ class AppService:
         data['date_start'] = timestamp_to_datetime_str(data['date_start'])
         data['is_failed'] = 1 if data.get('is_failed') else 0
 
-        r, content = self.post('send-proxy-log', data=data, use_auth=True)
+        r, content = self.requests.post('send-proxy-log', data=data, use_auth=True)
         if r.status_code != 201:
             logging.warning('Send proxy log failed: %s' % content)
             return None
@@ -296,7 +299,7 @@ class AppService:
         data['is_failed'] = 1 if data.get('is_failed') else 0
         data['is_finished'] = 1
         proxy_log_id = data.get('proxy_log_id') or 0
-        r, content = self.patch('finish-proxy-log', pk=proxy_log_id, data=data)
+        r, content = self.requests.patch('finish-proxy-log', pk=proxy_log_id, data=data)
 
         if r.status_code != 200:
             logging.warning('Finish proxy log failed: %s' % proxy_log_id)
@@ -327,7 +330,7 @@ class AppService:
             output = d['output'].encode('utf-8', 'ignore')
             d['output'] = base64.b64encode(output).decode("utf-8")
 
-        result, content = self.post('send-command-log', data=data)
+        result, content = self.requests.post('send-command-log', data=data)
         if result.status_code != 201:
             logging.warning('Send command log failed: %s' % content)
             return False
@@ -349,7 +352,7 @@ class AppService:
             if d.get('output') and isinstance(d['output'], str):
                 d['output'] = d['output'].encode('utf-8')
             d['output'] = base64.b64encode(d['output'])
-        result, content = self.post('send-record-log', data=data)
+        result, content = self.requests.post('send-record-log', data=data)
         if result.status_code != 201:
             logging.warning('Send record log failed: %s' % content)
             return False
@@ -371,27 +374,6 @@ class AppService:
     #     user = user_service.is_authenticated()
     #     return user
 
-
-    def login(self, data):
-        """用户登录Terminal时需要向Jumpserver进行认证, 登陆成功后返回用户和token
-        data = {
-            'username': 'admin',
-            'password': 'admin',
-            'public_key': 'public key string',
-            'login_type': 'ST',  # (('ST', 'SSH Terminal'),
-                                 #  ('WT', 'Web Terminal'))
-            'remote_addr': '2.2.2.2',  # User ip address not app address
-        }
-        """
-        r, content = self.post('user-auth', data=data, use_auth=False)
-        if r.status_code == 200:
-            self.token = content['token']
-            self.user = content['user']
-            self.auth(self.token)
-            return self.user,  self.token
-        else:
-            return None, None
-
     @cached(TTLCache(maxsize=100, ttl=60))
     def get_user_assets(self, user):
         """获取用户被授权的资产列表
@@ -399,7 +381,7 @@ class AppService:
          'system_users_granted': [{'id': 1, 'username': 'x',..}]
         ]
         """
-        r, content = self.get('user-assets', pk=user['id'], use_auth=True)
+        r, content = self.requests.get('user-assets', pk=user['id'], use_auth=True)
         if r.status_code == 200:
             assets = content
         else:
@@ -416,7 +398,7 @@ class AppService:
         """获取用户授权的资产组列表
         [{'name': 'x', 'comment': 'x', 'assets_amount': 2}, ..]
         """
-        r, content = self.get('user-asset-groups', pk=user['id'], uassetsse_auth=True)
+        r, content = self.requests.get('user-asset-groups', pk=user['id'], uassetsse_auth=True)
         if r.status_code == 200:
             asset_groups = content
         else:
@@ -429,7 +411,7 @@ class AppService:
         """获取用户授权的资产组列表及下面的资产
         [{'name': 'x', 'comment': 'x', 'assets': []}, ..]
         """
-        r, content = self.get('user-asset-groups-assets', pk=user['id'], use_auth=True)
+        r, content = self.requests.get('user-asset-groups-assets', pk=user['id'], use_auth=True)
         if r.status_code == 200:
             asset_groups_assets = content
         else:
@@ -443,7 +425,7 @@ class AppService:
 
         :param asset_group_id: 资产组id
         """
-        r, content = self.get('assets-of-group', use_auth=True,
+        r, content = self.requests.get('assets-of-group', use_auth=True,
                               pk=asset_group_id)
         if r.status_code == 200:
             assets = content

coco/utils.py

@@ -91,7 +91,8 @@ def content_md5(data):
     """
     if isinstance(data, str):
         data = hashlib.md5(data.encode('utf-8'))
-    return base64.b64encode(data.digest())
+    value = base64.b64encode(data.digest())
+    return value.decode('utf-8')
 
 
 _STRPTIME_LOCK = threading.Lock()
@@ -126,7 +127,7 @@ def iso8601_to_unixtime(time_string):
 
 def make_signature(access_key_secret, date=None):
     if isinstance(date, bytes):
-        date = date.decode("utf-8")
+        date = bytes.decode(date)
     if isinstance(date, int):
         date_gmt = http_date(date)
     elif date is None: