Pytest mock ssh device with interactive shell based on asyncssh

[evgenoir]

For testing purposes, I made a server based on asyncssh that emulates the interactive shell of network devices. The server itself works fine if you run it separately, but does not work if you try to run it in the same event loop with the asyncssh client.

import sys
import asyncio
import asyncssh
from typing import Optional


class MockSshDevice(asyncssh.SSHServer):
    """ Create mock SSH-device """
    def __init__(self, port: int = 8022):
        self._port = port
        self._prompt = ""
        self._cmd_resp_map = {}
        self._ssh_acceptor = None

    def connection_made(self, conn: asyncssh.SSHServerConnection):
        print(f"SSH connection received from {conn.get_extra_info('peername')[0]}")

    def connection_lost(self, exc: Optional[Exception]):
        if exc:
            print('SSH connection error: ' + str(exc), file=sys.stderr)
        else:
            print('SSH connection closed')

    def begin_auth(self, username: str) -> bool:
        # Without authentication
        return False

    @property
    def port(self) -> int:
        """ Get used SSH-port """
        return self._port

    @property
    def prompt(self) -> str:
        """ Get device prompt """
        return self._prompt

    @prompt.setter
    def prompt(self, value: str):
        """ Set device prompt """
        if type(value) is not str:
            raise RuntimeError("MockSshDevice.prompt must be of type string")

        self._prompt = value

    @property
    def cmd_resp_map(self) -> dict:
        """ Get command to response map """
        return self._cmd_resp_map

    @cmd_resp_map.setter
    def cmd_resp_map(self, value: dict):
        """ Set command to response map """
        if type(value) is not dict:
            raise RuntimeError("MockSshDevice.cmd_resp_map must be of type dict")
        # All keys must be non-empty strings
        if not all(value.keys()):
            raise RuntimeError("MockSshDevice.cmd_resp_map must not contain empty commands")

        # Remove leading and trailing new lines & spaces in commands & their responses
        value = {cmd.strip(): response.strip() for cmd, response in value.items()}

        self._cmd_resp_map = value

    async def handle_request(self, process: asyncssh.SSHServerProcess):
        """ Handle incoming SSH-request """
        if self.prompt:
            process.stdout.write(self.prompt)

        try:
            async for command in process.stdin:
                command = command.rstrip('\n')
                if command:
                    # Handle supported command
                    if command in self.cmd_resp_map:
                        # Handle command with non-empty output
                        if self.cmd_resp_map[command]:
                            process.stdout.write(self.cmd_resp_map[command] + '\n')
                    # Handle unsupported command
                    else:
                        process.stdout.write("Unsupported command\n")

                    if self.prompt:
                        process.stdout.write(self.prompt)

        except asyncssh.BreakReceived:
            process.stdout.write('\n')
            process.exit(0)

    async def start(self):
        """ Start mock SSH-device """
        self._ssh_acceptor = await asyncssh.create_server(
            MockSshDevice,
            '',
            self.port,
            server_host_keys=['keys/private_key'],
            process_factory=self.handle_request
        )

    async def stop(self):
        """ Stop mock SSH-device """
        self._ssh_acceptor.close()
        await self._ssh_acceptor.wait_closed()

    async def __aenter__(self):
        await self.start()

    async def __aexit__(self, exc_type, exc, tb):
        await self.stop()


async def main():
    device = MockSshDevice()
    device.prompt = "SomePrompt#"
    device.cmd_resp_map = {
        "terminal length 0": "",
        "show version": """
Cisco IOS Software, c6848x Software (c6848x-ADVENTERPRISEK9_NPE-M), Version 15.5(1)SY2, RELEASE SOFTWARE (fc6)
Technical Support: http://www.cisco.com/techsupport
Copyright (c) 1986-2018 by Cisco Systems, Inc.
Compiled Fri 14-Sep-18 21:41 by prod_rel_team
"""
    }

    async with device:
        async with asyncssh.connect('localhost', 8022, known_hosts=None) as conn:
            result = await conn.run("terminal length 0", check=True)
            print(result.stdout)
            result = await conn.run("show version", check=True)
            print(result.stdout)


