From ad16d8b532e7c5271c8417b754b6b7b43faf2de7 Mon Sep 17 00:00:00 2001
From: ibuler <ibuler@qq.com>
Date: Fri, 21 Aug 2015 23:45:41 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96connect.py=20=E5=92=8C=20api.?=
 =?UTF-8?q?py?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 connect.py        |  70 ++++----
 jlog/models.py    |   2 +-
 jumpserver/api.py | 395 ++++++----------------------------------------
 juser/models.py   |   5 +-
 4 files changed, 90 insertions(+), 382 deletions(-)

diff --git a/connect.py b/connect.py
index ddfd50cfc..fe83afcd4 100644
--- a/connect.py
+++ b/connect.py
@@ -26,8 +26,9 @@ os.environ['DJANGO_SETTINGS_MODULE'] = 'jumpserver.settings'
 if django.get_version() != '1.6':
     django.setup()
 from jlog.models import Log
-from jumpserver.api import CONF, BASE_DIR, ServerError, User, UserGroup, Asset, BisGroup
+from jumpserver.api import CONF, BASE_DIR, ServerError, User, UserGroup, Asset, get_object
 from jumpserver.api import CRYPTOR, logger, is_dir
+from jumpserver.api import BisGroup as AssetGroup
 
 try:
     import termios
@@ -38,7 +39,7 @@ except ImportError:
     sys.exit()
 
 log_dir = os.path.join(BASE_DIR, 'logs')
-login_user = User(username=getpass.getuser())
+login_user = get_object(User, username=getpass.getuser())
 
 
 def color_print(msg, color='red', exits=False):
@@ -98,14 +99,14 @@ class Jtty(object):
         timestamp_start = int(time.time())
         date_start = time.strftime('%Y%m%d', time.localtime(timestamp_start))
         time_start = time.strftime('%H%M%S', time.localtime(timestamp_start))
-        today_connect_log_dir = os.path.join(tty_log_dir, date_start)
         log_filename = '%s_%s_%s.log' % (self.username, self.ip, time_start)
+        today_connect_log_dir = os.path.join(tty_log_dir, date_start)
         log_file_path = os.path.join(today_connect_log_dir, log_filename)
         dept_name = self.user.dept.name
 
         pid = os.getpid()
         pts = os.popen("ps axu | grep %s | grep -v grep | awk '{ print $7 }'" % pid).read().strip()
-        remote_ip = os.popen("who | grep %s | awk '{ print $5 }'" % pts).read().strip('()\n')
+        ip_list = os.popen("who | grep %s | awk '{ print $5 }'" % pts).read().strip('()\n')
 
         try:
             is_dir(today_connect_log_dir)
@@ -117,7 +118,7 @@ class Jtty(object):
         except IOError:
             raise ServerError('Create logfile failed, Please modify %s permission.' % today_connect_log_dir)
 
