@@ -0,0 +1,521 @@
1+use std::{
2+ env,
3+ fs,
4+ io::{self, Write},
5+ net::{SocketAddr, TcpStream},
6+ os::unix::{
7+ fs::PermissionsExt,
8+ io::OwnedFd,
9+ },
10+ path::{Path, PathBuf},
11+ process::{Command, Stdio},
12+ sync::{
13+ Arc, Mutex,
14+ atomic::{AtomicBool, Ordering},
15+ },
16+ thread,
17+};
18+
19+use serde::{Deserialize, Serialize};
20+use uuid::Uuid;
21+use std::sync::mpsc::Sender;
22+
23+use crate::{
24+ session_manager::INSTANCE_DIR_NAME,
25+ tui::{self, AppState},
26+ vm::{self, DirectoryShare, LoginAction, VmInput},
27+};
28+
29+const INSTANCE_TOML: &str = "instance.toml";
30+const SSH_KEY_NAME: &str = "ssh_key";
31+const SERIAL_LOG_NAME: &str = "serial.log";
32+const SSH_GUEST_DIR: &str = "/root/.vibebox";
33+const DEFAULT_SSH_USER: &str = "vibebox";
34+const SSH_CONNECT_RETRIES: usize = 20;
35+const SSH_CONNECT_DELAY_MS: u64 = 500;
36+
37+#[derive(Debug, Clone, Serialize, Deserialize)]
38+struct InstanceConfig {
39+ #[serde(default = "default_ssh_user")]
40+ ssh_user: String,
41+ #[serde(default)]
42+ sudo_password: String,
43+ #[serde(default)]
44+ vm_ipv4: Option<String>,
45+}
46+
47+fn default_ssh_user() -> String {
48+ DEFAULT_SSH_USER.to_string()
49+}
50+
51+pub fn run_with_ssh(
52+ args: vm::CliArgs,
53+ app: Arc<Mutex<AppState>>,
54+) -> Result<(), Box<dyn std::error::Error>> {
55+ let project_root = env::current_dir()?;
56+ let project_name = project_root
57+ .file_name()
58+ .ok_or("Project directory has no name")?
59+ .to_string_lossy()
60+ .into_owned();
61+ let instance_dir = ensure_instance_dir(&project_root)?;
62+ let (ssh_key, _ssh_pub) = ensure_ssh_keypair(&instance_dir)?;
63+
64+ let config = load_or_create_instance_config(&instance_dir)?;
65+ let config = Arc::new(Mutex::new(config));
66+
67+ let extra_shares = vec![DirectoryShare::new(
68+ instance_dir.clone(),
69+ SSH_GUEST_DIR.into(),
70+ true,
71+ )?];
72+
73+ let extra_login_actions = build_ssh_login_actions(
74+ &config,
75+ &project_name,
76+ SSH_GUEST_DIR,
77+ SSH_KEY_NAME,
78+ );
79+
80+ vm::run_with_args_and_extras(
81+ args,
82+ |output_monitor, vm_output_fd, vm_input_fd| {
83+ spawn_ssh_io(
84+ app.clone(),
85+ config.clone(),
86+ instance_dir.clone(),
87+ ssh_key.clone(),
88+ output_monitor,
89+ vm_output_fd,
90+ vm_input_fd,
91+ )
92+ },
93+ extra_login_actions,
94+ extra_shares,
95+ )
96+}
97+
98+fn ensure_instance_dir(project_root: &Path) -> Result<PathBuf, io::Error> {
99+ let instance_dir = project_root.join(INSTANCE_DIR_NAME);
100+ fs::create_dir_all(&instance_dir)?;
101+ Ok(instance_dir)
102+}
103+
104+fn ensure_ssh_keypair(instance_dir: &Path) -> Result<(PathBuf, PathBuf), Box<dyn std::error::Error>> {
105+ let private_key = instance_dir.join(SSH_KEY_NAME);
106+ let public_key = instance_dir.join(format!("{SSH_KEY_NAME}.pub"));
107+
108+ if private_key.exists() && public_key.exists() {
109+ return Ok((private_key, public_key));
110+ }
111+
112+ if private_key.exists() {
113+ let _ = fs::remove_file(&private_key);
114+ }
115+ if public_key.exists() {
116+ let _ = fs::remove_file(&public_key);
117+ }
118+
119+ let status = Command::new("ssh-keygen")
120+ .args([
121+ "-t",
122+ "ed25519",
123+ "-N",
124+ "",
125+ "-f",
126+ private_key
127+ .to_str()
128+ .ok_or("ssh key path not utf-8")?,
129+ "-C",
130+ "vibebox",
131+ ])
132+ .status()?;
133+
134+ if !status.success() {
135+ return Err("ssh-keygen failed".into());
136+ }
137+
138+ fs::set_permissions(&private_key, fs::Permissions::from_mode(0o600))?;
139+ fs::set_permissions(&public_key, fs::Permissions::from_mode(0o644))?;
140+
141+ Ok((private_key, public_key))
142+}
143+
144+fn load_or_create_instance_config(
145+ instance_dir: &Path,
146+) -> Result<InstanceConfig, Box<dyn std::error::Error>> {
147+ let config_path = instance_dir.join(INSTANCE_TOML);
148+ let mut config = if config_path.exists() {
149+ let raw = fs::read_to_string(&config_path)?;
150+ toml::from_str::<InstanceConfig>(&raw)?
151+ } else {
152+ InstanceConfig {
153+ ssh_user: default_ssh_user(),
154+ sudo_password: String::new(),
155+ vm_ipv4: None,
156+ }
157+ };
158+
159+ let mut changed = false;
160+ if config.ssh_user.trim().is_empty() {
161+ config.ssh_user = default_ssh_user();
162+ changed = true;
163+ }
164+
165+ if config.sudo_password.trim().is_empty() {
166+ config.sudo_password = generate_password();
167+ changed = true;
168+ }
169+
170+ if !config_path.exists() || changed {
171+ write_instance_config(&config_path, &config)?;
172+ }
173+
174+ Ok(config)
175+}
176+
177+fn write_instance_config(
178+ path: &Path,
179+ config: &InstanceConfig,
180+) -> Result<(), Box<dyn std::error::Error>> {
181+ let data = toml::to_string_pretty(config)?;
182+ fs::write(path, data)?;
183+ fs::set_permissions(path, fs::Permissions::from_mode(0o600))?;
184+ Ok(())
185+}
186+
187+fn generate_password() -> String {
188+ Uuid::now_v7().simple().to_string()
189+}
190+
191+fn extract_ipv4(line: &str) -> Option<String> {
192+ let mut current = String::new();
193+ let mut best: Option<String> = None;
194+
195+ for ch in line.chars().chain(std::iter::once(' ')) {
196+ if ch.is_ascii_digit() || ch == '.' {
197+ current.push(ch);
198+ } else if !current.is_empty() {
199+ if is_ipv4_candidate(¤t) {
200+ best = Some(current.clone());
201+ break;
202+ }
203+ current.clear();
204+ }
205+ }
206+
207+ best
208+}
209+
210+fn is_ipv4_candidate(candidate: &str) -> bool {
211+ let parts: Vec<&str> = candidate.split('.').collect();
212+ if parts.len() != 4 {
213+ return false;
214+ }
215+ for part in parts {
216+ if part.is_empty() || part.len() > 3 {
217+ return false;
218+ }
219+ if part.parse::<u8>().is_err() {
220+ return false;
221+ }
222+ }
223+ true
224+}
225+
226+fn ssh_port_open(ip: &str) -> bool {
227+ let addr: SocketAddr = match format!("{ip}:22").parse() {
228+ Ok(addr) => addr,
229+ Err(_) => return false,
230+ };
231+ TcpStream::connect_timeout(&addr, std::time::Duration::from_millis(500)).is_ok()
232+}
233+
234+fn build_ssh_login_actions(
235+ config: &Arc<Mutex<InstanceConfig>>,
236+ project_name: &str,
237+ guest_dir: &str,
238+ key_name: &str,
239+) -> Vec<LoginAction> {
240+ let config_guard = config.lock().expect("config mutex poisoned");
241+ let ssh_user = config_guard.ssh_user.clone();
242+ let sudo_password = config_guard.sudo_password.clone();
243+ drop(config_guard);
244+
245+ let key_path = format!("{guest_dir}/{key_name}.pub");
246+
247+ let mask_cache = format!(
248+ "if [ -d /root/{project_name}/.vibebox ]; then mount -t tmpfs tmpfs /root/{project_name}/.vibebox; fi"
249+ );
250+
251+ let setup_lines = [
252+ "if ! command -v sshd >/dev/null 2>&1; then apt-get update && apt-get install -y openssh-server sudo; fi",
253+ "systemctl enable ssh >/dev/null 2>&1 || true",
254+ &format!("id -u {ssh_user} >/dev/null 2>&1 || useradd -m -s /bin/bash {ssh_user}"),
255+ &format!("echo \"{ssh_user}:{sudo_password}\" | chpasswd"),
256+ &format!("usermod -aG sudo {ssh_user}"),
257+ &format!("install -d -m 700 /home/{ssh_user}/.ssh"),
258+ &format!("install -m 600 {key_path} /home/{ssh_user}/.ssh/authorized_keys"),
259+ &format!("chown -R {ssh_user}:{ssh_user} /home/{ssh_user}/.ssh"),
260+ "mkdir -p /etc/ssh/sshd_config.d",
261+ "cat >/etc/ssh/sshd_config.d/vibebox.conf <<'VIBEBOX_SSHD'",
262+ "PasswordAuthentication no",
263+ "KbdInteractiveAuthentication no",
264+ "ChallengeResponseAuthentication no",
265+ "PubkeyAuthentication yes",
266+ "PermitRootLogin no",
267+ &format!("AllowUsers {ssh_user}"),
268+ "VIBEBOX_SSHD",
269+ "systemctl restart ssh",
270+ "echo VIBEBOX_SSH_READY",
271+ ];
272+ let setup = setup_lines.join("\n");
273+
274+ let ip_probe = "while true; do ip=$(ip -4 -o addr show scope global | awk '{print $4}' | cut -d/ -f1 | head -n 1); if [ -n \"$ip\" ]; then echo VIBEBOX_IPV4=$ip; break; fi; sleep 1; done";
275+
276+ vec![
277+ LoginAction::Send(mask_cache),
278+ LoginAction::Send(setup),
279+ LoginAction::Send(ip_probe.to_string()),
280+ ]
281+}
282+
283+fn spawn_ssh_io(
284+ app: Arc<Mutex<AppState>>,
285+ config: Arc<Mutex<InstanceConfig>>,
286+ instance_dir: PathBuf,
287+ ssh_key: PathBuf,
288+ output_monitor: Arc<vm::OutputMonitor>,
289+ vm_output_fd: OwnedFd,
290+ vm_input_fd: OwnedFd,
291+) -> vm::IoContext {
292+ let io_control = vm::IoControl::new();
293+
294+ let log_path = instance_dir.join(SERIAL_LOG_NAME);
295+ let log_file = fs::OpenOptions::new()
296+ .create(true)
297+ .append(true)
298+ .open(&log_path)
299+ .ok()
300+ .map(|file| Arc::new(Mutex::new(file)));
301+
302+ let ssh_connected = Arc::new(AtomicBool::new(false));
303+ let ssh_started = Arc::new(AtomicBool::new(false));
304+ let ssh_ready = Arc::new(AtomicBool::new(false));
305+ let input_tx_holder: Arc<Mutex<Option<Sender<VmInput>>>> =
306+ Arc::new(Mutex::new(None));
307+
308+ let instance_path = instance_dir.join(INSTANCE_TOML);
309+ let config_for_output = config.clone();
310+ let log_for_output = log_file.clone();
311+ let ssh_connected_for_output = ssh_connected.clone();
312+ let ssh_ready_for_output = ssh_ready.clone();
313+
314+ let mut line_buf = String::new();
315+
316+ let on_line = {
317+ let app = app.clone();
318+ move |line: &str| {
319+ if line == ":help" {
320+ if let Ok(mut locked) = app.lock() {
321+ let _ = tui::render_commands_component(&mut locked);
322+ }
323+ return true;
324+ }
325+ false
326+ }
327+ };
328+
329+ let on_output = move |bytes: &[u8]| {
330+ if ssh_connected_for_output.load(Ordering::SeqCst) {
331+ if let Some(log) = &log_for_output {
332+ if let Ok(mut file) = log.lock() {
333+ let _ = file.write_all(bytes);
334+ }
335+ }
336+ }
337+
338+ let text = String::from_utf8_lossy(bytes);
339+ line_buf.push_str(&text);
340+
341+ while let Some(pos) = line_buf.find('\n') {
342+ let mut line = line_buf[..pos].to_string();
343+ if line.ends_with('\r') {
344+ line.pop();
345+ }
346+ line_buf.drain(..=pos);
347+
348+ let cleaned = line.trim_start_matches(|c: char| c == '\r' || c == ' ');
349+
350+ if let Some(pos) = cleaned.find("VIBEBOX_IPV4=") {
351+ let ip_raw = &cleaned[(pos + "VIBEBOX_IPV4=".len())..];
352+ let ip = extract_ipv4(ip_raw).unwrap_or_default();
353+ if !ip.is_empty() {
354+ if let Ok(mut cfg) = config_for_output.lock() {
355+ if cfg.vm_ipv4.as_deref() != Some(ip.as_str()) {
356+ cfg.vm_ipv4 = Some(ip.clone());
357+ let _ = write_instance_config(&instance_path, &cfg);
358+ eprintln!("[vibebox] detected vm IPv4: {ip}");
359+ }
360+ }
361+ }
362+ }
363+
364+ if cleaned.contains("VIBEBOX_SSH_READY") {
365+ ssh_ready_for_output.store(true, Ordering::SeqCst);
366+ eprintln!("[vibebox] sshd ready");
367+ }
368+ }
369+ };
370+
371+ let io_ctx = vm::spawn_vm_io_with_hooks(
372+ output_monitor,
373+ vm_output_fd,
374+ vm_input_fd,
375+ io_control.clone(),
376+ on_line,
377+ on_output,
378+ );
379+
380+ *input_tx_holder.lock().unwrap() = Some(io_ctx.input_tx.clone());
381+
382+ let ssh_ready_for_thread = ssh_ready.clone();
383+ let ssh_started_for_thread = ssh_started.clone();
384+ let ssh_connected_for_thread = ssh_connected.clone();
385+ let io_control_for_thread = io_control.clone();
386+ let ssh_key_for_thread = ssh_key.clone();
387+ let config_for_thread = config.clone();
388+ let input_tx_holder_for_thread = input_tx_holder.clone();
389+
390+ thread::spawn(move || {
391+ loop {
392+ if ssh_started_for_thread.load(Ordering::SeqCst) {
393+ break;
394+ }
395+ if !ssh_ready_for_thread.load(Ordering::SeqCst) {
396+ thread::sleep(std::time::Duration::from_millis(200));
397+ continue;
398+ }
399+
400+ let ip = config_for_thread
401+ .lock()
402+ .ok()
403+ .and_then(|cfg| cfg.vm_ipv4.clone());
404+ let ssh_user = config_for_thread
405+ .lock()
406+ .map(|cfg| cfg.ssh_user.clone())
407+ .unwrap_or_else(|_| DEFAULT_SSH_USER.to_string());
408+
409+ if let Some(ip) = ip {
410+ if ssh_started_for_thread
411+ .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
412+ .is_err()
413+ {
414+ break;
415+ }
416+
417+ eprintln!("[vibebox] starting ssh to {ssh_user}@{ip}");
418+ io_control_for_thread.request_terminal_restore();
419+ io_control_for_thread.set_forward_output(false);
420+ io_control_for_thread.set_forward_input(false);
421+ ssh_connected_for_thread.store(true, Ordering::SeqCst);
422+
423+ let mut attempts = 0usize;
424+ loop {
425+ attempts += 1;
426+ if !ssh_port_open(&ip) {
427+ eprintln!(
428+ "[vibebox] waiting for ssh port on {ip} (attempt {attempts}/{SSH_CONNECT_RETRIES})"
429+ );
430+ if attempts >= SSH_CONNECT_RETRIES {
431+ ssh_connected_for_thread.store(false, Ordering::SeqCst);
432+ eprintln!(
433+ "[vibebox] ssh port not ready after {SSH_CONNECT_RETRIES} attempts"
434+ );
435+ break;
436+ }
437+ thread::sleep(std::time::Duration::from_millis(
438+ SSH_CONNECT_DELAY_MS,
439+ ));
440+ continue;
441+ }
442+
443+ eprintln!(
444+ "[vibebox] starting ssh to {ssh_user}@{ip} (attempt {attempts}/{SSH_CONNECT_RETRIES})"
445+ );
446+ let status = Command::new("ssh")
447+ .args([
448+ "-i",
449+ ssh_key_for_thread
450+ .to_str()
451+ .unwrap_or(".vibebox/ssh_key"),
452+ "-o",
453+ "IdentitiesOnly=yes",
454+ "-o",
455+ "StrictHostKeyChecking=no",
456+ "-o",
457+ "UserKnownHostsFile=/dev/null",
458+ "-o",
459+ "GlobalKnownHostsFile=/dev/null",
460+ "-o",
461+ "PasswordAuthentication=no",
462+ "-o",
463+ "BatchMode=yes",
464+ "-o",
465+ "LogLevel=ERROR",
466+ "-o",
467+ "ConnectTimeout=5",
468+ ])
469+ .arg(format!("{ssh_user}@{ip}"))
470+ .stdin(Stdio::inherit())
471+ .stdout(Stdio::inherit())
472+ .stderr(Stdio::inherit())
473+ .status();
474+
475+ match status {
476+ Ok(status) if status.success() => {
477+ ssh_connected_for_thread.store(false, Ordering::SeqCst);
478+ eprintln!("[vibebox] ssh exited with status: {status}");
479+ if let Some(tx) = input_tx_holder_for_thread.lock().unwrap().clone() {
480+ let _ =
481+ tx.send(VmInput::Bytes(b"systemctl poweroff\n".to_vec()));
482+ }
483+ break;
484+ }
485+ Ok(status) if status.code() == Some(255) => {
486+ eprintln!("[vibebox] ssh connection failed: {status}");
487+ if attempts >= SSH_CONNECT_RETRIES {
488+ ssh_connected_for_thread.store(false, Ordering::SeqCst);
489+ eprintln!(
490+ "[vibebox] ssh failed after {SSH_CONNECT_RETRIES} attempts"
491+ );
492+ break;
493+ }
494+ thread::sleep(std::time::Duration::from_millis(500));
495+ continue;
496+ }
497+ Ok(status) => {
498+ ssh_connected_for_thread.store(false, Ordering::SeqCst);
499+ eprintln!("[vibebox] ssh exited with status: {status}");
500+ if let Some(tx) = input_tx_holder_for_thread.lock().unwrap().clone() {
501+ let _ =
502+ tx.send(VmInput::Bytes(b"systemctl poweroff\n".to_vec()));
503+ }
504+ break;
505+ }
506+ Err(err) => {
507+ ssh_connected_for_thread.store(false, Ordering::SeqCst);
508+ eprintln!("[vibebox] failed to start ssh: {err}");
509+ break;
510+ }
511+ }
512+ }
513+ break;
514+ }
515+
516+ thread::sleep(std::time::Duration::from_millis(200));
517+ }
518+ });
519+
520+ io_ctx
521+}