Merge pull request telekom#1 from phoenixszg/main

mb_netmgmt/__init__.py

@@ -18,7 +18,7 @@
 # along with mb-netmgmt. If not, see <https://www.gnu.org/licenses/
 
 """Network Management Protocols for Mountebank"""
-__version__ = "0.0.65"
+__version__ = "0.0.66"
 
 import os
 import subprocess

mb_netmgmt/__main__.py

@@ -121,6 +121,21 @@ def get_proxy(self, stub):
         return stub["responses"][0].get("proxy")
 
 
+def get_cli_patterns():
+    patterns = []
+    patterns.append(
+        b"[\r\n\x1b\[K][\-\w+\.:/]+(?:\([^\)]+\))?[>#] ?$"
+    )  # based on IOS driver of Exscript
+    patterns.append(
+        b"[\r\n\x00\x1b\[K]RP/\d+/(?:RS?P)?\d+\/CPU\d+:[^#]+(?:\([^\)]+\))?#$"
+    )  # based on IOS XR driver of Exscript
+    patterns.append(
+        b"[\r\n\x00\x1b\[K](?P<text>[\w/ .:,\(\)\-\?]*)(?P<default>\[[\w/.,():\-]*\])?(?(default)(?P<end1>(?:\?|: ?| |)$)|(?P<end2>: $))"
+    )  # Interactive prompt
+    patterns.append(b"[\r\n\x00\x1b\[K] --More-- $")  # Terminal paging
+    return patterns
+
+
 def disable_algorithms(disabled_algorithms):
     # https://github.com/ncclient/ncclient/issues/526#issuecomment-1096563028
     class MonkeyPatchedTransport(paramiko.Transport):

mb_netmgmt/ssh.py

@@ -17,12 +17,13 @@
 # You should have received a copy of the GNU General Public License
 # along with mb-netmgmt. If not, see <https://www.gnu.org/licenses/
 
+import re
 from socketserver import BaseRequestHandler
 from socketserver import ThreadingTCPServer as Server
 
 import paramiko
 
-from mb_netmgmt.__main__ import Protocol
+from mb_netmgmt.__main__ import Protocol, get_cli_patterns
 
 stopped = False
 
@@ -88,19 +89,20 @@ def respond(self, response, request_id):
         return response
 
     def read_proxy_response(self):
-        message = self.read_message(
-            self.channel.upstream, [self.channel.command_prompt]
-        )
+        prompt_patterns = [self.channel.command_prompt]
+        prompt_patterns += get_cli_patterns()
+        message = self.read_message(self.channel.upstream, prompt_patterns)
         return {"response": message.decode()}
 
-    def read_message(self, channel, terminators):
+    def read_message(self, channel, patterns):
         message = b""
         end_of_message = False
         while not end_of_message and not stopped:
             message += channel.recv(1024)
-            for terminator in terminators:
-                if terminator in message:
+            for pattern in patterns:
+                if re.findall(pattern, message):
                     end_of_message = True
+                    break
         return message
 
 

mb_netmgmt/telnet.py

@@ -22,7 +22,7 @@
 import time
 from socketserver import StreamRequestHandler, TCPServer
 
-from mb_netmgmt.__main__ import Protocol
+from mb_netmgmt.__main__ import Protocol, get_cli_patterns
 
 Server = TCPServer
 
@@ -42,7 +42,8 @@ def handle(self):
             pass
         self.wfile.write(self.command_prompt)
         request = None
-        while request != b"exit\r\n":
+        self.stopped = False
+        while not self.stopped:
             request, request_id = self.read_request()
             if not request:
                 return
@@ -59,9 +60,14 @@ def read_request(self):
 
     def respond(self, response, request_id):
         self.wfile.write(response["response"].encode())
+        if response["response"].encode() in [b"exit\r\n\r", b"quit\r\n\r"]:
+            self.stopped = True
 
     def read_proxy_response(self):
-        return {"response": self.telnet.read_until(self.command_prompt).decode()}
+        prompt_patterns = [self.command_prompt]
+        prompt_patterns += get_cli_patterns()
+        _, _, response = self.telnet.expect(prompt_patterns)
+        return {"response": response.decode()}
 
     def handle_username_prompt(self):
         username_prompt = b"Username: "

test/test_mb_netmgmt.py

@@ -1,5 +1,6 @@
 import io
 import os
+import re
 from threading import Thread
 from urllib.parse import urlparse
 
@@ -10,13 +11,37 @@
 from ncclient.transport.ssh import MSG_DELIM
 
 from mb_netmgmt import mb, netconf, ssh, use_scalar_strings, yaml
-from mb_netmgmt.__main__ import create_server
+from mb_netmgmt.__main__ import create_server, get_cli_patterns
 
 port = 8081
 prompt = b"prompt#"
 mock_response = f"""<rpc-reply xmlns="{BASE_NS_1_0}">
   <blubb/>
 </rpc-reply>"""
+cli_responses = [
+    (b"\rIOS-1>", True),
+    (b"\rIOS-1#", True),
+    (b"\rIOS-1(config)#", True),
+    (b"\rIOS-1(config-if)#", True),
+    (b"\rRP/0/8/CPU0:IOSXR-2>", True),
+    (b"\rRP/0/8/CPU0:IOSXR-2#", True),
+    (b"\rRP/0/8/CPU0:IOSXR-2(config)#", True),
+    (b"\rRP/0/8/CPU0:IOSXR-2(config-if)#", True),
+    (
+        b"\rKUncommitted changes found, commit them before exiting(yes/no/cancel)? [cancel]:",
+        True,
+    ),
+    (b"\rProtocol [ipv4]: ", True),
+    (b"\rTarget IP address: ", True),
+    (b"\rHost name or IP address (control-c to abort): []?", True),
+    (b"\rDestination file name (control-c to abort): [running-config]?", True),
+    (b"\rDelete net/node0_8_CPU0/disk0:/c12k-mini.vm-4.3.2[confirm]", True),
+    (b"\rDestination filename [/net/node0_8_CPU0/disk0:/c12k-mini.vm-4.3.2]?", True),
+    (b"\rReload hardware module ? [no,yes] ", True),
+    (b"\rDo you wish to continue?[confirm(y/n)]", True),
+    (b"\r --More-- ", True),
+    (b"\rSending 5, 100-byte ICMP Echos to 8.8.8.8, timeout is 2 seconds:", False),
+]
 
 
 @pytest.mark.parametrize("protocol", ["http", "snmp", "telnet", "netconf"])
@@ -202,3 +227,13 @@ def test_use_scalar_strings(base, result):
     yaml.dump(base, s)
     s.seek(0)
     assert s.read() == result
+
+
+@pytest.mark.parametrize("cli_response,result", cli_responses)
+def test_cli_patterns(cli_response, result):
+    matched = False
+    for pattern in get_cli_patterns():
+        if re.findall(pattern, cli_response):
+            matched = True
+            break
+    assert matched == result