From 3b140d30292e631a3f70a4c651040c4841bb2876 Mon Sep 17 00:00:00 2001
From: Gabor <82544733+phoenixszg@users.noreply.github.com>
Date: Thu, 22 Jun 2023 10:48:35 +0200
Subject: [PATCH] Define CLI patterns centrally (#3)

* update telnet.py

* update telnet.py

* update ssh.py

* update telnet.py

* update telnet.py

* update telnet.py

* update telnet.py

* update telnet.py

* match for initial command as well

* define cli patterns centrally

* comment

---------
---
 mb_netmgmt/__main__.py |  7 +++++++
 mb_netmgmt/ssh.py      | 10 +++-------
 mb_netmgmt/telnet.py   |  8 ++------
 3 files changed, 12 insertions(+), 13 deletions(-)

diff --git a/mb_netmgmt/__main__.py b/mb_netmgmt/__main__.py
index a6f8734..4f8fdf6 100755
--- a/mb_netmgmt/__main__.py
+++ b/mb_netmgmt/__main__.py
@@ -120,6 +120,13 @@ def save_key(self, proxy):
     def get_proxy(self, stub):
         return stub["responses"][0].get("proxy")
 
+    def get_cli_patterns(self):
+        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
diff --git a/mb_netmgmt/ssh.py b/mb_netmgmt/ssh.py
index f347b58..8be0ed1 100644
--- a/mb_netmgmt/ssh.py
+++ b/mb_netmgmt/ssh.py
@@ -89,13 +89,9 @@ def respond(self, response, request_id):
         return response
 
     def read_proxy_response(self):
-        patterns = []
-        patterns.append(self.channel.command_prompt)  # Initial command prompt
-        patterns.append(b"[\r\n\x1b\[K][\-\w+\.:/]+(?:\([^\)]+\))?[>#] ?$")  # IOS
-        patterns.append(b"[\r\n\x00\x1b\[K]RP/\d+/(?:RS?P)?\d+\/CPU\d+:[^#]+(?:\([^\)]+\))?#$")  # IOS XR
-        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
-        message = self.read_message(self.channel.upstream, patterns)
+        prompt_patterns = [self.channel.command_prompt]
+        prompt_patterns += self.get_cli_patterns()
+        message = self.read_message(self.channel.upstream, prompt_patterns)
         return {"response": message.decode()}
 
     def read_message(self, channel, patterns):
diff --git a/mb_netmgmt/telnet.py b/mb_netmgmt/telnet.py
index ceb82bf..0648956 100644
--- a/mb_netmgmt/telnet.py
+++ b/mb_netmgmt/telnet.py
@@ -64,12 +64,8 @@ def respond(self, response, request_id):
             self.stopped = True
 
     def read_proxy_response(self):
-        prompt_patterns = []
-        prompt_patterns.append(self.command_prompt)  # Initial command prompt
-        prompt_patterns.append(b"[\r\n][\-\w+\.:/]+(?:\([^\)]+\))?[>#] ?$")  # IOS
-        prompt_patterns.append(b"[\r\n\x00]RP/\d+/(?:RS?P)?\d+\/CPU\d+:[^#]+(?:\([^\)]+\))?#$")  # IOS XR
-        prompt_patterns.append(b"[\r\n\x00](?P<text>[\w/ .:,\(\)\-\?]*)(?P<default>\[[\w/.:\-]*\])?(?(default)(?P<end1>(?:\?|: ?|)$)|(?P<end2>: $))")  # Interactive prompt
-        prompt_patterns.append(b"[\r\n\x00] --More-- $")  # Terminal paging
+        prompt_patterns = [self.command_prompt]
+        prompt_patterns += self.get_cli_patterns()
         _, _, response = self.telnet.expect(prompt_patterns)
         return {"response": response.decode()}