Spaces:
Paused
Paused
| import os | |
| import time | |
| import subprocess | |
| import socket | |
| import threading | |
| import paramiko | |
| import io | |
| import logging | |
| # Configure Logging | |
| logging.basicConfig(level=logging.INFO) | |
| logger = logging.getLogger("azure_bridge") | |
| class AzureSFTPBridge: | |
| def __init__(self, hostname="ssh.my-robot.dev", local_port=2222, username="azureuser"): | |
| self.hostname = hostname | |
| self.local_port = local_port | |
| self.username = username | |
| self.tunnel_process = None | |
| self.ssh_key = None | |
| self._load_key() | |
| def _load_key(self): | |
| """Load SSH Key from Env Var""" | |
| key_content = os.environ.get("SSH_KEY") | |
| if not key_content: | |
| logger.error("SSH_KEY environment variable not found!") | |
| return | |
| # Determine key type and load | |
| try: | |
| self.ssh_key = paramiko.RSAKey.from_private_key(io.StringIO(key_content)) | |
| logger.info("Loaded RSA Key successfully.") | |
| except Exception: | |
| try: | |
| self.ssh_key = paramiko.Ed25519Key.from_private_key(io.StringIO(key_content)) | |
| logger.info("Loaded Ed25519 Key successfully.") | |
| except Exception as e: | |
| logger.error(f"Failed to load SSH Key: {e}") | |
| def start_tunnel(self): | |
| """Starts cloudflared access tcp in background.""" | |
| logger.info(f"Starting Cloudflare Tunnel to {self.hostname} on localhost:{self.local_port}...") | |
| cmd = [ | |
| "cloudflared", | |
| "access", "tcp", | |
| "--hostname", self.hostname, | |
| "--url", f"tcp://localhost:{self.local_port}" | |
| ] | |
| self.tunnel_process = subprocess.Popen( | |
| cmd, | |
| stdout=subprocess.PIPE, | |
| stderr=subprocess.PIPE | |
| ) | |
| # Wait for port to be open | |
| retries = 10 | |
| for i in range(retries): | |
| if self._check_port(): | |
| logger.info("Tunnel established successfully.") | |
| return True | |
| time.sleep(1) | |
| logger.error("Failed to establish tunnel.") | |
| return False | |
| def _check_port(self): | |
| try: | |
| with socket.create_connection(("localhost", self.local_port), timeout=1): | |
| return True | |
| except (socket.timeout, ConnectionRefusedError): | |
| return False | |
| def get_sftp(self): | |
| """Connects via Paramiko and returns SFTPClient.""" | |
| if not self.ssh_key: | |
| raise ValueError("SSH Key not loaded.") | |
| if not self._check_port(): | |
| logger.warning("Tunnel port closed. Restarting tunnel...") | |
| self.start_tunnel() | |
| try: | |
| logger.info("Connecting to SFTP via Tunnel...") | |
| transport = paramiko.Transport(("localhost", self.local_port)) | |
| transport.connect(username=self.username, pkey=self.ssh_key) | |
| sftp = paramiko.SFTPClient.from_transport(transport) | |
| logger.info("SFTP Connection Established.") | |
| return sftp | |
| except Exception as e: | |
| logger.error(f"SFTP Connection Failed: {e}") | |
| raise | |
| # --- High Level API --- | |
| def list_files(self, remote_path): | |
| with self.get_sftp() as sftp: | |
| return sftp.listdir(remote_path) | |
| def upload_file(self, local_path, remote_path): | |
| with self.get_sftp() as sftp: | |
| sftp.put(local_path, remote_path) | |
| logger.info(f"Uploaded {local_path} -> {remote_path}") | |
| def download_file(self, remote_path, local_path): | |
| with self.get_sftp() as sftp: | |
| sftp.get(remote_path, local_path) | |
| logger.info(f"Downloaded {remote_path} -> {local_path}") | |