if __name__ == '__main__':
    asyncio.run(main())

The client code blocks on the first await conn.run(...) while the server continues to run, which can be verified by running ssh localhost -p 8022 on the Linux command line. This makes me very sad, because I want to use this server as a fixture in pytest and want to use it like this:

@pytest.fixture
def device():
    yield MockSshDevice()


@pytest.mark.asyncio
async def test(device):
    device.prompt = "SomePrompt#"
    device.cmd_resp_map = {
        "show version": """
Cisco IOS Software, c6848x Software (c6848x-ADVENTERPRISEK9_NPE-M), Version 15.5(1)SY2, RELEASE SOFTWARE (fc6)
Technical Support: http://www.cisco.com/techsupport
Copyright (c) 1986-2018 by Cisco Systems, Inc.
Compiled Fri 14-Sep-18 21:41 by prod_rel_team
""",
        "terminal length 0": ""
    }

    async with device:
        async with asyncssh.connect('localhost', 8022, known_hosts=None) as conn:
            result = await conn.run("terminal length 0", check=True)
            assert result.stdout == "..."
            result = await conn.run("show version", check=True)
            assert result.stdout == "..."

Could you tell me how to overcome this difficulty and achieve the desired behavior?

[ronf]

Hello...

I see a couple of potential problems with this code that may explain why you're getting blocked here.

First, your server handle_request() code appears to be reading commands from stdin but your client code is passing the commands in as a "command" argument (creating an SSH "exec" session, not a "shell" session). The client code is not writing anything to stdin and not closing stdin, so the server code is going to block forever waiting for input on stdin. Also, note that by using a conn.run() in your client, each command you run will be executed as a separate SSH session, with its own stdin/stdout/stderr. Normally, a client would open a single SSH session (usually a "shell" session) if it plans on running multiple consecutive commands. In fact, it would need to do this for something like "terminal length 0" to be applied correctly, since that command only affects the terminal session it is run on.

To fix this, you'd want to use conn.create_process() instead of conn.run() on the client side and write your commands to process.stdin (with a newline at the end) and not pass the "command" argument. If you want to wait on the process closing on the client side, you'll also want to make sure that you write EOF to stdin before doing so, or explicitly close the client process when done writing commands. You could also do this by using the client process object you get back as a context manager when you replace conn.run() with conn.create_process().

Second, once this is fixed, the handle_request() function exits without doing a process.close(), or closing stdout/stderr explicitly, so the client code would end up blocking waiting for the server process to finish sending output and it never will. Adding this explicit close at the bottom of the function (or in a "try..finally" block) should solve this issue.

I don't have any direct experience with using pytest fixtures with asyncio code, so there may be additional things to consider there. However, I'd start with fixing the above issues and see how far you get. You may also find the info in https://tonybaloney.github.io/posts/async-test-patterns-for-pytest-and-unittest.html to be useful.

[evgenoir]

This is definitely the fastest answer I've ever seen on github ) I followed your advice and now everything works as expected:

import re
import sys
import socket
import asyncio
import asyncssh
from typing import Optional


class MockSshDevice(asyncssh.SSHServer):
    """ Create mock SSH-device """
    def __init__(self, port: int = 8022):
        self._port = port
        self._prompt = ""
        self._cmd_resp_map = {}
        self._ssh_acceptor = None

    def connection_made(self, conn: asyncssh.SSHServerConnection):
        print(f"SSH connection received from {conn.get_extra_info('peername')[0]}")

    def connection_lost(self, exc: Optional[Exception]):
        if exc:
            print('SSH connection error: ' + str(exc), file=sys.stderr)
        else:
            print('SSH connection closed')

    def begin_auth(self, username: str) -> bool:
        # Without authentication
        return False

    @property
    def port(self) -> int:
        """ Get used SSH-port """
        return self._port

    @property
    def prompt(self) -> str:
        """ Get device prompt """
        return self._prompt

    @prompt.setter
    def prompt(self, value: str):
        """ Set device prompt """
        if type(value) is not str:
            raise RuntimeError("MockSshDevice.prompt must be of type string")

        self._prompt = value

    @property
    def cmd_resp_map(self) -> dict:
        """ Get command to response map """
        return self._cmd_resp_map

    @cmd_resp_map.setter
    def cmd_resp_map(self, value: dict):
        """ Set command to response map """
        if type(value) is not dict:
            raise RuntimeError("MockSshDevice.cmd_resp_map must be of type dict")
        # All keys must be non-empty strings
        if not all(value.keys()):
            raise RuntimeError("MockSshDevice.cmd_resp_map must not contain empty commands")

        # Remove leading and trailing new lines & spaces in commands & their responses
        value = {cmd.strip(): response.strip() for cmd, response in value.items()}

        self._cmd_resp_map = value

    async def handle_request(self, process: asyncssh.SSHServerProcess):
        """ Handle incoming SSH-request """
        if self.prompt:
            process.stdout.write(self.prompt)

        try:
            async for command in process.stdin:
                command = command.rstrip('\n')
                if command:
                    # Handle supported command
                    if command in self.cmd_resp_map:
                        # Handle command with non-empty output
                        if self.cmd_resp_map[command]:
                            process.stdout.write(self.cmd_resp_map[command] + '\n')
                    # Handle unsupported command
                    else:
                        process.stdout.write("Unsupported command\n")

                    if self.prompt:
                        process.stdout.write(self.prompt)

        except asyncssh.BreakReceived:
            process.stdout.write('\n')
            process.exit(0)

        finally:
            process.close()

    async def start(self):
        """ Start mock SSH-device """
        self._ssh_acceptor = await asyncssh.create_server(
            MockSshDevice,
            '',
            self.port,
            family=socket.AF_INET,  # IPv4 socket
            server_host_keys=['keys/private_key'],
            process_factory=self.handle_request
        )

    async def stop(self):
        """ Stop mock SSH-device """
        self._ssh_acceptor.close()
        await self._ssh_acceptor.wait_closed()

    async def __aenter__(self):
        await self.start()

    async def __aexit__(self, exc_type, exc, tb):
        await self.stop()


async def main():
    device = MockSshDevice()
    device.prompt = "SomePrompt#"
    device.cmd_resp_map = {
        "terminal length 0": "",
        "show version": """
Cisco IOS Software, c6848x Software (c6848x-ADVENTERPRISEK9_NPE-M), Version 15.5(1)SY2, RELEASE SOFTWARE (fc6)
Technical Support: http://www.cisco.com/techsupport
Copyright (c) 1986-2018 by Cisco Systems, Inc.
Compiled Fri 14-Sep-18 21:41 by prod_rel_team
"""
    }

    async with device:
        async with asyncssh.connect('localhost', 8022, known_hosts=None) as conn:
            async with conn.create_process() as process:
                result = await process.stdout.readuntil(device.prompt)
                assert result == device.prompt
                process.stdin.write("terminal length 0\n")
                result = await process.stdout.readuntil(device.prompt)
                assert result == device.prompt
                process.stdin.write("show version\n")
                result = await process.stdout.readuntil(device.prompt)
                assert re.search(rf"Cisco IOS Software.*{device.prompt}", result, re.DOTALL)


if __name__ == '__main__':
    asyncio.run(main())

Thank you!

[ronf]

My pleasure - glad to hear you got it working!

[evgenoir]

