Fixes to netconf on SSH tunnels

connector/docs/changelog/undistributed/netconf_sshtunnel_logging.rst

@@ -0,0 +1,7 @@
+--------------------------------------------------------------------------------
+Fix
+--------------------------------------------------------------------------------
+* yang.connector
+    * Modified Netconf:
+        * Captured SSH tunnel setup and ncclient session logs in the per-connection NETCONF log file.
+        * Avoided attaching the pyATS tasklog handler when its stream is unavailable.

connector/src/yang/connector/netconf.py

@@ -3,6 +3,7 @@
 import re
 import time
 import atexit
+import copy
 import logging
 import subprocess
 import datetime
@@ -127,6 +128,29 @@ def stream(self):
         return self._pyats_handlers.tasklog.stream
 
 
+class NetconfLogForwardingHandler(logging.Handler):
+    """Forward helper log records to the connection log file."""
+
+    def __init__(self, log):
+        super().__init__()
+        self.log = log
+
+    def emit(self, record):
+        if getattr(record, '_yang_connector_forwarded', False):
+            return
+        if not self.log.isEnabledFor(record.levelno):
+            return
+
+        forwarded_record = copy.copy(record)
+        forwarded_record._yang_connector_forwarded = True
+        for handler in self.log.handlers:
+            if not isinstance(handler, logging.FileHandler):
+                continue
+            if record.levelno < handler.level:
+                continue
+            handler.handle(forwarded_record)
+
+
 class NetconfSessionLogHandler(logging.Handler):
     """Logging handler that pretty prints ncclient XML."""
 
@@ -153,6 +177,11 @@ def emit(self, record):
                 # Unable to handle record so leave it unchanged
                 pass
 
+        session = getattr(record, 'session', None)
+        log = getattr(session, '_yang_connector_log', None)
+        if log:
+            NetconfLogForwardingHandler(log).emit(record)
+
 
 nccl.addHandler(NetconfSessionLogHandler())
 
@@ -317,7 +346,9 @@ def configure_logging(self):
         self.log = logging.getLogger('netconf.%s' % logger_name)
 
         # workaround for double invocation that somehow happens in robot
-        self.log.handlers.clear()
+        for handler in self.log.handlers[:]:
+            self.log.removeHandler(handler)
+            handler.close()
         self.log.filters.clear()
 
         # default log level
@@ -367,7 +398,9 @@ def convert(string):
             pass
         else:
             # we're in pyATS, use pyATS loggers
-            if not self.no_pyats_tasklog:
+            tasklog = getattr(managed_handlers, 'tasklog', None)
+            tasklog_stream = getattr(tasklog, 'stream', None)
+            if not self.no_pyats_tasklog and tasklog_stream is not None:
                 pta = pyATS_TaskLog_Adapter()
                 nsf = NetconfScreenFormatter(fmt=TaskLogFormatter.MESSAGE_FORMAT)
                 nsf.MAX_LINES = self.settings.get('NETCONF_SCREEN_LOGGING_MAX_LINES', 40)
@@ -378,6 +411,12 @@ def convert(string):
         # if debug_mode is True, enable debug mode
         if self.debug:
             self.log.setLevel(logging.DEBUG)
