From 5a00736f23c7981b7995eb33c4018376775f6f7f Mon Sep 17 00:00:00 2001 From: oliver Date: Tue, 23 Jun 2026 22:35:37 +0800 Subject: [PATCH] fix(bastion): subclass netmiko driver class on Python 3.14 ConnectHandler is a factory function in current Netmiko; subclass the CLASS_MAPPER driver (e.g. ZteZxrosSSH) when wrapping a pre-authenticated Paramiko session. Co-authored-by: Cursor --- netx_api/ne_session_factory.py | 16 ++++++++++++++-- tests/test_bastion_hop.py | 7 +++++++ 2 files changed, 21 insertions(+), 2 deletions(-) diff --git a/netx_api/ne_session_factory.py b/netx_api/ne_session_factory.py index f438916..d4dc3e6 100644 --- a/netx_api/ne_session_factory.py +++ b/netx_api/ne_session_factory.py @@ -231,6 +231,17 @@ def _bastion_ssh_connect( ) +def _netmiko_driver_class(device_type: str) -> type: + """Resolve Netmiko driver class (ConnectHandler is a factory func, not a base class).""" + from netmiko.ssh_dispatcher import CLASS_MAPPER + + dt = str(device_type or "").strip() + cls = CLASS_MAPPER.get(dt) + if cls is None: + raise ValueError(f"unsupported_device_type: {dt}") + return cls + + def _netmiko_over_ssh_client( ssh_client: paramiko.SSHClient, *, @@ -243,8 +254,9 @@ def _netmiko_over_ssh_client( session_timeout: int | None, ) -> ConnectHandler: """Netmiko session over an already-authenticated SSH client (bastion protocol proxy).""" + base_cls = _netmiko_driver_class(device_type) - class _PreauthConnectHandler(ConnectHandler): + class _PreauthSession(base_cls): def establish_connection(self, width: int = 511, height: int = 1000) -> None: from netmiko.channel import SSHChannel @@ -267,7 +279,7 @@ def _netmiko_over_ssh_client( enable_secret=enable_secret, session_timeout=session_timeout, ) - return _PreauthConnectHandler(**dev) + return _PreauthSession(**dev) def _base_connect_kwargs( diff --git a/tests/test_bastion_hop.py b/tests/test_bastion_hop.py index 8470435..94833ac 100644 --- a/tests/test_bastion_hop.py +++ b/tests/test_bastion_hop.py @@ -64,6 +64,13 @@ class BastionTemplateTests(unittest.TestCase): "ssh ZTE-FIVIE@ca-admin@114.1.198.1@10.34.145.25", ) + def test_netmiko_driver_class_resolves_zte(self) -> None: + from netmiko.zte.zte_zxros import ZteZxrosSSH + + from netx_api.ne_session_factory import _netmiko_driver_class + + self.assertIs(_netmiko_driver_class("zte_zxros_ssh"), ZteZxrosSSH) + class BastionConnectRoutingTests(unittest.TestCase): @patch("netx_api.ne_session_factory._connect_via_bastion")