Spaces:
Running
Running
| """Keep a local PostgreSQL SSH tunnel open without executing remote commands.""" | |
| from __future__ import annotations | |
| import os | |
| import select | |
| import socketserver | |
| import threading | |
| import paramiko | |
| SSH_HOST = os.getenv("EGODATA_SSH_HOST", "192.168.110.41") | |
| SSH_PORT = int(os.getenv("EGODATA_SSH_PORT", "22")) | |
| SSH_USER = os.getenv("EGODATA_SSH_USER", "Stouch") | |
| SSH_PASSWORD = os.environ.get("EGODATA_SSH_PASSWORD", "") | |
| LOCAL_HOST = os.getenv("EGODATA_DB_LOCAL_HOST", "127.0.0.1") | |
| LOCAL_PORT = int(os.getenv("EGODATA_DB_LOCAL_PORT", "15432")) | |
| REMOTE_HOST = os.getenv("EGODATA_DB_REMOTE_HOST", "172.19.0.2") | |
| REMOTE_PORT = int(os.getenv("EGODATA_DB_REMOTE_PORT", "5432")) | |
| class ForwardServer(socketserver.ThreadingTCPServer): | |
| allow_reuse_address = True | |
| daemon_threads = True | |
| def __init__(self, server_address, handler_class, transport): | |
| super().__init__(server_address, handler_class) | |
| self.transport = transport | |
| class Handler(socketserver.BaseRequestHandler): | |
| def handle(self): | |
| peer = self.request.getpeername() | |
| channel = self.server.transport.open_channel( | |
| "direct-tcpip", (REMOTE_HOST, REMOTE_PORT), peer | |
| ) | |
| if channel is None: | |
| raise OSError("SSH server refused the PostgreSQL forwarding channel") | |
| try: | |
| sockets = [self.request, channel] | |
| while sockets: | |
| readable, _, _ = select.select(sockets, [], [], 1.0) | |
| for source in readable: | |
| payload = source.recv(64 * 1024) | |
| if not payload: | |
| sockets.clear() | |
| break | |
| target = channel if source is self.request else self.request | |
| target.sendall(payload) | |
| finally: | |
| channel.close() | |
| def main() -> None: | |
| if not SSH_PASSWORD: | |
| raise SystemExit("EGODATA_SSH_PASSWORD is required") | |
| client = paramiko.SSHClient() | |
| client.set_missing_host_key_policy(paramiko.AutoAddPolicy()) | |
| client.connect( | |
| SSH_HOST, | |
| port=SSH_PORT, | |
| username=SSH_USER, | |
| password=SSH_PASSWORD, | |
| look_for_keys=False, | |
| allow_agent=False, | |
| timeout=15, | |
| banner_timeout=15, | |
| auth_timeout=15, | |
| ) | |
| transport = client.get_transport() | |
| if transport is None or not transport.is_active(): | |
| client.close() | |
| raise SystemExit("SSH transport is not active") | |
| server = ForwardServer((LOCAL_HOST, LOCAL_PORT), Handler, transport) | |
| print( | |
| f"PostgreSQL tunnel: {LOCAL_HOST}:{LOCAL_PORT} -> " | |
| f"{REMOTE_HOST}:{REMOTE_PORT} via {SSH_USER}@{SSH_HOST}", | |
| flush=True, | |
| ) | |
| try: | |
| server.serve_forever(poll_interval=0.5) | |
| finally: | |
| server.server_close() | |
| client.close() | |
| if __name__ == "__main__": | |
| main() | |