+            nccl.setLevel(logging.DEBUG)
+        else:
+            nccl.setLevel(logging.INFO)
+
+    def configure_session_logging(self):
+        self.session._yang_connector_log = self.log
 
     def connect(self):
         '''connect
@@ -489,6 +528,7 @@ def connect(self):
 
         if not self.session.is_alive():
             self._session = transport.SSHSession(self._device_handler)
+        self.configure_session_logging()
 
         # default values
         defaults = {
@@ -533,6 +573,9 @@ def connect(self):
         # support sshtunnel
         if 'sshtunnel' in defaults:
             from unicon.sshutils import sshtunnel
+            tunnel_logger = logging.getLogger('unicon.sshutils')
+            tunnel_log_handler = NetconfLogForwardingHandler(self.log)
+            tunnel_logger.addHandler(tunnel_log_handler)
             try:
                 tunnel_port = sshtunnel.auto_tunnel_add(self.device, self.via)
                 if tunnel_port:
@@ -543,6 +586,8 @@ def connect(self):
                 raise AttributeError("Cannot add ssh tunnel. \
                 Connection %s may not have ip/host or port.\n%s"
                                      % (self.via, err))
+            finally:
+                tunnel_logger.removeHandler(tunnel_log_handler)
             del defaults['sshtunnel']
 
         defaults = {k: getattr(self, k, v) for k, v in defaults.items()}

connector/src/yang/connector/tests/test_yang.py

@@ -1,12 +1,16 @@
 #!/bin/env python
 """ Unit tests for the yang.connector cisco-shared package. """
 
+import logging
+import os
+import tempfile
 import unittest
 from ncclient import manager
 from ncclient import transport
 from ncclient.devices.default import DefaultDeviceHandler
 from pyats.topology import loader
 from pyats.connections import BaseConnection
+from pyats.datastructures import AttrDict
 from unittest.mock import Mock, patch
 import yang.connector
 
@@ -32,6 +36,7 @@ def connected(self):
         return self._connected
 
     def connect(self, **kwargs):
+        self.connect_kwargs = kwargs
         self._connected = True
         self.transport = MyTransportSession()
 
@@ -62,6 +67,7 @@ def connected(self):
         return self._connected
 
     def connect(self, **kwargs):
+        self.connect_kwargs = kwargs
         if kwargs['username'] == 'admin' and kwargs['password'] == 'admin':
             self._connected = True
             self.transport = MyTransportSession()
@@ -197,6 +203,152 @@ def test_connect_4(self):
         expected_value = False
         self.assertEqual(generated_value, expected_value)
 
+    def test_connect_sshtunnel_logging(self):
+        yaml = \
+            'devices:\n' \
+            '    dummy:\n' \
+            '        os: iosxe\n' \
+            '        type: dummy_device\n' \
+            '        connections:\n' \
+            '            netconf:\n' \
+            '                class: yang.connector.Netconf\n' \
+            '                protocol: netconf\n' \
+            '                ip : "1.2.3.4"\n' \
+            '                port: 830\n' \
+            '                username: admin\n' \
+            '                password: admin\n' \
+            '                sshtunnel:\n' \
+            '                    host: proxy\n'
+
+        testbed = loader.load(yaml)
+        device = testbed.devices['dummy']
+        device.connections.netconf.sshtunnel = AttrDict(
+            device.connections.netconf.sshtunnel)
+        device.connections.netconf.sshtunnel.tunnel_ip = '127.0.0.1'
+        logfile = tempfile.mktemp(suffix='.log')
+
+        nc_device = yang.connector.Netconf(device=device,
+                                           alias='nc',
+                                           via='netconf',
+                                           logfile=logfile,
+                                           log_stdout=False,
+                                           no_pyats_tasklog=True)
+        nc_device._session = MySSHSession()
+
+        def add_tunnel(device, via):
+            logging.getLogger('unicon.sshutils').info(
+                "Device '%s' connection '%s' via new SSH tunnel %s:%s",
+                device.name, via, '127.0.0.1', 123)
+            return 123
+
+        try:
+            with patch('unicon.sshutils.sshtunnel.auto_tunnel_add',
+                       side_effect=add_tunnel):
+                nc_device.connect()
+
+            with open(logfile) as log_file:
+                log_content = log_file.read()
+
+            self.assertIn('via new SSH tunnel', log_content)
+            self.assertIn('NETCONF CONNECTED', log_content)
+            self.assertEqual(nc_device.session.connect_kwargs['host'],
+                             '127.0.0.1')
+            self.assertEqual(nc_device.session.connect_kwargs['port'], 123)
+        finally:
+            if os.path.exists(logfile):
+                os.remove(logfile)
+
+    def test_ncclient_session_logging(self):
+        logfile = tempfile.mktemp(suffix='.log')
+        nc_device = yang.connector.Netconf(device=self.device,
+                                           alias='nc',
+                                           via='netconf',
+                                           logfile=logfile,
+                                           log_stdout=False,
+                                           no_pyats_tasklog=True)
+        nc_device._session = MySSHSession()
+
+        try:
+            nc_device.connect()
+            logging.getLogger('ncclient.transport.ssh').info(
+                'Sending:\n%s', b'<hello/>',
+                extra={'session': nc_device.session})
+
+            with open(logfile) as log_file:
+                log_content = log_file.read()
+
+            self.assertIn('Sending:', log_content)
+            self.assertIn('<hello/>', log_content)
+        finally:
+            if os.path.exists(logfile):
+                os.remove(logfile)
+
+    def test_log_forwarding_only_writes_file_handlers(self):
+        class FailingHandler(logging.Handler):
+            def emit(self, record):
+                raise AssertionError('non-file handler should not be used')
+
+        logfile = tempfile.mktemp(suffix='.log')
+        log = logging.getLogger('test.netconf.forwarding')
+        log.handlers.clear()
+        log.setLevel(logging.INFO)
+        log.propagate = False
+
+        file_handler = logging.FileHandler(logfile)
+        file_handler.setFormatter(yang.connector.netconf.NetconfFormatter())
+        log.addHandler(FailingHandler())
+        log.addHandler(file_handler)
+
+        try:
+            record = logging.LogRecord(
+                'unicon.sshutils', logging.INFO, __file__, 0,
+                'Adding local tunnel %s', ('127.0.0.1:123',), None)
+            handler = yang.connector.netconf.NetconfLogForwardingHandler(log)
+            handler.emit(record)
+
+            with open(logfile) as log_file:
+                log_content = log_file.read()
+
+            self.assertIn('Adding local tunnel 127.0.0.1:123', log_content)
+        finally:
+            for handler in log.handlers[:]:
+                log.removeHandler(handler)
+                handler.close()
+            if os.path.exists(logfile):
+                os.remove(logfile)
+
+    def test_configure_logging_skips_unavailable_tasklog(self):
+        from pyats.log import managed_handlers
+
+        logfile = tempfile.mktemp(suffix='.log')
+        nc_device = yang.connector.Netconf(device=self.device,
+                                           alias='nc',
+                                           via='netconf',
+                                           logfile=logfile,
+                                           log_stdout=False)
+        original_stream = managed_handlers.tasklog.stream
+
+        try:
+            managed_handlers.tasklog.stream = None
+            nc_device.configure_logging()
+
+            tasklog_handlers = [
+                handler for handler in nc_device.log.handlers
+                if isinstance(
+                    handler,
+                    yang.connector.netconf.pyATS_TaskLog_Adapter)
+            ]
+            self.assertEqual(tasklog_handlers, [])
+
+            nc_device.log.info('tasklog stream unavailable')
+        finally:
+            managed_handlers.tasklog.stream = original_stream
+            for handler in nc_device.log.handlers[:]:
+                nc_device.log.removeHandler(handler)
+                handler.close()
+            if os.path.exists(logfile):
+                os.remove(logfile)
+
     @patch('yang.connector.netconf.RawRPC', new=MyRawRPC)
     def test_request(self):
         self.nc_device._session = MySSHSession()