Secure Access with Paramiko: A First Step Towards Data Democratization

At Privacy Dynamics, the privacy and security of our customers' data is the core of our mission. Our goal is to help them activate the data they need while minimizing PII sprawl. As such, read and write functions putting data in transit must be secure. Our latest feature, SSH Tunnel-enabled connections, ensures secure transit and is now available for our cloud customers.

Implementing the SSH tunnel support was challenging, as we needed a reliable, simple, and, importantly, testable solution. Using out-of-band SSH implementations (e.g. a Docker container) is an option but makes visibility under test challenging. There are some libraries available for exposing mock SSH servers to pytest, but they are designed for mocking an SSH server you send commands to and receive a response (e.g. `ls -la). There are no purpose-built tools for mocking an SSH server to tunnel connections through.

For this feature, we were able to leverage the Paramiko Python package. In this article, we describe our implementation, detailing how Paramiko was used to build a custom SSH server, integrating a spyable mock, and seamlessly weaving it into our server to test SSH tunnels.

Harnessing Paramiko for SSH Server

Our custom-built SSH server, SSHAdapter``, uses Paramiko for user authentication and managing TCP/IP requests, creating a simple SSH tunneling TCP/IP proxy. The flexibility of the Paramiko library allowed us to leverage its ServerInterface to build our SSH server, as well as manage user authentication and TCP/IP requests. Here’s how ourSSHAdapter gets initialized:

class SSHAdapter(paramiko.ServerInterface):
    def __init__(self):
        self.event = threading.Event()
        self.handle_tcpip_request = Mock(side_effect=self._handle_tcpip_request)

Implementing Paramiko's ServerInterface, our server authenticates users through both password-based and public key methods:

def check_auth_password(self, username, password):
    if (username == ssh_server_username) and (password == ssh_server_password):
        return paramiko.AUTH_SUCCESSFUL
    return paramiko.AUTH_FAILED
def check_auth_publickey(self, username, key):
header, key_string, comment = TestConfig.ssh_tunnel_public_key.split(" ")
public_key = paramiko.RSAKey(data=base64.b64decode(key_string))
if (username == ssh_server_username) and (key == public_key):
return paramiko.AUTH_SUCCESSFUL
return paramiko.AUTH_FAILED

And most importantly, we use ServerInterface to implement a proxy for TCP requests between clients and servers. check_channel_direct_tcpip_request() is implemented by spawning a new thread to handle a new SSH session:

def check_channel_direct_tcpip_request(self, chanid, origin, destination):
    forward_thread = threading.Thread(
        target=self.handle_tcpip_request, args=(chanid, origin, destination)
    )
    forward_thread.start()
    return paramiko.OPEN_SUCCEEDED

def _handle_tcpip_request(self, chanid, origin, destination):
    transport = self.transport
    while True:
        chan = transport.accept(1000)
        if chan is None:
            continue
        logging.info(f"Received connection from {origin}")
        break
    remote_host, remote_port = destination
    try:
        remote_sock = socket.create_connection((remote_host, remote_port))
    except Exception as e:
        logging.error(
            f"Unable to establish connection to {remote_host}:{remote_port}: {str(e)}"
        )
        return
    logging.info(f"Connected to {remote_host}:{remote_port}")
    while True:
        r, w, x = select.select([chan, remote_sock], [], [])
        if chan in r:
            x = chan.recv(1024)
            if len(x) == 0:
                break
            remote_sock.send(x)
        if remote_sock in r:
            x = remote_sock.recv(1024)
            if len(x) == 0:
                break
            chan.send(x)
    chan.close()
    remote_sock.close()

We've integrated the spyable mock into our testing framework directly inside the SSHAdapter class. The handle_tcpip_request method is the mock that helps us track and control the tunnelling calls:

self.handle_tcpip_request = Mock(side_effect=self._handle_tcpip_request)

This strategy, enabled by Paramiko, provides a layer of abstraction, simplifying our tests and providing insights into the number of times the TCP/IP request being called.

Running an In-Process Server as a Pytest Fixture

Our Paramiko-powered SSH server operates within the same process but on a different thread, allowing it to run concurrently with other processes without requiring separate deployment.

def start_ssh_server(future):
    ssh_adapter = SSHAdapter()
    server = SSHServer(
        ssh_adapter,
        ssh_server_host,
        ssh_server_port,
        ssh_server_username,
        ssh_server_password,
    )
    future.set_result(server)
    server.start()
server_thread = threading.Thread(
target=start_ssh_server, args=(future,), daemon=True
)
server_thread.start()

Further, by making the server available as a pytest fixture, we facilitate its use across different tests, fostering a flexible and reusable testing environment:

@pytest.fixture(scope="session", autouse=True)
def ssh_server():
    future = concurrent.futures.Future()
    server_thread = threading.Thread(
        target=start_ssh_server, args=(future,), daemon=True
    )
    server_thread.start()
    server = future.result()  # Get the SSHServer instance from the future
    yield server

Here is a simple pytest that uses the ssh_server() fixture to ensure the connection to the SSH server is working correctly:

from sshtunnel import SSHTunnelForwarder
def test_ssh_tunneling(ssh_server):
    with SSHTunnelForwarder(
      (ssh_server.host, ssh_server.port),
      ssh_username=ssh_server.username,
      ssh_password=ssh_server.password,
      remote_bind_address=("localhost", 5432), # Destination host and port
    ) as tunnel:
      object.network_call_through_tunnel()
      assert ssh_server.ssh_adapter.handle_tcpip_request.call_count == 1

In this test, we use Paramiko to connect to the SSH server and assert that the connection was successful. If there is an exception during the connection, the test will fail.

In Summary

Thorough testing of our Paramiko-driven SSH tunneling feature is a cornerstone of our commitment to data security at Privacy Dynamics. We scrutinize every step of our process to ensure data remains secure, including new connections, request forwarding, and data flow handling. These steps provide a protective layer safeguarding customer data regardless of where that data is in the job process.

Paramiko has proven to be a key component in achieving our data security objectives. It enabled us to establish a robust data security architecture with a custom SSH server, integrated spyable mock, in-process server, and comprehensive testing on SSH tunneling.

If you would like to learn more about setting up secure data pipelines between Prod and your developer environments, we’d love to chat. You can book some time with us here.