Using Python Context Managers for SSH connections

In this post, I will cover basic usage of Python’s context managers to connect to a network device using SSH. I will use them to abstract the connection establishment and teardown logic that is needed when making an SSH connection.

Note: This post will not cover context manager details, as great explanations can already be found online. Instead, this article is focused on showing how they can be used to facilitate the SSH connection logic to networking devices.

One of the simplest ways to connect to a network device programitically using Python is with Kirk Byer’s netmiko library. netmiko is a multi-vendor SSH Python library which makes connecting to network devices via SSH a breeze.  This library adds some vendor specific logic to paramiko, which is the de-facto SSH library in Python.  netmiko simplifies connecting to a network device via SSH and taking actions on the device – it allows you to use simple method calls such as “send_command” to execute commands on a device, and will properly parse the response according to the device being connected to. The vendor specific logic in netmiko allows the responses received from each device to be parsed according to the device type (i.e Juniper devices are parsed differently than Arista)

Context managers allow us to make the connection establishment process even simpler.  Establishing a connection to a device consists of the following workflow:

1. Create an object representing the device we are going to connect to. This “network device” object contains attributes such as IP address, SSH port, username, password, and hostname. For example, you can read in the data from a json file containing your network device data, create a dictionary with key-value pairs,  read it from a database, etc. The end result being that we create an object with the attributes required to connect to it.
2. Instantiate a netmiko/paramiko object to which we pass in the network device object.  The netmiko/paramiko object will use the network device attributes to establish an SSH session to the device.
3. Establish the SSH session.
4. Take actions on the device (send commands, parse output, etc)
5. Disconnect from the device.

The code for the workflow described would look something like this:

This code is simple, and it works well for connecting to a single device from the Python interpreter shell. But what if you wanted to connect to hundreds of devices? What if you were writing a library with several tools which all need to the same basic function of connecting to a device via SSH? The code would start to get very repetitive if this same routine needs to be done in several places within the code base. What about error handling? You would not want your script that is processing hundreds of devices to crash due to an unhanded error, or have to have the same repetitive error handling code all over the place.

Here is where context managers come into play! Does this syntax look familiar?

The “open” builtin function Python was implemented to support the Context Manager protocol (called using “with”). Open can also be called in the traditional way, however we need to close the file handle manually after we are done.

However, if there is an exception that occurs before the close() call, the program will terminate and the file handle will remain open. To take care of that we would need to use a try…except…finally block so that the file handle is closed regardless of how the program terminates (successfully or with errors).

This gets ugly, repetitive, and easy to forget. Thankfully context managers take care of this administrative logic for us.  When we call open using “with”, we are calling open’s context manager. No matter what happens in the body of the “with” block, the file handle will always get closed.  Open’s context manager takes care of this for us, without us having to manually handle it in our code.

Context managers can be used for anything that requires a setup and teardown logic  in-between the actual “body” of the code.  What is our end goal when we open a file? We either want to read the contents of the file or write something to the file.  Our code will look much cleaner if we can separate the setup/teardown code (opening the file, handling errors, closing the file) from the “doing something” code (reading or writing to the file).

This is perfect for our SSH connection logic. We can make an SSH context manager that takes care of doing all of the session setup and teardown, and we can ensure that the SSH session will be closed even if exceptions are raised.  Opening an SSH session can be thought of as opening a file handle, or a socket. Our setup logic consists of building the netmiko/paramiko object, passing it our network_device object, and establishing the connection (opening the socket). If there is an error, the setup logic should be able to handle it. The teardown logic consists of closing the SSH session (i,e cleaning up, closing the socket).

The Python standard library ships with the “contextlib” module, which has tools that make it very simple to make your own context managers. We make use of the “contextmanager” decorator to “wrap” our soon to be context manager function.  The “ssh_manager” function presented above takes care of creating the required netmiko SSH object, connecting to the end device, and handling setup errors, and closing the session. We can call our context manager as follows:

When the first line of this code runs :

the ssh_manager context manager function is triggered.  The ssh_manager function will run up until the “yield” statement, yielding the netmiko SSH session handle. The SSH handle is stored in the “conn” variable (as conn) Now the body of the with block will run, which in our case we are sending the “show ip route” command to the network device and printing it. As soon as the with block finishes execution, control is passed back to the the ssh_manager function right after the yield statement. The rest of the ssh_manager function completes (teardown code).

Now the code in all our of functions/scripts can simply call the “ssh_manager” context manager and pass it in a network_device object with the required attributes in order to establish a connection. This greatly simplifies the rest of our code and makes it re-usable.

More information about context managers can be found on PEP 343