From fd713e0e5ca778a5c9659fe5e44080daaf6b4fb8 Mon Sep 17 00:00:00 2001
From: kelianchun_miller <kelianchun@qq.com>
Date: Fri, 25 Mar 2016 21:08:52 +0800
Subject: [PATCH] pull issue #120  remove ps1

---
 connect.py    | 16 +++++++---------
 run_server.py | 11 ++++++-----
 2 files changed, 13 insertions(+), 14 deletions(-)

diff --git a/connect.py b/connect.py
index c3d575249..7ff8d344d 100755
--- a/connect.py
+++ b/connect.py
@@ -92,7 +92,8 @@ class Tty(object):
         self.remote_ip = ''
         self.login_type = login_type
         self.vim_flag = False
-        self.ps1_pattern = re.compile('\[?.*@.*\]?[\$#]\s')
+        self.vim_end_flag = False
+        self.vim_end_pattern = re.compile(r'\x1b\[\?1049', re.X)
         self.vim_pattern = re.compile(r'\W?vi[m]?\s.* | \W?fg\s.*', re.X)
         self.vim_data = ''
         self.stream = None
@@ -122,7 +123,7 @@ class Tty(object):
         :return:返回去除PS1或者mysql字符串的结果
         """
         result = None
-        match = self.ps1_pattern.split(command)
+        match = re.compile('\[?.*@.*\]?[\$#]\s').split(command)
         if match:
             # 只需要最后的一个PS1后面的字符串
             result = match[-1].strip()
@@ -303,7 +304,6 @@ class SshTty(Tty):
         data = ''
         input_str = ''
         input_mode = False
-        vim_end_flag = False
         try:
             tty.setraw(sys.stdin.fileno())
             tty.setcbreak(sys.stdin.fileno())
@@ -334,8 +334,6 @@ class SshTty(Tty):
                             except OSError as msg:
                                 if msg.errno == errno.EAGAIN:
                                     continue
-                        #sys.stdout.write(x)
-                        #sys.stdout.flush()
                         now_timestamp = time.time()
                         log_time_f.write('%s %s\n' % (round(now_timestamp-pre_timestamp, 4), len(x)))
                         log_time_f.flush()
@@ -364,13 +362,13 @@ class SshTty(Tty):
                         if input_str != x:
                             data += input_str
                         if self.vim_flag:
-                            match = re.compile(r'\x1b\[\?1049', re.X).findall(self.vim_data)
+                            match = self.vim_end_pattern.findall(self.vim_data)
                             if match:
-                                if vim_end_flag or len(match) == 2:
+                                if self.vim_end_flag or len(match) == 2:
                                     self.vim_flag = False
-                                    vim_end_flag = False
+                                    self.vim_end_flag = False
                                 else:
-                                    vim_end_flag = True
+                                    self.vim_end_flag = True
                         else:
                             data = self.deal_command(data)[0:200]
                             if len(data) > 0:
diff --git a/run_server.py b/run_server.py
index 8ec0d46b8..00a6eb8cf 100755
--- a/run_server.py
+++ b/run_server.py
@@ -364,12 +364,13 @@ class WebTerminalHandler(tornado.websocket.WebSocketHandler):
             self.term.input_mode = True
             if str(jsondata['data']) in ['\r', '\n', '\r\n']:
                 if self.term.vim_flag:
-                    match = self.term.ps1_pattern.search(self.term.vim_data)
+                    match = re.compile(r'\x1b\[\?1049', re.X).findall(self.vim_data)
                     if match:
-                        self.term.vim_flag = False
-                        result = self.term.deal_command(self.term.data)[0:200]
-                        if len(result) > 0:
-                            TtyLog(log=self.log, datetime=datetime.datetime.now(), cmd=result).save()
+                        if self.term.vim_end_flag or len(match) == 2:
+                            self.term.vim_flag = False
+                            self.term.vim_end_flag = False
+                        else:
+                            self.term.vim_end_flag = True
                 else:
                     result = self.term.deal_command(self.term.data)[0:200]
                     if len(result) > 0: