File size: 1,650 Bytes
cf9d671
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
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