fix regression for ssh and netconf

posntw · Feb 19, 2023 · 8c0151e · 8c0151e
1 parent ae7fd8c
commit 8c0151e
Show file tree

Hide file tree

Showing 6 changed files with 63 additions and 46 deletions.
diff --git a/conftest.py b/conftest.py
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.60"
+__version__ = "0.0.61"
 
 import os
 import subprocess
@@ -43,9 +43,11 @@ def mb(imposters, loglevel="info"):
 
 
 def start_mb(loglevel="info"):
+    mb_netmgmt_dir = os.path.dirname(__file__)
     return subprocess.Popen(
         ["mb", "--loglevel", loglevel, "--pidfile", "/tmp/mb.pid"],
-        cwd=os.path.dirname(__file__),
+        cwd=mb_netmgmt_dir,
+        env=dict(os.environ, PYTHONPATH=mb_netmgmt_dir + "/.."),
     )
 
 

diff --git a/mb_netmgmt/__main__.py b/mb_netmgmt/__main__.py
@@ -92,6 +92,7 @@ def post_request(self, request):
         return response.json()
 
     def get_to(self):
+        self.key_filename = None
         try:
             imposter_response = requests.get(
                 self.server.callback_url.replace("/_requests", "")
@@ -104,16 +105,17 @@ def get_to(self):
                 self.save_key(proxy)
                 disable_algorithms(proxy.get("disabled_algorithms", {}))
                 return urlparse(proxy["to"])
-        except IndexError:
+        except (IndexError, AttributeError):
             pass
 
     def save_key(self, proxy):
-        self.keyfile = tempfile.NamedTemporaryFile("w")
+        keyfile = tempfile.NamedTemporaryFile("w")
         try:
-            self.keyfile.write(proxy["key"])
-            self.keyfile.flush()
+            keyfile.write(proxy["key"])
+            keyfile.flush()
+            self.key_filename = keyfile.name
         except KeyError:
-            pass
+            self.key_filename = None
 
     def get_proxy(self, stub):
         return stub["responses"][0].get("proxy")

diff --git a/mb_netmgmt/netconf.py b/mb_netmgmt/netconf.py
@@ -51,7 +51,9 @@ def setup(self):
 
     def handle(self):
         self.callback_url = self.server.callback_url
-        transport = start_server(self.request)
+        transport = start_server(
+            self.request, self.get_to(), self.key_filename, self.handle_request
+        )
         self.channel = transport.accept()
         self.open_upstream()
         self.session._connected = True
@@ -69,7 +71,7 @@ def open_upstream(self):
             port=to.port or PORT_NETCONF_DEFAULT,
             username=to.username,
             password=to.password,
-            key_filename=self.keyfile.name,
+            key_filename=self.key_filename,
             hostkey_verify=False,
             timeout=60,
         )

diff --git a/mb_netmgmt/ssh.py b/mb_netmgmt/ssh.py
@@ -53,7 +53,7 @@ def check_channel_pty_request(
             pixelwidth,
             pixelheight,
         )
-        channel.command_prompt = handle_prompt(transport.handle_request)
+        channel.command_prompt = handle_prompt(self.handle_request)
         return True
 
     def check_channel_shell_request(*args):
@@ -66,7 +66,9 @@ def check_channel_subsystem_request(*args):
 class Handler(BaseRequestHandler, Protocol):
     def handle(self):
         self.callback_url = self.server.callback_url
-        transport = start_server(self.request, self.get_to(), self.keyfile.name)
+        transport = start_server(
+            self.request, self.get_to(), self.key_filename, self.handle_request
+        )
         self.channel = transport.accept()
         while not stopped:
             request, request_id = self.read_request()
@@ -125,13 +127,14 @@ def handle_prompt(handle_request):
     return command_prompt
 
 
-def start_server(request, to, key_filename):
+def start_server(request, to, key_filename, handle_request):
     t = paramiko.Transport(request)
     t.add_server_key(paramiko.DSSKey.generate())
     t.add_server_key(paramiko.ECDSAKey.generate())
     t.add_server_key(paramiko.RSAKey.generate(4096))
     t.to = to
     t.key_filename = key_filename
-    t.handle_request = self.handle_request
-    t.start_server(server=ParamikoServer())
+    paramiko_server = ParamikoServer()
+    paramiko_server.handle_request = handle_request
+    t.start_server(server=paramiko_server)
     return t
diff --git a/test/test_mb_netmgmt.py b/test/test_mb_netmgmt.py
@@ -10,7 +10,7 @@
 from ncclient.transport.session import BASE_NS_1_0, to_ele
 from ncclient.transport.ssh import MSG_DELIM
 
-port = 8080
+port = 8081
 prompt = b"prompt"
 mock_response = f"""<rpc-reply xmlns="{BASE_NS_1_0}">
   <blubb/>
@@ -33,21 +33,25 @@ def test_ssh():
 
 def test_ssh_proxy():
     with mb(
-        imposter(
-            "ssh",
-            [
-                prompt_stub(),
-                {
-                    "responses": [
-                        {
-                            "proxy": {
-                                "to": f"ssh://{os.environ['NETCONF_USERNAME']}:{os.environ['NETCONF_PASSWORD']}@localhost"
-                            }
-                        },
-                    ]
-                },
-            ],
-        )
+        [
+            {"protocol": "ssh", "port": 2222, "stubs": [prompt_stub()]},
+            {
+                "protocol": "ssh",
+                "port": port,
+                "stubs": [
+                    prompt_stub(),
+                    {
+                        "responses": [
+                            {
+                                "proxy": {
+                                    "to": f"ssh://{os.environ['NETCONF_USERNAME']}:{os.environ['NETCONF_PASSWORD']}@localhost:2222"
+                                }
+                            },
+                        ]
+                    },
+                ],
+            },
+        ]
     ):
         client = connect_ssh()
         chan = client.invoke_shell()
@@ -56,7 +60,7 @@ def test_ssh_proxy():
 
 
 def prompt_stub():
-    return {"responses": [{"is": {"response": prompt}}]}
+    return {"responses": [{"is": {"response": prompt.decode()}}]}
 
 
 def test_create_ssh_server():
@@ -158,21 +162,25 @@ def test_netconf_default_response():
 
 def test_netconf_private_key():
     with mb(
-        imposter(
-            "netconf",
-            [
-                {
-                    "responses": [
-                        {
-                            "proxy": {
-                                "to": f"netconf://{os.environ['NETCONF_USERNAME']}@{os.environ['NETCONF_HOSTNAME']}",
-                                "key": os.environ["NETCONF_KEY"],
-                            }
-                        },
-                    ]
-                },
-            ],
-        )
+        [
+            {"protocol": "netconf", "port": 830, "stubs": []},
+            {
+                "protocol": "netconf",
+                "port": port,
+                "stubs": [
+                    {
+                        "responses": [
+                            {
+                                "proxy": {
+                                    "to": f"netconf://{os.environ['NETCONF_USERNAME']}@localhost",
+                                    "key": os.environ["NETCONF_KEY"],
+                                }
+                            },
+                        ]
+                    },
+                ],
+            },
+        ]
     ):
         ncclient.manager.connect(
             host="localhost", port=port, password="", hostkey_verify=False