Merge pull request #134 from kelianchun/dev

修复命令截取bug

connect.py

@@ -16,6 +16,8 @@ import readline
 import django
 import paramiko
 import errno
+import pyte
+import operator
 import struct, fcntl, signal, socket, select
 from io import open as copen
 import uuid
@@ -90,8 +92,11 @@ class Tty(object):
         self.remote_ip = ''
         self.login_type = login_type
         self.vim_flag = False
-        self.ps1_pattern = re.compile('\[.*@.*\][\$#]')
+        self.ps1_pattern = re.compile('\[.*@.*\][\$#]\s')
+        self.vim_pattern = re.compile(r'\Wvi[m]+\s.* | \Wfg\s.*', re.X)
         self.vim_data = ''
+        self.stream = pyte.ByteStream()
+        self.screen = None
 
     @staticmethod
     def is_output(strings):
@@ -101,135 +106,50 @@ class Tty(object):
                 return True
         return False
 
-    @staticmethod
+    def command_parser(self, command):
-    def remove_obstruct_char(cmd_str):
-        '''删除一些干扰的特殊符号'''
-        control_char = re.compile(r'\x07 | \x1b\[1P | \r ', re.X)
-        cmd_str = control_char.sub('',cmd_str.strip())
-        patch_char = re.compile('\x08\x1b\[C')      #删除方向左右一起的按键
-        while patch_char.search(cmd_str):
-            cmd_str = patch_char.sub('', cmd_str.rstrip())
-        return cmd_str
-
-    @staticmethod
-    def deal_backspace(match_str, result_command, pattern_str, backspace_num):
-        '''
-        处理删除确认键
-        '''
-        if backspace_num > 0:
-            if backspace_num > len(result_command):
-                result_command += pattern_str
-                result_command = result_command[0:-backspace_num]
-            else:
-                result_command = result_command[0:-backspace_num]
-                result_command += pattern_str
-        del_len = len(match_str)-3
-        if del_len > 0:
-            result_command = result_command[0:-del_len]
-        return result_command, len(match_str)
-    
-    @staticmethod
-    def deal_replace_char(match_str,result_command,backspace_num):
-        '''
-        处理替换命令
-        '''
-        str_lists = re.findall(r'(?<=\x1b\[1@)\w',match_str)
-        tmp_str =''.join(str_lists)
-        result_command_list = list(result_command)
-        if len(tmp_str) > 1:
-            result_command_list[-backspace_num:-(backspace_num-len(tmp_str))] = tmp_str
-        elif len(tmp_str) > 0:
-            if result_command_list[-backspace_num] == ' ':
-                result_command_list.insert(-backspace_num, tmp_str)
-            else:
-                result_command_list[-backspace_num] = tmp_str
-        result_command = ''.join(result_command_list)
-        return result_command, len(match_str)
-    
-    def remove_control_char(self, result_command):    
         """
-        处理日志特殊字符
+        处理命令中如果有ps1或者mysql的特殊情况,极端情况下会有ps1和mysql
+        :param command:要处理的字符传
+        :return:返回去除PS1或者mysql字符串的结果
         """
-        control_char = re.compile(r"""
+        result = None
-                \x1b[ #%()*+\-.\/]. |
+        match = self.ps1_pattern.split(command)
-                \r |                                               #匹配 回车符(CR)
+        if match:
-                (?:\x1b\[|\x9b) [ -?]* [@-~] |                     #匹配 控制顺序描述符(CSI)... Cmd
+            # 只需要最后的一个PS1后面的字符串
-                (?:\x1b\]|\x9d) .*? (?:\x1b\\|[\a\x9c]) | \x07 |   #匹配 操作系统指令(OSC)...终止符或振铃符(ST|BEL)
+            result = match[-1].strip()
-                (?:\x1b[P^_]|[\x90\x9e\x9f]) .*? (?:\x1b\\|\x9c) | #匹配 设备控制串或私讯或应用程序命令(DCS|PM|APC)...终止符(ST)
-                \x1b.                                              #匹配 转义过后的字符
-                [\x80-\x9f] | (?:\x1b\]0.*) | \[.*@.*\][\$#] | (.*mysql>.*)      #匹配 所有控制字符
-                """, re.X)
-        result_command = control_char.sub('', result_command.strip())
- 
-        if not self.vim_flag:
-            if result_command.startswith('vi') or result_command.startswith('fg'):
-                self.vim_flag = True
-            return result_command.decode('utf8',"ignore")
         else:
-            return ''
+            # PS1没找到,查找mysql
+            match = re.split('mysql>\s', command)
+            if match:
+                # 只需要最后一个mysql后面的字符串
+                result = match[-1].strip()
+        return result
 
-    def deal_command(self, str_r):
+    def deal_command(self):
         """
-            处理命令中特殊字符
+        处理截获的命令
+        :return:返回最后的处理结果
         """