-        log = Log(user=self.username, host=self.ip, remote_ip=remote_ip, dept_name=dept_name,
+        log = Log(user=self.username, host=self.ip, remote_ip=ip_list, dept_name=dept_name,
                   log_path=log_file_path, start_time=datetime.datetime.now(), pid=pid)
         log_file.write('Start time is %s\n' % datetime.datetime.now())
         log.save()
@@ -164,18 +165,17 @@ class Jtty(object):
             log_file.write('End time is %s' % datetime.datetime.now())
             log_file.close()
             log.is_finished = True
-            log.log_finished = False
+            log.handle_finished = False
             log.end_time = datetime.datetime.now()
             log.save()
 
     def get_connect_item(self):
-        port = int(self.asset.port)
-
+        """获取连接需要的参数，也就是服务ip, 端口, 用户账号和密码"""
         if not self.asset.is_active:
-            raise ServerError('Host %s is not active.' % self.ip)
+            raise ServerError('该主机被禁用 Host %s is not active.' % self.ip)
 
         if not self.user.is_active:
-            raise ServerError('User %s is not active.' % self.username)
+            raise ServerError('该用户被禁用 User %s is not active.' % self.username)
 
         login_type_dict = {
             'L': self.user.ldap_pwd,
@@ -183,37 +183,38 @@ class Jtty(object):
 
         if self.asset.login_type in login_type_dict:
             password = CRYPTOR.decrypt(login_type_dict[self.asset.login_type])
-            return self.username, password, self.ip, port
+            return self.username, password, self.ip, int(self.asset.port)
 
         elif self.asset.login_type == 'M':
             username = self.asset.username
             password = CRYPTOR.decrypt(self.asset.password)
-            return username, password, self.ip, port
+            return username, password, self.ip, int(self.asset.port)
 
         else:
-            raise ServerError('Login type is not in ["L", "M"]')
+            raise ServerError('不支持的服务器登录方式 Login type is not in ["L", "M"]')
 
     def connect(self):
         """
         Connect server.
+        连接服务器
         """
         username, password, ip, port = self.get_connect_item()
         logger.debug("username: %s, password: %s, ip: %s, port: %s" % (username, password, ip, port))
         ps1 = "PS1='[\u@%s \W]\$ '\n" % self.ip
-        login_msg = "clear;echo -e '\\033[32mLogin %s done. Enjoy it.\\033[0m'\n" % self.ip
+        login_msg = "clear;echo -e '\\033[32mLogin %s done. Enjoy it.\\033[0m'\n" % ip
 
-        # Make a ssh connection
+        # 发起ssh连接请求 Make a ssh connection
         ssh = paramiko.SSHClient()
         ssh.load_system_host_keys()
         ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
         try:
             ssh.connect(ip, port=port, username=username, password=password, compress=True)
         except paramiko.ssh_exception.AuthenticationException, paramiko.ssh_exception.SSHException:
-            raise ServerError('Authentication Error.')
+            raise ServerError('认证错误 Authentication Error.')
         except socket.error:
-            raise ServerError('Connect SSH Socket Port Error, Please Correct it.')
+            raise ServerError('端口可能不对 Connect SSH Socket Port Error, Please Correct it.')
 
-        # Make a channel and set windows size
+        # 获取连接的隧道并设置窗口大小 Make a channel and set windows size
         global channel
         win_size = self.get_win_size()
         self.chan = channel = ssh.invoke_shell(height=win_size[0], width=win_size[1])
@@ -222,7 +223,7 @@ class Jtty(object):
         except:
             pass
 
-        # Set PS1 and msg it
+        # 设置PS1并提示 Set PS1 and msg it
         channel.send(ps1)
         channel.send(login_msg)
 
@@ -235,6 +236,7 @@ class Jtty(object):
 
 
 def verify_connect(user, option):
+    """鉴定用户是否有该主机权限 或 匹配到的ip是否唯一"""
     ip_matched = []
     try:
         assets_info = login_user.get_asset_info()
@@ -255,16 +257,20 @@ def verify_connect(user, option):
     ip_matched = list(set(ip_matched))
 
     if len(ip_matched) > 1:
+        ip_comment = {}
         for ip in ip_matched:
-            if assets_info[ip][2]:
-                print '%-15s -- %s' % (ip, assets_info[ip][2])
+            ip_comment[ip] = assets_info[ip][2]
+
+        for ip in sorted(ip_comment):
+            if ip_comment[ip]:
+                print '%-15s -- %s' % (ip, ip_comment[ip])
             else:
                 print '%-15s' % ip
-            print ''
+        print ''
     elif len(ip_matched) < 1:
-        color_print('No Permission or No host.', 'red')
+        color_print('没有该主机，或者您没有该主机的权限 No Permission or No host.', 'red')
     else:
-        asset = Asset(ip=ip_matched[0]).asset
+        asset = get_object(Asset, ip=ip_matched[0])
         jtty = Jtty(user, asset)
         jtty.connect()
 
@@ -351,12 +357,13 @@ def print_prompt():
 #             multi_remote_exec_cmd(hosts, username, cmd)
 
 
-if __name__ == '__main__':
-    if not login_user.validate():
-        color_print(u'没有该用户 No that user.', exits=True)
+def main():
+    if not login_user:  # 判断用户是否存在
+        color_print(u'没有该用户，或许你是以root运行的 No that user.', exits=True)
 
     print_prompt()
     gid_pattern = re.compile(r'^g\d+$')
+
     try:
         while True:
             try:
@@ -374,8 +381,8 @@ if __name__ == '__main__':
                 continue
             elif gid_pattern.match(option):
                 gid = option[1:].strip()
-                asset_group = JassetGroup(id=gid)
-                if asset_group.validate() and asset_group.is_permed(user=login_user.user):
+                asset_group = get_object(AssetGroup, id=gid)
+                if asset_group and asset_group.is_permed(user=login_user):
                     asset_group.get_asset_info(printable=True)
                 continue
             elif option in ['E', 'e']:
@@ -390,3 +397,8 @@ if __name__ == '__main__':
                     color_print(e, 'red')
     except IndexError:
         pass
+
+if __name__ == '__main__':
+    main()
+
+
diff --git a/jlog/models.py b/jlog/models.py
index 3d2a319df..baaffb5a7 100644
--- a/jlog/models.py
+++ b/jlog/models.py
@@ -10,7 +10,7 @@ class Log(models.Model):
     start_time = models.DateTimeField(null=True)
     pid = models.IntegerField(max_length=10)
     is_finished = models.BooleanField(default=False)
-    log_finished = models.BooleanField(default=False)
+    handle_finished = models.BooleanField(default=False)
     end_time = models.DateTimeField(null=True)
 
     def __unicode__(self):
diff --git a/jumpserver/api.py b/jumpserver/api.py
index 811759f6d..2b133585a 100644
--- a/jumpserver/api.py
+++ b/jumpserver/api.py
@@ -41,16 +41,18 @@ SEND_IP = CONF.get('base', 'ip')
 SEND_PORT = CONF.get('base', 'port')
 MAIL_FROM = CONF.get('mail', 'email_host_user')
 
-log_level = CONF.get('base', 'log')
-log_level_total = {'debug': logging.DEBUG, 'info': logging.INFO, 'warning': logging.WARN, 'error': logging.ERROR,
-                   'critical': logging.CRITICAL}
-logger = logging.getLogger('jumpserver')
-logger.setLevel(logging.DEBUG)
-fh = logging.FileHandler(JLOG_FILE)
-fh.setLevel(log_level_total.get(log_level, logging.DEBUG))
-formatter = logging.Formatter('%(asctime)s - %(filename)s - %(levelname)s - %(message)s')
-fh.setFormatter(formatter)
-logger.addHandler(fh)
+
+def set_log(level):
+    log_level_total = {'debug': logging.DEBUG, 'info': logging.INFO, 'warning': logging.WARN, 'error': logging.ERROR,
+                       'critical': logging.CRITICAL}
+    logger = logging.getLogger('jumpserver')
+    logger.setLevel(logging.DEBUG)
+    fh = logging.FileHandler(JLOG_FILE)
+    fh.setLevel(log_level_total.get(level, logging.DEBUG))
+    formatter = logging.Formatter('%(asctime)s - %(filename)s - %(levelname)s - %(message)s')
+    fh.setFormatter(formatter)
+    logger.addHandler(fh)
+    return logger
 
 
 class LDAPMgmt():
@@ -205,68 +207,45 @@ def get_object(model, **kwargs):
     try:
         the_object = model.objects.get(**kwargs)
     except ObjectDoesNotExist:
-        raise ServerError('Object get %s failed.' % str(kwargs.values()))
+        the_object = None
     return the_object
 
 
-def require_login(func):
-    """要求登录的装饰器"""
-    def _deco(request, *args, **kwargs):
-        if not request.session.get('user_id'):
-            return HttpResponseRedirect('/login/')
-        return func(request, *args, **kwargs)
+def require_role(role='user'):
+    """
+    要求用户是某种角色 ["super", "admin", "user"]
+    """
+    def _deco(func):
+        def __deco(request, *args, **kwargs):
+            if role == 'user':
+                if not request.session.get('user_id'):
+                    return HttpResponseRedirect('/login/')
+            elif role == 'admin':
+                if request.session.get('role_id', 0) != 1:
+                    return HttpResponseRedirect('/')
+            elif role == 'super':
+                if request.session.get('role_id', 0) != 2:
+                    return HttpResponseRedirect('/')
+            return func(request, *args, **kwargs)
+        return __deco()
     return _deco
 
 
-def require_super_user(func):
-    """要求是超级管理员"""
-    def _deco(request, *args, **kwargs):
-        if not request.session.get('user_id'):
-            return HttpResponseRedirect('/login/')
-
-        if request.session.get('role_id', 0) != 2:
-            return HttpResponseRedirect('/')
-        return func(request, *args, **kwargs)
-    return _deco
-
-
-def require_admin(func):
-    """要求是管理员"""
-    def _deco(request, *args, **kwargs):
-        if not request.session.get('user_id'):
-            return HttpResponseRedirect('/login/')
-
-        if request.session.get('role_id', 0) < 1:
-            return HttpResponseRedirect('/')
-        return func(request, *args, **kwargs)
-    return _deco
-
-
-def is_super_user(request):
-    """要求请求是超级管理员"""
-    if request.session.get('role_id') == 2:
+def is_role_request(request, role='user'):
+    """
+    :param request: 请求
+    :param role: 角色
+    :return: bool
+    要求请求角色正确
+    """
+    role_all = {'user': 0, 'admin': 1, 'super': 2}
+    if request.session.get('role_id') == role_all.get(role, '0'):
         return True
     else:
         return False
 
 
-def is_group_admin(request):
-    """要求请求是组管理员"""
-    if request.session.get('role_id') == 1:
-        return True
-    else:
-        return False
-
-
-def is_common_user(request):
-    """要求用户是普通用户"""
-    if request.session.get('role_id') == 0:
-        return True
-    else:
-        return False
-
-
-@require_login
+@require_role
 def get_session_user_dept(request):
     user_id = request.session.get('user_id', 0)
     user = User.objects.filter(id=user_id)
@@ -276,7 +255,7 @@ def get_session_user_dept(request):
         return user, dept
 
 
-@require_login
+@require_role
 def get_session_user_info(request):
     user_id = request.session.get('user_id', 0)
     user = User.objects.filter(id=user_id)
@@ -310,295 +289,6 @@ def api_user(request):
 #         return HttpResponseRedirect('/login/')
 
 
-# def user_group_perm_asset_group_api(user_group):
-#     asset_group_list = []
-#     perm_list = user_group.perm_set.all()
-#     for perm in perm_list:
-#         asset_group_list.append(perm.asset_group)
-#     return asset_group_list
-
-
-# class Juser(object):
-#     """
-#     Jumpserver user class
-#     用户类
-#     """
-#
-#     def __init__(self, username=None, uid=None):
-#         if username:
-#             user = User.objects.filter(username=username)
-#         elif uid:
-#             user = User.objects.filter(id=uid)
-#         else:
-#             user = ''
-#
-#         if user:
-#             user = user[0]
-#             self.user = user
-#             self.id = user.id
-#             # self.id = user.id
-#             # self.username = user.username
-#             # self.name = user.name
-#             self.group = user.group.all()
-#         else:
-#             self.id = None
-#
-#     def __repr__(self):
-#         if self.id:
-#             return '<%s Juser instance>' % getattr(self.user, 'username')
-#         else:
-#             return 'None'
-#
-#     def __getattr__(self, item):
-#         if self.id:
-#             return getattr(self.user, item)
-#         else:
-#             return None
-#
-#     def validate(self):
-#         """
-#         Validate is or not a true user
-#         鉴定用户
-#         """
-#         if self.id:
-#             return True
-#         else:
-#             return False
-#
-#     def get_asset_group(self):
-#         """
-#         Get user host_groups.
-#         获取用户有权限的主机组
-#         """
-#         host_group_list = []
-#         perm_list = []
-#         user_group_all = self.user.group.all()
-#         for user_group in user_group_all:
-#             perm_list.extend(user_group.perm_set.all())
-#
-#         for perm in perm_list:
-#             host_group_list.append(perm.asset_group)
-#
-#         return host_group_list
-#
-#     def get_asset_group_info(self, printable=False):
-#         """
-#         Get or print asset group info
-#         获取或打印用户授权资产组
-#         """
-#         asset_groups_info = {}
-#         asset_groups = self.get_asset_group()
-#
-#         for asset_group in asset_groups:
-#             asset_groups_info[asset_group.id] = [asset_group.name, asset_group.comment]
-#
-#         if printable:
-#             for group_id in asset_groups_info:
-#                 if asset_groups_info[group_id][1]:
-#                     print "[%3s] %s -- %s" % (group_id,
-#                                               asset_groups_info[group_id][0],
-#                                               asset_groups_info[group_id][1])
-#                 else:
-#                     print "[%3s] %s" % (group_id, asset_groups_info[group_id][0])
-#                 print ''
-#         else:
-#             return asset_groups_info
-#
-#     def get_asset(self):
-#         """
-#         Get the assets of under the user control.
-#         获取主机列表
-#         """
-#         assets = []
-#         asset_groups = self.get_asset_group()
-#
-#         for asset_group in asset_groups:
-#             assets.extend(asset_group.asset_set.all())
-#
-#         return assets
-#
-#     def get_asset_info(self, printable=False):
-#         """
-#         Get or print the user asset info
-#         获取或打印用户资产信息
-#         """
-#         assets_info = {}
-#         assets = self.get_asset()
-#
-#         for asset in assets:
-#             asset_alias = AssetAlias.objects.filter(user=self.user, asset=asset)
-#             if asset_alias and asset_alias[0].alias != '':
-#                 assets_info[asset.ip] = [asset.id, asset.ip, str(asset_alias[0].alias)]
-#             else:
-#                 assets_info[asset.ip] = [asset.id, asset.ip, str(asset.comment)]
-#
-#         if printable:
-#             ips = assets_info.keys()
-#             ips.sort()
-#             for ip in ips:
-#                 if assets_info[ip][2]:
-#                     print '%-15s -- %s' % (ip, assets_info[ip][2])
-#                 else:
-#                     print '%-15s' % ip
-#             print ''
-#         else:
-#             return assets_info
-#
-
-# class Jasset(object):
-#     """
-#     Jumpserver asset class
-#     Jumpserver资产类
-#     """
-#     def __init__(self, ip=None, id=None):
-#         if ip:
-#             asset = Asset.objects.filter(ip=ip)
-#         elif id:
-#             asset = Asset.objects.filter(id=id)
-#         else:
-#             asset = ''
-#
-#         if asset:
-#             asset = asset[0]
-#             self.asset = asset
-#             self.id = asset.id
-#         else:
-#             self.id = None
-#
-#     def __repr__(self):
-#         if self.id:
-#             return '<%s Jasset instance>' % self.asset.ip
-#         else:
-#             return 'None'
-#
-#     def __getattr__(self, item):
-#         if self.id:
-#             return getattr(self.asset, item)
-#         else:
-#             return None
-#
-#     def validate(self):
-#         """
-#         Validate is or not a true asset
-#         判断是否存在
-#         """
-#         if self.id:
-#             return True
-#         else:
-#             return False
-#
-#     def get_user(self):
-#         perm_list = []
-#         asset_group_all = self.bis_group.all()
-#         for asset_group in asset_group_all:
-#             perm_list.extend(asset_group.perm_set.all())
-#
-#         user_group_list = []
-#         for perm in perm_list:
-#             user_group_list.append(perm.user_group)
-#
-#         user_permed_list = []
-#         for user_group in user_group_list:
-#             user_permed_list.extend(user_group.user_set.all())
-#         user_permed_list = list(set(user_permed_list))
-#         return user_permed_list
-
-
-# class JassetGroup(object):
-#     """
-#     Jumpserver AssetGroup class
-#     Jumpserver 资产组类
-#     """
-#     def __init__(self, name=None, id=None):
-#         if id:
-#             asset_group = BisGroup.objects.filter(id=int(id))
-#         elif name:
-#             asset_group = BisGroup.objects.filter(name=name)
-#         else:
-#             asset_group = ''
-#
-#         if asset_group:
-#             asset_group = asset_group[0]
-#             self.asset_group = asset_group
-#             # self.name = asset_group.name
-#             self.id = asset_group.id
-#         else:
-#             self.id = None
-#
-#     def __repr__(self):
-#         if self.id:
-#             return '<%s JassetGroup instance>' % self.name
-#         else:
-#             return 'None'
-#
-#     def validate(self):
-#         """
-#         Validate it is a true asset group or not
-#         鉴定是否为真是存在的组
-#         """
-#         if self.id:
-#             return True
-#         else:
-#             return False
-#
-#     def get_asset(self):
-#         return self.asset_group.asset_set.all()
-#
-#     def get_asset_info(self, printable=False):
-#         assets = self.get_asset()
-#         for asset in assets:
-#             if asset.comment:
-#                 print '%-15s -- %s' % (asset.ip, asset.comment)
-#             else:
-#                 print '%-15s' % asset.ip
-#         print ''
-#
-#     def get_asset_num(self):
-#         return len(self.get_asset())
-#
-#     def get_user_group(self):
-#         perm_list = self.asset_group.perm_set.all()
-#         user_group_list = []
-#         for perm in perm_list:
-#             user_group_list.append(perm.user_group)
-#         return user_group_list
-#
-#     def get_user(self):
-#         user_list = []
-#         user_group_list = self.get_user_group()
-#         for user_group in user_group_list:
-#             user_list.extend(user_group.user_set.all())
-#         return user_list
-#
-#     def is_permed(self, user=None, user_group=None):
-#         if user:
-#             if user in self.get_user():
-#                 return True
-#
-#         if user_group:
-#             if user_group in self.get_user_group():
-#                 return True
-#         return False
-
-
-# def asset_perm_api(asset):
-#     if asset:
-#         perm_list = []
-#         asset_group_all = asset.bis_group.all()
-#         for asset_group in asset_group_all:
-#             perm_list.extend(asset_group.perm_set.all())
-#
-#         user_group_list = []
-#         for perm in perm_list:
-#             user_group_list.append(perm.user_group)
-#
-#         user_permed_list = []
-#         for user_group in user_group_list:
-#             user_permed_list.extend(user_group.user_set.all())
-#         user_permed_list = list(set(user_permed_list))
-#         return user_permed_list
-
-
 def validate(request, user_group=None, user=None, asset_group=None, asset=None, edept=None):
     dept = get_session_user_dept(request)[1]
     if edept:
@@ -694,6 +384,7 @@ def bash(cmd):
 
 
 def is_dir(dir_name, username='root', mode=0755):
+    """目录存在，如果不存在就建立，并且权限正确"""
     if not os.path.isdir(dir_name):
         os.makedirs(dir_name)
         bash("chown %s:%s '%s'" % (username, username, dir_name))
@@ -720,3 +411,5 @@ if LDAP_ENABLE:
 else:
     ldap_conn = None
 
+log_level = CONF.get('base', 'log')
+logger = set_log(log_level)
\ No newline at end of file
diff --git a/juser/models.py b/juser/models.py
index 31dd9b663..a507e1c63 100644
--- a/juser/models.py
+++ b/juser/models.py
@@ -19,6 +19,9 @@ class UserGroup(models.Model):
     def __unicode__(self):
         return self.name
 
+    def get_user(self):
+        return self.user_set.all()
+
 
 class User(models.Model):
     USER_ROLE_CHOICES = (
@@ -104,7 +107,7 @@ class User(models.Model):
         assets = self.get_asset()
 
         for asset in assets:
-            asset_alias = AssetAlias.objects.filter(user=self.user, asset=asset)
+            asset_alias = AssetAlias.objects.filter(user=self, asset=asset)
             if asset_alias and asset_alias[0].alias != '':
                 assets_info[asset.ip] = [asset.id, asset.ip, str(asset_alias[0].alias)]
             else: