src/tunnel.py

import io
import os

import paramiko
import streamlit as st
from sshtunnel import SSHTunnelForwarder


@st.cache_resource()
def start_server():
    server = SSHTunnelManager()
    server.start()
    return server


class SSHTunnelManager:
    def __init__(self):
        pkey = paramiko.RSAKey.from_private_key(io.StringIO(os.getenv('PRIVATE_KEY')))
    
        self.server = SSHTunnelForwarder(
            ssh_address_or_host=os.getenv('SERVER_DNS_NAME'),
            ssh_username="ec2-user",
            ssh_pkey=pkey,  
            local_bind_address=("127.0.0.1", int(os.getenv('LOCAL_PORT'))),
            remote_bind_address=("127.0.0.1", 8000)
        )

        self._is_starting = False
        self._is_running = False

    def update_status(self):
        if not self._is_starting:
            self.server.check_tunnels()
            self._is_running = list(self.server.tunnel_is_up.values())[0]
        else:
            self._is_running = False

    def is_starting(self):
        self.update_status()
        return self._is_starting

    def is_running(self):
        self.update_status()
        return self._is_running

    def is_down(self):
        self.update_status()
        return (not self._is_running) and (not self._is_starting)

    def start(self, *args, **kwargs):
        if not self._is_starting:
            self._is_starting = True
            self.server.start(*args, **kwargs)
            self._is_starting = False

    def restart(self, *args, **kwargs):
        if not self._is_starting:
            self._is_starting = True
            self.server.restart(*args, **kwargs)
            self._is_starting = False