db_client_with_ssh_tunnel.py

import logging

from pymongo import MongoClient
from sshtunnel import SSHTunnelForwarder


# local
SSH_PASSPHRASE = 'you ssh passphrase (if any)'
LOCAL_BIND_PORT = 27017
SSH_KEYPATH = '/your/path/to/key/file'
SSH_USERNAME = '...'

# remote server (bastion host)
SSH_REMOTE_IP = '...'
SSH_REMOTE_PORT = 27018

# private server
SSH_PRIVATE_IP = '...'
SSH_PRIVATE_PORT = 27019

# DB credential
DB_CRED = {"username": '...', "password": '...', 'host': f'127.0.0.1:{LOCAL_BIND_PORT}'}
DB_NAME = '...'


class MongoClientWithSSHTunnel:
    def __init__(self) -> None:
        self.logger = logging.getLogger(__name__)
        self.client = None
        self.ssh_tunnel_forwarder = None

    def create_ssh_tunnel_forwarder(self):
        self.ssh_tunnel_forwarder = SSHTunnelForwarder(
            (SSH_REMOTE_IP, SSH_REMOTE_PORT),
            ssh_username=SSH_USERNAME,
            ssh_pkey=SSH_KEYPATH,
            ssh_private_key_password=SSH_PASSPHRASE,
            remote_bind_address=(
                SSH_PRIVATE_IP,
                SSH_PRIVATE_PORT,
            ),
            local_bind_address=('0.0.0.0', LOCAL_BIND_PORT),
        )

    def start_ssh_tunneling(self) -> None:
        if not self.ssh_tunnel_forwarder:
            self.logger.info("SSH tunnel forwarder not exists")
            return
        try:
            self.ssh_tunnel_forwarder.start()
        except Exception as e:
            self.logger.error(f'something went wrong when start ssh tunneling, err={e}')

    def get_db(self):
        mongo_uri = "mongodb://{username}:{password}@{host}".format(**DB_CRED)
        self.client = MongoClient(mongo_uri)
        db = None
        try:
            self.client.server_info()
            db = self.client[DB_NAME]
        except Exception as e:
            self.logger.error(f'something went wrong when getting the db instance, err={e}')
        return db

    def start(self):
        self.create_ssh_tunnel_forwarder()
        self.start_ssh_tunneling()

    def stop(self):
        try:
            self.client.close()
            self.ssh_tunnel_forwarder.stop()
        except Exception as e:
            self.logger.error(f'something went wrong when closing the connection, err={e}')


if __name__ == "__main__":
    client = MongoClientWithSSHTunnel()

    client.start()

    # should be able to connect to the private server
    db = client.get_db()
    if db:
        print(db.list_collection_names()) 

    client.stop()