use std::{ env, fs, io::{self}, net::{SocketAddr, TcpStream}, os::unix::{fs::PermissionsExt, net::UnixStream}, path::{Path, PathBuf}, process::{Command, Stdio}, sync::{Arc, Mutex}, thread, time::{Duration, Instant}, }; use serde::{Deserialize, Serialize}; use time::{OffsetDateTime, format_description::well_known::Rfc3339}; use uuid::Uuid; use crate::{ commands, session_manager::{INSTANCE_DIR_NAME, INSTANCE_FILENAME}, vm::{self, LoginAction}, }; const SSH_KEY_NAME: &str = "ssh_key"; pub(crate) const VM_ROOT_LOG_NAME: &str = "vm_root.log"; pub(crate) const STATUS_FILE_NAME: &str = "status.txt"; pub(crate) const DEFAULT_SSH_USER: &str = "vibecoder"; const SSH_CONNECT_RETRIES: usize = 30; const SSH_CONNECT_DELAY_MS: u64 = 500; const SSH_SETUP_SCRIPT: &str = include_str!("ssh.sh"); #[derive(Debug, Clone, Serialize, Deserialize)] pub(crate) struct InstanceConfig { #[serde(default)] id: String, #[serde(default = "default_ssh_user")] ssh_user: String, #[serde(default)] sudo_password: String, #[serde(default)] last_active: Option, #[serde(default)] pub(crate) vm_ipv4: Option, } impl InstanceConfig { pub(crate) fn ssh_user_display(&self) -> String { if self.ssh_user.trim().is_empty() { DEFAULT_SSH_USER.to_string() } else { self.ssh_user.clone() } } } fn default_ssh_user() -> String { DEFAULT_SSH_USER.to_string() } pub fn run_with_ssh(manager_conn: UnixStream) -> Result<(), Box> { let project_root = env::current_dir()?; tracing::info!(root = %project_root.display(), "starting ssh session"); let instance_dir = ensure_instance_dir(&project_root)?; tracing::debug!(instance_dir = %instance_dir.display(), "instance dir ready"); let (ssh_key, _ssh_pub) = ensure_ssh_keypair(&instance_dir)?; let config = load_or_create_instance_config(&instance_dir)?; let ssh_user = config.ssh_user.clone(); tracing::debug!(ssh_user = %ssh_user, "loaded instance config"); let _manager_conn = manager_conn; wait_for_vm_ipv4(&instance_dir, Duration::from_secs(480))?; let ip = load_or_create_instance_config(&instance_dir)? .vm_ipv4 .ok_or("VM IPv4 not available")?; tracing::info!(ip = %ip, "vm ipv4 ready"); run_ssh_session(ssh_key, ssh_user, ip) } pub fn ensure_instance_dir(project_root: &Path) -> Result { let instance_dir = project_root.join(INSTANCE_DIR_NAME); fs::create_dir_all(&instance_dir)?; Ok(instance_dir) } pub(crate) fn ensure_ssh_keypair( instance_dir: &Path, ) -> Result<(PathBuf, PathBuf), Box> { let private_key = instance_dir.join(SSH_KEY_NAME); let public_key = instance_dir.join(format!("{SSH_KEY_NAME}.pub")); if private_key.exists() && public_key.exists() { return Ok((private_key, public_key)); } if private_key.exists() { let _ = fs::remove_file(&private_key); } if public_key.exists() { let _ = fs::remove_file(&public_key); } let status = Command::new("ssh-keygen") .args([ "-t", "ed25519", "-N", "", "-f", private_key.to_str().ok_or("ssh key path not utf-8")?, "-C", "vibebox", ]) .stdin(Stdio::null()) .stdout(Stdio::inherit()) .stderr(Stdio::inherit()) .status()?; if !status.success() { return Err("ssh-keygen failed".into()); } fs::set_permissions(&private_key, fs::Permissions::from_mode(0o600))?; fs::set_permissions(&public_key, fs::Permissions::from_mode(0o644))?; Ok((private_key, public_key)) } pub(crate) fn load_or_create_instance_config( instance_dir: &Path, ) -> Result> { let config_path = instance_dir.join(INSTANCE_FILENAME); let mut config = if config_path.exists() { let raw = fs::read_to_string(&config_path)?; toml::from_str::(&raw)? } else { InstanceConfig { id: String::new(), ssh_user: default_ssh_user(), sudo_password: String::new(), last_active: None, vm_ipv4: None, } }; let mut changed = false; if config.ssh_user.trim().is_empty() { config.ssh_user = default_ssh_user(); changed = true; } if config.id.trim().is_empty() { config.id = Uuid::now_v7().to_string(); changed = true; } if config.sudo_password.trim().is_empty() { config.sudo_password = generate_password(); changed = true; } if !config_path.exists() || changed { write_instance_config(&config_path, &config)?; } Ok(config) } fn read_instance_config( instance_dir: &Path, ) -> Result, Box> { let config_path = instance_dir.join(INSTANCE_FILENAME); if !config_path.exists() { return Ok(None); } let raw = fs::read_to_string(&config_path)?; let config = toml::from_str::(&raw)?; Ok(Some(config)) } pub fn read_instance_vm_ip( instance_dir: &Path, ) -> Result, Box> { let config = read_instance_config(instance_dir)?; Ok(config.and_then(|cfg| cfg.vm_ipv4)) } pub fn read_instance_ssh_user( instance_dir: &Path, ) -> Result, Box> { let config = read_instance_config(instance_dir)?; Ok(config .map(|cfg| cfg.ssh_user) .filter(|user| !user.trim().is_empty())) } pub fn touch_last_active(instance_dir: &Path) -> Result<(), Box> { let mut config = load_or_create_instance_config(instance_dir)?; let now = OffsetDateTime::now_utc().format(&Rfc3339)?; config.last_active = Some(now); write_instance_config(&instance_dir.join(INSTANCE_FILENAME), &config)?; Ok(()) } pub(crate) fn write_instance_config( path: &Path, config: &InstanceConfig, ) -> Result<(), Box> { let data = toml::to_string_pretty(config)?; fs::write(path, data)?; fs::set_permissions(path, fs::Permissions::from_mode(0o600))?; Ok(()) } fn generate_password() -> String { Uuid::now_v7().simple().to_string() } pub(crate) fn extract_ipv4(line: &str) -> Option { let mut current = String::new(); let mut best: Option = None; for ch in line.chars().chain(std::iter::once(' ')) { if ch.is_ascii_digit() || ch == '.' { current.push(ch); } else if !current.is_empty() { if is_ipv4_candidate(¤t) { best = Some(current.clone()); break; } current.clear(); } } best } fn wait_for_vm_ipv4( instance_dir: &Path, timeout: Duration, ) -> Result<(), Box> { let start = Instant::now(); let mut next_log_at = start + Duration::from_secs(10); let mut next_status_check = start; tracing::info!("waiting for vm ipv4"); let status_path = instance_dir.join(STATUS_FILE_NAME); let mut last_status: Option = None; let mut status_missing = true; let mut once_hint = false; loop { let config = load_or_create_instance_config(instance_dir)?; if config.vm_ipv4.is_some() { let _ = fs::remove_file(&status_path); return Ok(()); } if start.elapsed() > timeout { let _ = fs::remove_file(&status_path); return Err("Timed out waiting for VM IPv4".into()); } let now = Instant::now(); if now >= next_status_check { match fs::read_to_string(&status_path) { Ok(status) => { status_missing = false; let status = status.trim().to_string(); if !status.is_empty() && last_status.as_deref() != Some(status.as_str()) { tracing::info!("[background]: {}", status); last_status = Some(status); next_log_at = now + Duration::from_secs(20); } } Err(_) => { status_missing = true; } } next_status_check = now + Duration::from_millis(500); } if now >= next_log_at { let waited = start.elapsed(); if waited.as_secs() > 15 && !once_hint { tracing::info!( "if vibebox is just initialized in this directory, it might take up to 1 minute depending on your machine, and then you can enjoy secure & speed vibecoding! go pack!" ); once_hint = true; } if status_missing { tracing::info!("still waiting for vm ipv4, {}s elapsed", waited.as_secs(),); } next_log_at += Duration::from_secs(20); } thread::sleep(Duration::from_millis(200)); } } fn run_ssh_session( ssh_key: PathBuf, ssh_user: String, ip: String, ) -> Result<(), Box> { let mut attempts = 0usize; loop { attempts += 1; if !ssh_port_open(&ip) { tracing::debug!(attempts, "ssh port doesn't open yet"); tracing::info!( attempts, ip = %ip, "waiting for ssh port ({}/{})", attempts, SSH_CONNECT_RETRIES ); if attempts >= SSH_CONNECT_RETRIES { return Err( format!("ssh port not ready after {SSH_CONNECT_RETRIES} attempts").into(), ); } thread::sleep(Duration::from_millis(SSH_CONNECT_DELAY_MS)); continue; } tracing::info!( attempts, user = %ssh_user, ip = %ip, "starting ssh ({}/{})", attempts, SSH_CONNECT_RETRIES ); let status = Command::new("ssh") .args([ "-i", ssh_key.to_str().unwrap_or(".vibebox/ssh_key"), "-o", "IdentitiesOnly=yes", "-o", "StrictHostKeyChecking=no", "-o", "UserKnownHostsFile=/dev/null", "-o", "GlobalKnownHostsFile=/dev/null", "-o", "PasswordAuthentication=no", "-o", "BatchMode=yes", "-o", "LogLevel=ERROR", "-o", "ConnectTimeout=5", ]) .env_remove("LC_CTYPE") .env_remove("LC_ALL") .env_remove("LANG") .arg(format!("{ssh_user}@{ip}")) .stdin(Stdio::inherit()) .stdout(Stdio::inherit()) .stderr(Stdio::inherit()) .status(); match status { Ok(status) if status.success() => { tracing::info!(status = %status, "ssh exited"); break; } Ok(status) if status.code() == Some(255) => { tracing::warn!(status = %status, "ssh connection failed"); if attempts >= SSH_CONNECT_RETRIES { return Err(format!("ssh failed after {SSH_CONNECT_RETRIES} attempts").into()); } thread::sleep(Duration::from_millis(500)); } Ok(status) => { tracing::info!(status = %status, "ssh exited"); break; } Err(err) => { tracing::error!(error = %err, "failed to start ssh"); return Err(format!("failed to start ssh: {err}").into()); } } } Ok(()) } fn is_ipv4_candidate(candidate: &str) -> bool { let parts: Vec<&str> = candidate.split('.').collect(); if parts.len() != 4 { return false; } for part in parts { if part.is_empty() || part.len() > 3 { return false; } if part.parse::().is_err() { return false; } } true } fn ssh_port_open(ip: &str) -> bool { let addr: SocketAddr = match format!("{ip}:22").parse() { Ok(addr) => addr, Err(_) => return false, }; TcpStream::connect_timeout(&addr, std::time::Duration::from_millis(500)).is_ok() } pub(crate) fn build_ssh_login_actions( config: &Arc>, project_name: &str, project_guest_dir: &str, guest_dir: &str, key_name: &str, home_links_script: &str, ) -> Vec { let config_guard = config.lock().expect("config mutex poisoned"); let ssh_user = config_guard.ssh_user.clone(); let sudo_password = config_guard.sudo_password.clone(); drop(config_guard); let key_path = format!("{guest_dir}/{key_name}.pub"); let setup_script = SSH_SETUP_SCRIPT .replace("__SSH_USER__", &ssh_user) .replace("__SUDO_PASSWORD__", &sudo_password) .replace("__PROJECT_NAME__", project_name) .replace("__PROJECT_GUEST_DIR__", project_guest_dir) .replace("__KEY_PATH__", &key_path) .replace("__VIBEBOX_SHELL_SCRIPT__", &commands::render_shell_script()) .replace("__VIBEBOX_HOME_LINKS__", home_links_script); let setup = vm::script_command_from_content("ssh_setup", &setup_script) .expect("ssh setup script contained invalid marker"); vec![LoginAction::Send(setup)] }