Update sdk

coco/app.py

@@ -32,7 +32,7 @@ class Coco:
         'LOG_LEVEL': 'INFO',
         'LOG_DIR': os.path.join(BASE_DIR, 'logs'),
         'SESSION_DIR': os.path.join(BASE_DIR, 'sessions'),
-        'ASSET_SORT_BY': 'hostname',  # hostname, ip
+        'ASSET_LIST_SORT_BY': 'hostname',  # hostname, ip
         'SSH_PASSWORD_AUTH': True,
         'SSH_PUBLIC_KEY_AUTH': True,
         'HEARTBEAT_INTERVAL': 5,

coco/models.py

@@ -9,52 +9,78 @@ BUF_SIZE = 4096
 
 
 class Decoder:
-    def __init__(self, **kwargs):
-        for attr, val in kwargs.items():
-            setattr(self, attr, val)
-
     @classmethod
-    def from_json(cls, json_str):
-        json_dict = json.loads(json_str)
-        return cls(**json_dict)
+    def from_json(cls, json_dict):
+        self = cls()
+        for k, v in json_dict.items():
+            if isinstance(getattr(self, k, None), datetime.datetime):
+                v = datetime.datetime.strptime(v, "%Y-%m-%d %H:%M:%S")
+            setattr(self, k, v)
+        return self
 
     @classmethod
-    def from_multi_json(cls, json_list):
-        json_dict_list = json.loads(json_list)
-        return [cls(**json_dict) for json_dict in json_dict_list]
+    def from_multi_json(cls, json_dict_list):
+        return [cls.from_json(json_dict) for json_dict in json_dict_list]
 
 
 class User(Decoder):
-    id = ""
+    id = 0
     username = ""
     name = ""
+    email = ""
+    is_active = False
+    is_superuser = False
+    role = "User"
+    groups = []
+    wechat = ""
+    phone = ""
+    comment = ""
+    date_expired = datetime.datetime.now()
 
     def __str__(self):
         return self.name
-    __repr__ = __str__
+
+    def __repr__(self):
+        return self.name
 
 
 class Asset(Decoder):
-    id = ""
+    id = 0
     hostname = ""
     ip = ""
     port = 22
+    system_users_granted = []
+    is_active = False
+    system_users_join = ""
+
+    @classmethod
+    def from_json(cls, json_dict):
+        system_users_granted = SystemUser.from_multi_json(json_dict["system_users_granted"])
+        json_dict["system_users_granted"] = system_users_granted
+        return super().from_json(json_dict)
 
     def __str__(self):
         return self.hostname
-    __repr__ = __str__
+
+    def __repr__(self):
+        return self.hostname
 
 
 class SystemUser(Decoder):
-    id = ""
+    id = 0
     name = ""
     username = ""
+    protocol = "ssh"
+    auth_method = "P"
+    comment = ""
     password = ""
     private_key = None
 
     def __str__(self):
         return self.name
-    __repr__ = __str__
+
+    def __repr__(self):
+        return self.name
 
 
 class Request:

coco/sdk.py

@@ -16,7 +16,7 @@ from cachetools import cached, TTLCache
 from .auth import AppAccessKey, AccessKeyAuth
 from .utils import sort_assets, PKey, timestamp_to_datetime_str
 from .exception import RequestError, ResponseError
-
+from .models import User, Asset
 
 _USER_AGENT = 'jms-sdk-py'
 CACHED_TTL = os.environ.get('CACHED_TTL', 30)
@@ -160,7 +160,7 @@ class AppService:
     def valid_auth(self):
         delay = 1
         while delay < 300:
-            if self.heatbeat() is None:
+            if self.heartbeat() is None:
                 msg = "Access key is not valid or need admin " \
                       "accepted, waiting %d s" % delay
                 logger.info(msg)
@@ -205,7 +205,7 @@ class AppService:
             logging.error('Register terminal {} failed unknown: {}'.format(self.app.name, resp.json()))
             sys.exit()
 
-    def heatbeat(self):
+    def heartbeat(self):
         """和Jumpserver维持心跳, 当Terminal断线后,jumpserver可以知晓
 
         Todo: Jumpserver发送的任务也随heatbeat返回, 并执行,如 断开某用户
@@ -220,6 +220,29 @@ class AppService:
         else:
             return None
 
+    def check_user_credential(self, username, password="", pubkey="",
+                              remote_addr="8.8.8.8", login_type='ST'):
+        data = {
+            'username': username,
+            'password': password,
+            'public_key': pubkey,
+            'remote_addr': remote_addr,
+            'login_type': login_type,
+        }
+        try:
+            resp = self.requests.post('user-auth', data=data, use_auth=False)
+        except (ResponseError, RequestError):
+            return None
+
+        if resp.status_code == 200:
+            user = User.from_json(resp.json()["user"])
+            return user
+        else:
+            return None
+
+    def check_user_cookie(self, session_id, csrf_token):
+        pass
+
     def validate_user_asset_permission(self, user_id, asset_id, system_user_id):
         """验证用户是否有登录该资产的权限"""
         params = {
@@ -358,22 +381,6 @@ class AppService:
             return False
         return True
 
-    # Todo: 或许没什么用
-    # def check_user_authentication(self, token=None, session_id=None,
-    #                               csrf_token=None):
-    #     """
-    #     用户登陆webterminal或其它网站时,检测用户cookie中的sessionid和csrf_token
-    #     是否合法, 如果合法返回用户,否则返回空
-    #     :param session_id: cookie中的 sessionid
-    #     :param csrf_token: cookie中的 csrftoken
-    #     :return: user object or None
-    #     """
-    #     user_service = UserService(endpoint=self.endpoint)
-    #     user_service.auth(token=token, session_id=session_id,
-    #                       csrf_token=csrf_token)
-    #     user = user_service.is_authenticated()
-    #     return user
-
     @cached(TTLCache(maxsize=100, ttl=60))
     def get_user_assets(self, user):
         """获取用户被授权的资产列表
@@ -381,25 +388,30 @@ class AppService:
          'system_users_granted': [{'id': 1, 'username': 'x',..}]
         ]
         """
-        r, content = self.requests.get('user-assets', pk=user['id'], use_auth=True)
-        if r.status_code == 200:
-            assets = content
+        try:
+            resp = self.requests.get('user-assets', pk=user.id, use_auth=True)
+        except (RequestError, ResponseError):
+            return []
+
+        if resp.status_code == 200:
+            assets = Asset.from_multi_json(resp.json())
         else:
-            assets = []
+            return []
 
-        assets = sort_assets(assets)
-        for asset in assets:
-            asset['system_users'] = \
-                [system_user for system_user in asset.get('system_users_granted')]
-        return to_dotmap(assets)
+        assets = sort_assets(assets, self.app.config["ASSET_LIST_SORT_BY"])
+        return assets
 
     @cached(TTLCache(maxsize=100, ttl=60))
     def get_user_asset_groups(self, user):
         """获取用户授权的资产组列表
         [{'name': 'x', 'comment': 'x', 'assets_amount': 2}, ..]
         """
-        r, content = self.requests.get('user-asset-groups', pk=user['id'], uassetsse_auth=True)
-        if r.status_code == 200:
+        try:
+            resp = self.requests.get('user-asset-groups', pk=user.id, use_auth=True)
+        except (ResponseError, RequestError):
+            return []
+
+        if resp.status_code == 200:
             asset_groups = content
         else:
             asset_groups = []

coco/utils.py

@@ -282,15 +282,10 @@ def split_string_int(s):
 
 
 def sort_assets(assets, order_by='hostname'):
-    if order_by == 'hostname':
-        key = lambda asset: split_string_int(asset['hostname'])
-        # print(assets)
-        # assets = sorted(assets, key=key)
-    elif order_by == 'ip':
-            assets = sorted(assets, key=lambda asset: [int(d) for d in asset['ip'].split('.') if d.isdigit()])
+    if order_by == 'ip':
+        assets = sorted(assets, key=lambda asset: [int(d) for d in asset.ip.split('.') if d.isdigit()])
     else:
-        key = lambda asset: asset.__getitem__(order_by)
-        assets = sorted(assets, key=key)
+        assets = sorted(assets, key=lambda asset: getattr(asset, order_by))
     return assets