Merge with zhuima

coco/auth.py

+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+#
+
+import os
+import six
+import logging
+from io import IOBase
+
+from . import utils
+from .exception import LoadAccessKeyError
+
+
+class AccessKeyAuth(object):
+    def __init__(self, access_key_id, access_key_secret):
+        self.id = access_key_id
+        self.secret = access_key_secret
+
+    def sign_request(self, req):
+        req.headers['Date'] = utils.http_date()
+        signature = utils.make_signature(self.secret)
+        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 OSError('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):
+        self.id = id
+        self.secret = secret
+
+    def clean(self, value, delimiter=':', silent=False):
+        try:
+            self.id, self.secret = value.split(delimiter)
+        except (AttributeError, ValueError) as e:
+            if not silent:
+                raise LoadAccessKeyError(e)
+        else:
+            return ':'.join([self.id, self.secret])
+
+    def load_from_env(self, env, delimiter=':', silent=False):
+        value = os.environ.get(env)
+        return self.clean(value, delimiter, silent)
+
+    def load_from_f(self, f, delimiter=':', silent=False):
+        value = ''
+        if isinstance(f, six.string_types) and os.path.isfile(f):
+            f = open(f)
+        if hasattr(f, 'read'):
+            for line in f:
+                if line and not line.strip().startswith('#'):
+                    value = line.strip()
+                    break
+            f.close()
+        return self.clean(value, delimiter, silent)
+
+    def save_to_f(self, f, silent=False):
+        if isinstance(f, six.string_types):
+            f = open(f, 'w')
+        try:
+            f.write(str('{0}:{1}'.format(self.id, self.secret)))
+        except IOError as e:
+            logging.error('Save access key error: {}'.format(e))
+            if not silent:
+                raise
+        finally:
+            f.close()
+
+    def __nonzero__(self):
+        return bool(self.id and self.secret)
+    __bool__ = __nonzero__
+
+    def __str__(self):
+        return '{0}:{1}'.format(self.id, self.secret)
+
+    __repr__ = __str__
+
+
+class ServiceAccessKey(AccessKey):
+    """使用Access key来认证"""
+
+    # 默认从配置文件中读取的设置
+    # 配置文件中ACCESS_KEY值的名称
+    conf_attr_var = 'ACCESS_KEY'
+    # 配置文件中配置环境变量的名称
+    conf_env_var = 'ACCESS_KEY_ENV'
+    # 配置文件中定义Access key store的位置
+    conf_store_var = 'ACCESS_KEY_STORE'
+
+    # 如果用户配置中没有设置, 方法中也没填入, 使用下面默认
+    default_key_env = 'ACCESS_KEY_ENV'
+    default_key_store = os.path.join(os.environ.get('HOME', ''), '.access_key')
+
+    def __init__(self, id=None, secret=None, config=None):
+        super(ServiceAccessKey, self).__init__(id=id, secret=secret)
+        self.config = config or {}
+        self._key_store = None
+        self._key_env = None
+
+    # 获取key store位置
+    @property
+    def key_store(self):
+        if self._key_store:
+            return self._key_store
+        elif self.conf_store_var in self.config:
+            return self.config[self.conf_store_var]
+        else:
+            return self.default_key_store
+
+    @key_store.setter
+    def key_store(self, value):
+        self._key_store = value
+
+    # 获取access key的环境变量名
+    @property
+    def key_env(self):
+        if self._key_env:
+            return self._key_env
+        elif self.conf_env_var in self.config:
+            return self.config[self.conf_env_var]
+        else:
+            return self.default_key_env
+
+    @key_env.setter
+    def key_env(self, value):
+        self._key_env = value
+
+    def load_from_conf_env(self, env=None, delimiter=':'):
+        if env is None:
+            env = self.key_env
+        return super(ServiceAccessKey, self).\
+            load_from_env(env, delimiter=delimiter)
+
+    def load_from_conf_setting(self, attr=None, delimiter=':', silent=False):
+        value = ''
+        if attr is None:
+            attr = self.conf_attr_var
+        if attr in self.config:
+            value = self.config.get(attr)
+        return self.clean(value, delimiter, silent)
+
+    def load_from_key_store(self, f=None, delimiter=':', silent=False):
+        if f is None:
+            f = self.key_store
+        return super(ServiceAccessKey, self).load_from_f(f, delimiter, silent)
+
+    def load_from_conf_all(self, **kwargs):
+        """Should return access_key_id, access_key_secret"""
+        for method in [self.load_from_conf_setting,
+                       self.load_from_key_store,
+                       self.load_from_conf_env]:
+            try:
+                return method(**kwargs)
+            except LoadAccessKeyError:
+                continue
+
+        if not (bool(self.id) and bool(self.secret)):
+            logging.error('Load access key failed')
+
+    def save_to_key_store(self, key_store=None, silent=True):
+        if key_store is None:
+            key_store = self.key_store
+        return super(ServiceAccessKey, self).save_to_f(key_store, silent)
\ No newline at end of file

coco/compat.py

+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+#
+
+"""
+兼容Python版本
+"""
+
+import sys
+
+is_py2 = (sys.version_info[0] == 2)
+is_py3 = (sys.version_info[0] == 3)
+
+
+try:
+    import simplejson as json
+except (ImportError, SyntaxError):
+    import json
+
+
+if is_py2:
+
+    def to_bytes(data):
+        """若输入为unicode， 则转为utf-8编码的bytes；其他则原样返回。"""
+        if isinstance(data, unicode):
+            return data.encode('utf-8')
+        else:
+            return data
+
+    def to_string(data):
+        """把输入转换为str对象"""
+        return to_bytes(data)
+
+    def to_unicode(data):
+        """把输入转换为unicode，要求输入是unicode或者utf-8编码的bytes。"""
+        if isinstance(data, bytes):
+            return data.decode('utf-8')
+        else:
+            return data
+
+    def stringify(input):
+        if isinstance(input, dict):
+            return dict([(stringify(key), stringify(value)) for key,value in input.iteritems()])
+        elif isinstance(input, list):
+            return [stringify(element) for element in input]
+        elif isinstance(input, unicode):
+            return input.encode('utf-8')
+        else:
+            return input
+
+    builtin_str = str
+    bytes = str
+    str = unicode
+
+
+elif is_py3:
+
+    def to_bytes(data):
+        """若输入为str（即unicode），则转为utf-8编码的bytes；其他则原样返回"""
+        if isinstance(data, str):
+            return data.encode(encoding='utf-8')
+        else:
+            return data
+
+    def to_string(data):
+        """若输入为bytes，则认为是utf-8编码，并返回str"""
+        if isinstance(data, bytes):
+            return data.decode('utf-8')
+        else:
+            return data
+
+    def to_unicode(data):
+        """把输入转换为unicode，要求输入是unicode或者utf-8编码的bytes。"""
+        return to_string(data)
+
+    def stringify(input):
+        return input
+
+    builtin_str = str
+    bytes = bytes
+    str = str
\ No newline at end of file

coco/config.py

@@ -261,3 +261,22 @@ class Config(dict):
 
     def __repr__(self):
         return '<%s %s>' % (self.__class__.__name__, dict.__repr__(self))
+
+
+API_URL_MAPPING = {
+    'terminal-register': '/api/applications/v1/terminal/register/',
+    'terminal-heatbeat': '/api/applications/v1/terminal/heatbeat/',
+    'send-proxy-log': '/api/audits/v1/proxy-log/receive/',
+    'finish-proxy-log': '/api/audits/v1/proxy-log/%s/',
+    'send-command-log': '/api/audits/v1/command-log/',
+    'send-record-log': '/api/audits/v1/record-log/',
+    'user-auth': '/api/users/v1/auth/',
+    'user-assets': '/api/perms/v1/user/%s/assets/',
+    'user-asset-groups': '/api/perms/v1/user/%s/asset-groups/',
+    'user-asset-groups-assets': '/api/perms/v1/user/my/asset-groups-assets/',
+    'assets-of-group': '/api/perms/v1/user/my/asset-group/%s/assets/',
+    'my-profile': '/api/users/v1/profile/',
+    'system-user-auth-info': '/api/assets/v1/system-user/%s/auth-info/',
+    'validate-user-asset-permission':
+        '/api/perms/v1/asset-permission/user/validate/',
+}

coco/exception.py

@@ -4,3 +4,11 @@
 class PermissionFailed(Exception):
     pass
 
+
+class LoadAccessKeyError(Exception):
+    pass
+
+
+class RequestError(Exception):
+    pass
+

coco/utils.py

 #!coding: utf-8
-import base64
-import calendar
-import os
-import re
 
-import paramiko
-from io import StringIO
+from __future__ import unicode_literals
+
 import hashlib
+import re
 import threading
-
+import base64
+import calendar
 import time
+import datetime
+from io import StringIO
 
 import pyte
+import pytz
+from email.utils import formatdate
+
+import paramiko
+from dotmap import DotMap
+
+
+try:
+    from Queue import Queue, Empty
+except ImportError:
+    from queue import Queue, Empty
+
+from .compat import to_string, to_bytes
+
 
 
 def ssh_key_string_to_obj(text):
@@ -238,3 +252,109 @@ def wrap_with_primary(text, bolder=False):
 
 def wrap_with_title(text):
     return wrap_with_color(text, color='black', background='green')
+
+
+def b64encode_as_string(data):
+    return to_string(base64.b64encode(data))
+
+
+
+def make_signature(access_key_secret, date=None):
+    if isinstance(date, bytes):
+        date = date.decode("utf-8")
+    if isinstance(date, int):
+        date_gmt = http_date(date)
+    elif date is None:
+        date_gmt = http_date(int(time.time()))
+    else:
+        date_gmt = date
+
+    data = str(access_key_secret) + "\n" + date_gmt
+    return content_md5(data)
+
+
+def split_string_int(s):
+    """Split string or int
+
+    example: test-01-02-db => ['test-', '01', '-', '02', 'db']
+    """
+    string_list = []
+    index = 0
+    pre_type = None
+    word = ''
+    for i in s:
+        if index == 0:
+            pre_type = int if i.isdigit() else str
+            word = i
+        else:
+            if pre_type is int and i.isdigit() or pre_type is str and not i.isdigit():
+                word += i
+            else:
+                string_list.append(word.lower() if not word.isdigit() else int(word))
+                word = i
+                pre_type = int if i.isdigit() else str
+        index += 1
+    string_list.append(word.lower() if not word.isdigit() else int(word))
+    return string_list
+
+
+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()])
+    else:
+        key = lambda asset: asset.__getitem__(order_by)
+        assets = sorted(assets, key=key)
+    return assets
+
+
+class PKey(object):
+    @classmethod
+    def from_string(cls, key_string):
+        try:
+            pkey = paramiko.RSAKey(file_obj=StringIO(key_string))
+            return pkey
+        except paramiko.SSHException:
+            try:
+                pkey = paramiko.DSSKey(file_obj=StringIO(key_string))
+                return pkey
+            except paramiko.SSHException:
+                return None
+
+
+def from_string(cls, key_string):
+    return cls(key_string=key_string).pkey
+
+
+def timestamp_to_datetime_str(ts):
+    datetime_format = '%Y-%m-%dT%H:%M:%S.%fZ'
+    dt = datetime.datetime.fromtimestamp(ts, tz=pytz.timezone('UTC'))
+    return dt.strftime(datetime_format)
+
+
+def to_dotmap(data):
+    """将接受dict转换为DotMap"""
+    if isinstance(data, dict):
+        data = DotMap(data)
+    elif isinstance(data, list):
+        data = [DotMap(d) for d in data]
+    else:
+        raise ValueError('Dict or list type required...')
+    return data
+
+
+class MultiQueue(Queue):
+    def mget(self, size=1, block=True, timeout=5):
+        items = []
+        for i in range(size):
+            try:
+                items.append(self.get(block=block, timeout=timeout))
+            except Empty:
+                break
+        return items
+
+
+