The joy was premature (( I tried to use paramiko.SSHClient instead asyncssh client with code above and and it didn't work:

async def main():
    device = MockSshDevice()
    device.prompt = "DevicePrompt#"
    device.command_sequence = [
        ("terminal length 0", device.prompt),
        ("show version", f"""
    Cisco IOS Software, c6848x Software (c6848x-ADVENTERPRISEK9_NPE-M), Version 15.5(1)SY2, RELEASE SOFTWARE (fc6)
    Technical Support: http://www.cisco.com/techsupport
    Copyright (c) 1986-2018 by Cisco Systems, Inc.
    Compiled Fri 14-Sep-18 21:41 by prod_rel_team
    {device.prompt}
    """),
        ("exit", device.TERMINATE_SHELL)
    ]

    async with device:
        ssh_client = paramiko.SSHClient()
        ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        ssh_client.connect('localhost', port=8022, username='admin', password='P@ssw0rd')  # exception here
        ssh_client.close()

I keep getting the exception:

Traceback (most recent call last):
  File "/home/admin/Projects/PyCharmProjects/mockSSH_test/asyncssh/question2.py", line 30, in main
    ssh_client.connect('localhost', port=8022, username='admin', password='P@ssw0rd')
  File "/home/admin/Projects/PyCharmProjects/mockSSH_test/venv/lib/python3.10/site-packages/paramiko/client.py", line 380, in connect
    raise NoValidConnectionsError(errors)
paramiko.ssh_exception.NoValidConnectionsError: [Errno None] Unable to connect to port 8022 on 127.0.0.1 or ::1

At the same time, if I put a breakpoint on the instruction ssh_client.connect('localhost', port=8022, username='admin', password='P@ssw0rd') and run the command netstat -aon | grep 8022 in the OS terminal, I see that the server is running and accepting connections. I suspect this is because synchronous and asynchronous code is mixed here, but I have no idea how to overcome this limitation and make everything work. I need this scenario to work too!

[ronf]

Yes - paramiko is threaded, and it's generally a bad idea to mix threaded and async code in the same Python interpreter. Is there a reason you weren't able to use AsyncSSH for the SSH client here?

From the error message, it looks like you were trying to use something called mockSSH in PyCharm, but is it a requirement to use that?

One possible solution would be to run the server in a separate process (with its own Python interpreter). As long as you don't need to communicate with that new process once it is started (other than to stop it), that should be fairly simple. I've never tried to do this as part of a pytest fixture, but I imagine it would be no different than using "subprocess" to run any other command. Just make sure to use the main "subprocess" module and not the one in asyncio. At that point, all the async code would be running in the subprocess.

Another option might be to run the threaded paramiko code using "run_in_executor", feeding the results back to an asyncio main thread. If you're bringing in pytest here, you'll need to make sure that this is done as async test code with only calls to asyncio, with anything threaded run from that asyncio event loop via an executor.

[evgenoir]

In general, I did as you said - launched my mock device in a separate process and this turned out to be the easiest solution.
As for using asyncssh instead of paramiko, I thought about it but I really miss the ability to wait for a command line prompt using a regex, and not a set of lines as it's done in SSHReader.readuntil().
In any case, I am very grateful to you for your advice.

[ronf]

Nothing requires that you read data a line at a time in SSHReader. You can do a plain read() instead of readline() or using an iterator, and then do whatever matching you want to do in the blocks of data you get back (including regex matches). You just need to specify a length in the read() call so that it doesn't wait indefinitely for data to arrive, and you may find that whatever you're searching for could be split across a read() boundary. That's what readuntil() is designed to solve, but there's no way to do incremental regex matching in Python efficiently, at least not with the standard "re" module. So, readuntil() only supports a list of static strings to match against. It uses "re" under the covers to speed that up, but since the strings are all fixed-size, there are ways to make that specific regex matching more efficient.

If you know the prompt will never appear in other command output, a call to readuntil() which specifically looks for the prompt should work well. You may also be able to send commands that output a specific "marker" line and search for that, to better guarantee that what you match on is unique, but that approach depends on whether the target device will let you queue up additional commands before you have read all of the previous command output.