Update console connection tools

1. Refactor code 2. Support interactive_shell Signed-off-by: bingwang <[email protected]>
sonic-net · Jan 17, 2021 · fca2f03 · fca2f03
1 parent 65c3293
commit fca2f03
Show file tree

Hide file tree

Showing 4 changed files with 68 additions and 36 deletions.
diff --git a/tests/common/connections/__init__.py b/tests/common/connections/__init__.py
@@ -1,34 +1 @@
-from tests.common.connections.base_console_conn import CONSOLE_SSH, CONSOLE_SSH_MENU_PORTS, CONSOLE_TELNET
-from telnet_console_conn import TelnetConsoleConn
-from ssh_console_conn import SSHConsoleConn
-
-__all__ = ["TelnetConsoleConn", "SSHConsoleConn"]
-
-ConsoleTypeMapper = {
-    CONSOLE_TELNET: TelnetConsoleConn,
-    CONSOLE_SSH: SSHConsoleConn,
-    CONSOLE_SSH_MENU_PORTS: SSHConsoleConn
-}
-
-def ConsoleHost(console_type,
-                console_host,
-                console_port,
-                sonic_username,
-                sonic_password,
-                console_username=None,
-                console_password=None,
-                timeout_s=100):
-    if not ConsoleTypeMapper.has_key(console_type):
-        raise ValueError("console type {} is not supported yet".format(console_type))
-    params = {
-        "console_host": console_host,
-        "console_port": console_port,
-        "console_type": console_type,
-        "sonic_username": sonic_username,
-        "sonic_password": sonic_password,
-        "console_username": console_username,
-        "console_password": console_password,
-        "timeout": timeout_s
-    }
-    return ConsoleTypeMapper[console_type](**params)
-
+__all__ = ["TelnetConsoleConn", "SSHConsoleConn", "ConsoleHost"]
diff --git a/tests/common/connections/base_console_conn.py b/tests/common/connections/base_console_conn.py
@@ -6,6 +6,14 @@
 from netmiko.cisco_base_connection import CiscoBaseConnection
 from netmiko.ssh_exception import NetMikoAuthenticationException
 
+# For interactive shell
+import sys
+import socket
+from paramiko.py3compat import u
+import termios
+import tty
+import select
+
 # All supported console types
 # Console login via telnet (mad console)
 CONSOLE_TELNET = "console_telnet"
@@ -81,3 +89,30 @@ def cleanup(self):
     def disconnect(self):
         super(BaseConsoleConn, self).disconnect()
 
+    def posix_shell(self):
+        oldtty = termios.tcgetattr(sys.stdin)
+        try:
+            tty.setraw(sys.stdin.fileno())
+            tty.setcbreak(sys.stdin.fileno())
+            self.remote_conn.settimeout(0.0)
+
+            while True:
+                r, w, e = select.select([self.remote_conn, sys.stdin], [], [])
+                if self.remote_conn in r:
+                    try:
+                        x = u(self.remote_conn.recv(1024))
+                        if len(x) == 0:
+                            sys.stdout.write("\r\n*** EOF\r\n")
+                            break
+                        sys.stdout.write(x)
+                        sys.stdout.flush()
+                    except socket.timeout:
+                        pass
+                if sys.stdin in r:
+                    x = sys.stdin.read(1)
+                    if len(x) == 0:
+                        break
+                    self.remote_conn.send(x)
+
+        finally:
+            termios.tcsetattr(sys.stdin, termios.TCSADRAIN, oldtty)
diff --git a/tests/common/connections/console_host.py b/tests/common/connections/console_host.py
@@ -0,0 +1,31 @@
+from base_console_conn import CONSOLE_SSH, CONSOLE_SSH_MENU_PORTS, CONSOLE_TELNET
+from telnet_console_conn import TelnetConsoleConn
+from ssh_console_conn import SSHConsoleConn
+
+ConsoleTypeMapper = {
+    CONSOLE_TELNET: TelnetConsoleConn,
+    CONSOLE_SSH: SSHConsoleConn,
+    CONSOLE_SSH_MENU_PORTS: SSHConsoleConn
+}
+
+def ConsoleHost(console_type,
+                console_host,
+                console_port,
+                sonic_username,
+                sonic_password,
+                console_username=None,
+                console_password=None,
+                timeout_s=100):
+    if not ConsoleTypeMapper.has_key(console_type):
+        raise ValueError("console type {} is not supported yet".format(console_type))
+    params = {
+        "console_host": console_host,
+        "console_port": console_port,
+        "console_type": console_type,
+        "sonic_username": sonic_username,
+        "sonic_password": sonic_password,
+        "console_username": console_username,
+        "console_password": console_password,
+        "timeout": timeout_s
+    }
+    return ConsoleTypeMapper[console_type](**params)
diff --git a/tests/common/connections/ssh_console_conn.py b/tests/common/connections/ssh_console_conn.py
@@ -1,7 +1,6 @@
-from tests.common.connections.base_console_conn import CONSOLE_SSH
 import time
 import re
-from base_console_conn import BaseConsoleConn
+from base_console_conn import BaseConsoleConn, CONSOLE_SSH
 from netmiko.ssh_exception import NetMikoAuthenticationException
 
 class SSHConsoleConn(BaseConsoleConn):