keep netconf capabilities from upstream

celinegrn · Dec 20, 2023 · e8b5a5d · e8b5a5d
1 parent 7f69806
commit e8b5a5d
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 10 deletions.
diff --git a/mb_netmgmt/__init__.py b/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.75"
+__version__ = "0.0.76"
 
 import os
 import subprocess

diff --git a/mb_netmgmt/__main__.py b/mb_netmgmt/__main__.py
@@ -68,11 +68,14 @@ def get_response(self, mb_response):
         except KeyError:
             proxy_response = self.read_proxy_response()
             logging.debug("proxy_response: %s", proxy_response)
-            response = requests.post(
-                mb_response["callbackURL"], json={"proxyResponse": proxy_response}
-            )
-            response.raise_for_status()
-            return response.json()
+            return self.post_proxy_response(mb_response, proxy_response)
+
+    def post_proxy_response(self, mb_response, proxy_response):
+        response = requests.post(
+            mb_response["callbackURL"], json={"proxyResponse": proxy_response}
+        )
+        response.raise_for_status()
+        return response.json()
 
     def send_upstream(self, request, request_id):
         raise NotImplementedError
@@ -130,7 +133,7 @@ def get_cli_patterns():
         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>[A-Z][\w\/ .:,>\(\)\-\?\"]*[^A-Z])(?P<default>\[[\w\/.,():\-]*\])?(?(default)(?P<end1>(?:\?|: ?| |)$)|(?P<end2>: $))"
+        b'[\r\n\x00\x1b\[K]+(?P<text>[A-Z][\w\/ .:,>\(\)\-\?"]*[^A-Z])(?P<default>\[[\w\/.,():\-]*\])?(?(default)(?P<end1>(?:\?|: ?| |)$)|(?P<end2>: $))'
     )  # Interactive prompt
     patterns.append(b"[\r\n\x00\x1b\[K] --More-- $")  # Terminal paging
     return patterns

diff --git a/mb_netmgmt/netconf.py b/mb_netmgmt/netconf.py
@@ -78,9 +78,16 @@ def open_upstream(self):
         )
 
     def handle_prompt(self):
-        hello = to_ele(
-            HelloHandler.build(DefaultDeviceHandler._BASE_CAPABILITIES, None)
-        )
+        mb_response = self.post_request({"rpc": ""})
+        try:
+            response = mb_response["response"]
+        except KeyError:
+            try:
+                capabilities = [c for c in self.manager.server_capabilities]
+            except AttributeError:
+                capabilities = DefaultDeviceHandler._BASE_CAPABILITIES
+            response = self.post_proxy_response(mb_response, capabilities)
+        hello = to_ele(HelloHandler.build(response, None))
 
         # A server sending the <hello> element MUST include a <session-id>
         # element containing the session ID for this NETCONF session.