-        str_r = self.remove_obstruct_char(str_r)
+        command = ''
-
+        # 从虚拟屏幕中获取处理后的数据
-        result_command = ''             # 最后的结果
+        for line in reversed(self.screen.buffer):
-        backspace_num = 0               # 光标移动的个数
+            line_data = "".join(map(operator.attrgetter("data"), line)).strip()
-        reach_backspace_flag = False    # 没有检测到光标键则为true
+            if len(line_data) > 0:
-        pattern_str = ''
+                parser_result = self.command_parser(line_data)
-        while str_r:
+                if parser_result is not None:
-            tmp = re.match(r'\s*\w+\s*', str_r)
+                    # 2个条件写一起会有错误的数据
-            if tmp:
+                    if len(parser_result) > 0:
-                str_r = str_r[len(str(tmp.group(0))):]
+                        command = parser_result
-                if reach_backspace_flag:
-                    pattern_str += str(tmp.group(0))
-                    continue
-                else:
-                    result_command += str(tmp.group(0))
-                    continue
-                
-            tmp = re.match(r'\x1b\[K[\x08]*', str_r)
-            if tmp:
-                result_command, del_len = self.deal_backspace(str(tmp.group(0)), result_command, pattern_str, backspace_num)
-                reach_backspace_flag = False
-                backspace_num = 0
-                pattern_str = ''
-                str_r = str_r[del_len:]
-                continue
-            
-            tmp = re.match(r'\x08+', str_r)
-            if tmp:
-                str_r = str_r[len(str(tmp.group(0))):]
-                if len(str_r) != 0:
-                    if reach_backspace_flag:
-                        result_command = result_command[0:-backspace_num] + pattern_str
-                        pattern_str = ''
-                    else:
-                        reach_backspace_flag = True
-                    backspace_num = len(str(tmp.group(0)))
-                    continue
                 else:
+                    command = line_data
                 break
-                
+        if command != '':
-            tmp = re.match(r'(\x1b\[1@\w)+', str_r)                           #处理替换的命令
+            # 判断用户输入的是否是vim 或者fg命令
-            if tmp:
+            if self.vim_pattern.search(command):
-                result_command,del_len = self.deal_replace_char(str(tmp.group(0)), result_command, backspace_num)
+                self.vim_flag = True
-                str_r = str_r[del_len:]
+        # 虚拟屏幕清空
-                backspace_num = 0
+        self.screen.reset()
-                continue
+        return command
-        
-            if reach_backspace_flag:
-                pattern_str += str_r[0]
-            else:
-                result_command += str_r[0]
-            str_r = str_r[1:]
-        
-        if backspace_num > 0:
-            result_command = result_command[0:-backspace_num] + pattern_str
-
-        result_command = self.remove_control_char(result_command)
-        return result_command
 
     def get_log(self):
         """
@@ -294,7 +214,7 @@ class Tty(object):
 
         # 发起ssh连接请求 Make a ssh connection
         ssh = paramiko.SSHClient()
-        #ssh.load_system_host_keys()
+        # ssh.load_system_host_keys()
         ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
         try:
             role_key = connect_info.get('role_key')
@@ -367,6 +287,7 @@ class SshTty(Tty):
         old_tty = termios.tcgetattr(sys.stdin)
         pre_timestamp = time.time()
         data = ''
+        input_str = ''
         input_mode = False
         try:
             tty.setraw(sys.stdin.fileno())
@@ -411,6 +332,8 @@ class SshTty(Tty):
                         if input_mode and not self.is_output(x):
                             data += x
 
+                        input_str = ''
+
                     except socket.timeout:
                         pass
 
@@ -420,19 +343,25 @@ class SshTty(Tty):
                     except OSError:
                         pass
                     input_mode = True
+                    input_str += x
                     if str(x) in ['\r', '\n', '\r\n']:
+                        # 这个是用来处理用户的复制操作
+                        if input_str != x:
+                            data += input_str
+                        self.stream.feed(data)
                         if self.vim_flag:
                             match = self.ps1_pattern.search(self.vim_data)
                             if match:
                                 self.vim_flag = False
-                                data = self.deal_command(data)[0:200]
+                                data = self.deal_command()[0:200]
                                 if len(data) > 0:
                                     TtyLog(log=log, datetime=datetime.datetime.now(), cmd=data).save()
                         else:
-                            data = self.deal_command(data)[0:200]
+                            data = self.deal_command()[0:200]
                             if len(data) > 0:
                                 TtyLog(log=log, datetime=datetime.datetime.now(), cmd=data).save()
                         data = ''
+                        input_str = ''
                         self.vim_data = ''
                         input_mode = False
 
@@ -466,6 +395,8 @@ class SshTty(Tty):
         win_size = self.get_win_size()
         #self.channel = channel = ssh.invoke_shell(height=win_size[0], width=win_size[1], term='xterm')
         self.channel = channel = transport.open_session()
+        self.screen = pyte.Screen(win_size[1], win_size[0])
+        self.stream.attach(self.screen)
         channel.get_pty(term='xterm', height=win_size[0], width=win_size[1])
         channel.invoke_shell()
         try: