feat: now added ssh

src/bin/vibebox-cli.rs

@@ -7,7 +7,7 @@ use std::{
 use color_eyre::Result;
 
 use vibebox::tui::{AppState, VmInfo};
-use vibebox::{tui, vm};
+use vibebox::{instance, tui, vm};
 
 fn main() -> Result<()> {
     color_eyre::install()?;
@@ -40,10 +40,8 @@ fn main() -> Result<()> {
         stdout.flush()?;
     }
 
-    vm::run_with_args(args, |output_monitor, vm_output_fd, vm_input_fd| {
-        tui::passthrough_vm_io(app.clone(), output_monitor, vm_output_fd, vm_input_fd)
-    })
-    .map_err(|err| color_eyre::eyre::eyre!(err.to_string()))?;
+    instance::run_with_ssh(args, app.clone())
+        .map_err(|err| color_eyre::eyre::eyre!(err.to_string()))?;
 
     Ok(())
 }

src/instance.rs

@@ -0,0 +1,521 @@
+use std::{
+    env,
+    fs,
+    io::{self, Write},
+    net::{SocketAddr, TcpStream},
+    os::unix::{
+        fs::PermissionsExt,
+        io::OwnedFd,
+    },
+    path::{Path, PathBuf},
+    process::{Command, Stdio},
+    sync::{
+        Arc, Mutex,
+        atomic::{AtomicBool, Ordering},
+    },
+    thread,
+};
+
+use serde::{Deserialize, Serialize};
+use uuid::Uuid;
+use std::sync::mpsc::Sender;
+
+use crate::{
+    session_manager::INSTANCE_DIR_NAME,
+    tui::{self, AppState},
+    vm::{self, DirectoryShare, LoginAction, VmInput},
+};
+
+const INSTANCE_TOML: &str = "instance.toml";
+const SSH_KEY_NAME: &str = "ssh_key";
+const SERIAL_LOG_NAME: &str = "serial.log";
+const SSH_GUEST_DIR: &str = "/root/.vibebox";
+const DEFAULT_SSH_USER: &str = "vibebox";
+const SSH_CONNECT_RETRIES: usize = 20;
+const SSH_CONNECT_DELAY_MS: u64 = 500;
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+struct InstanceConfig {
+    #[serde(default = "default_ssh_user")]
+    ssh_user: String,
+    #[serde(default)]
+    sudo_password: String,
+    #[serde(default)]
+    vm_ipv4: Option<String>,
+}
+
+fn default_ssh_user() -> String {
+    DEFAULT_SSH_USER.to_string()
+}
+
+pub fn run_with_ssh(
+    args: vm::CliArgs,
+    app: Arc<Mutex<AppState>>,
+) -> Result<(), Box<dyn std::error::Error>> {
+    let project_root = env::current_dir()?;
+    let project_name = project_root
+        .file_name()
+        .ok_or("Project directory has no name")?
+        .to_string_lossy()
+        .into_owned();
+    let instance_dir = ensure_instance_dir(&project_root)?;
+    let (ssh_key, _ssh_pub) = ensure_ssh_keypair(&instance_dir)?;
+
+    let config = load_or_create_instance_config(&instance_dir)?;
+    let config = Arc::new(Mutex::new(config));
+
+    let extra_shares = vec![DirectoryShare::new(
+        instance_dir.clone(),
+        SSH_GUEST_DIR.into(),
+        true,
+    )?];
+
+    let extra_login_actions = build_ssh_login_actions(
+        &config,
+        &project_name,
+        SSH_GUEST_DIR,
+        SSH_KEY_NAME,
+    );
+
+    vm::run_with_args_and_extras(
+        args,
+        |output_monitor, vm_output_fd, vm_input_fd| {
+            spawn_ssh_io(
+                app.clone(),
+                config.clone(),
+                instance_dir.clone(),
+                ssh_key.clone(),
+                output_monitor,
+                vm_output_fd,
+                vm_input_fd,
+            )
+        },
+        extra_login_actions,
+        extra_shares,
+    )
+}
+
+fn ensure_instance_dir(project_root: &Path) -> Result<PathBuf, io::Error> {
+    let instance_dir = project_root.join(INSTANCE_DIR_NAME);
+    fs::create_dir_all(&instance_dir)?;
+    Ok(instance_dir)
+}
+
+fn ensure_ssh_keypair(instance_dir: &Path) -> Result<(PathBuf, PathBuf), Box<dyn std::error::Error>> {
+    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",
+        ])
+        .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))
+}
+
+fn load_or_create_instance_config(
+    instance_dir: &Path,
+) -> Result<InstanceConfig, Box<dyn std::error::Error>> {
+    let config_path = instance_dir.join(INSTANCE_TOML);
+    let mut config = if config_path.exists() {
+        let raw = fs::read_to_string(&config_path)?;
+        toml::from_str::<InstanceConfig>(&raw)?
+    } else {
+        InstanceConfig {
+            ssh_user: default_ssh_user(),
+            sudo_password: String::new(),
+            vm_ipv4: None,
+        }
+    };
+
+    let mut changed = false;
+    if config.ssh_user.trim().is_empty() {
+        config.ssh_user = default_ssh_user();
+        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 write_instance_config(
+    path: &Path,
+    config: &InstanceConfig,
+) -> Result<(), Box<dyn std::error::Error>> {
+    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()
+}
+
+fn extract_ipv4(line: &str) -> Option<String> {
+    let mut current = String::new();
+    let mut best: Option<String> = 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(&current) {
+                best = Some(current.clone());
+                break;
+            }
+            current.clear();
+        }
+    }
+
+    best
+}
+
+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::<u8>().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()
+}
+
+fn build_ssh_login_actions(
+    config: &Arc<Mutex<InstanceConfig>>,
+    project_name: &str,
+    guest_dir: &str,
+    key_name: &str,
+) -> Vec<LoginAction> {
+    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 mask_cache = format!(
+        "if [ -d /root/{project_name}/.vibebox ]; then mount -t tmpfs tmpfs /root/{project_name}/.vibebox; fi"
+    );
+
+    let setup_lines = [
+        "if ! command -v sshd >/dev/null 2>&1; then apt-get update && apt-get install -y openssh-server sudo; fi",
+        "systemctl enable ssh >/dev/null 2>&1 || true",
+        &format!("id -u {ssh_user} >/dev/null 2>&1 || useradd -m -s /bin/bash {ssh_user}"),
+        &format!("echo \"{ssh_user}:{sudo_password}\" | chpasswd"),
+        &format!("usermod -aG sudo {ssh_user}"),
+        &format!("install -d -m 700 /home/{ssh_user}/.ssh"),
+        &format!("install -m 600 {key_path} /home/{ssh_user}/.ssh/authorized_keys"),
+        &format!("chown -R {ssh_user}:{ssh_user} /home/{ssh_user}/.ssh"),
+        "mkdir -p /etc/ssh/sshd_config.d",
+        "cat >/etc/ssh/sshd_config.d/vibebox.conf <<'VIBEBOX_SSHD'",
+        "PasswordAuthentication no",
+        "KbdInteractiveAuthentication no",
+        "ChallengeResponseAuthentication no",
+        "PubkeyAuthentication yes",
+        "PermitRootLogin no",
+        &format!("AllowUsers {ssh_user}"),
+        "VIBEBOX_SSHD",
+        "systemctl restart ssh",
+        "echo VIBEBOX_SSH_READY",
+    ];
+    let setup = setup_lines.join("\n");
+
+    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";
+
+    vec![
+        LoginAction::Send(mask_cache),
+        LoginAction::Send(setup),
+        LoginAction::Send(ip_probe.to_string()),
+    ]
+}
+
+fn spawn_ssh_io(
+    app: Arc<Mutex<AppState>>,
+    config: Arc<Mutex<InstanceConfig>>,
+    instance_dir: PathBuf,
+    ssh_key: PathBuf,
+    output_monitor: Arc<vm::OutputMonitor>,
+    vm_output_fd: OwnedFd,
+    vm_input_fd: OwnedFd,
+) -> vm::IoContext {
+    let io_control = vm::IoControl::new();
+
+    let log_path = instance_dir.join(SERIAL_LOG_NAME);
+    let log_file = fs::OpenOptions::new()
+        .create(true)
+        .append(true)
+        .open(&log_path)
+        .ok()
+        .map(|file| Arc::new(Mutex::new(file)));
+
+    let ssh_connected = Arc::new(AtomicBool::new(false));
+    let ssh_started = Arc::new(AtomicBool::new(false));
+    let ssh_ready = Arc::new(AtomicBool::new(false));
+    let input_tx_holder: Arc<Mutex<Option<Sender<VmInput>>>> =
+        Arc::new(Mutex::new(None));
+
+    let instance_path = instance_dir.join(INSTANCE_TOML);
+    let config_for_output = config.clone();
+    let log_for_output = log_file.clone();
+    let ssh_connected_for_output = ssh_connected.clone();
+    let ssh_ready_for_output = ssh_ready.clone();
+
+    let mut line_buf = String::new();
+
+    let on_line = {
+        let app = app.clone();
+        move |line: &str| {
+            if line == ":help" {
+                if let Ok(mut locked) = app.lock() {
+                    let _ = tui::render_commands_component(&mut locked);
+                }
+                return true;
+            }
+            false
+        }
+    };
+
+    let on_output = move |bytes: &[u8]| {
+        if ssh_connected_for_output.load(Ordering::SeqCst) {
+            if let Some(log) = &log_for_output {
+                if let Ok(mut file) = log.lock() {
+                    let _ = file.write_all(bytes);
+                }
+            }
+        }
+
+        let text = String::from_utf8_lossy(bytes);
+        line_buf.push_str(&text);
+
+        while let Some(pos) = line_buf.find('\n') {
+            let mut line = line_buf[..pos].to_string();
+            if line.ends_with('\r') {
+                line.pop();
+            }
+            line_buf.drain(..=pos);
+
+            let cleaned = line.trim_start_matches(|c: char| c == '\r' || c == ' ');
+
+            if let Some(pos) = cleaned.find("VIBEBOX_IPV4=") {
+                let ip_raw = &cleaned[(pos + "VIBEBOX_IPV4=".len())..];
+                let ip = extract_ipv4(ip_raw).unwrap_or_default();
+                if !ip.is_empty() {
+                    if let Ok(mut cfg) = config_for_output.lock() {
+                        if cfg.vm_ipv4.as_deref() != Some(ip.as_str()) {
+                            cfg.vm_ipv4 = Some(ip.clone());
+                            let _ = write_instance_config(&instance_path, &cfg);
+                            eprintln!("[vibebox] detected vm IPv4: {ip}");
+                        }
+                    }
+                }
+            }
+
+            if cleaned.contains("VIBEBOX_SSH_READY") {
+                ssh_ready_for_output.store(true, Ordering::SeqCst);
+                eprintln!("[vibebox] sshd ready");
+            }
+        }
+    };
+
+    let io_ctx = vm::spawn_vm_io_with_hooks(
+        output_monitor,
+        vm_output_fd,
+        vm_input_fd,
+        io_control.clone(),
+        on_line,
+        on_output,
+    );
+
+    *input_tx_holder.lock().unwrap() = Some(io_ctx.input_tx.clone());
+
+    let ssh_ready_for_thread = ssh_ready.clone();
+    let ssh_started_for_thread = ssh_started.clone();
+    let ssh_connected_for_thread = ssh_connected.clone();
+    let io_control_for_thread = io_control.clone();
+    let ssh_key_for_thread = ssh_key.clone();
+    let config_for_thread = config.clone();
+    let input_tx_holder_for_thread = input_tx_holder.clone();
+
+    thread::spawn(move || {
+        loop {
+            if ssh_started_for_thread.load(Ordering::SeqCst) {
+                break;
+            }
+            if !ssh_ready_for_thread.load(Ordering::SeqCst) {
+                thread::sleep(std::time::Duration::from_millis(200));
+                continue;
+            }
+
+            let ip = config_for_thread
+                .lock()
+                .ok()
+                .and_then(|cfg| cfg.vm_ipv4.clone());
+            let ssh_user = config_for_thread
+                .lock()
+                .map(|cfg| cfg.ssh_user.clone())
+                .unwrap_or_else(|_| DEFAULT_SSH_USER.to_string());
+
+            if let Some(ip) = ip {
+                if ssh_started_for_thread
+                    .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
+                    .is_err()
+                {
+                    break;
+                }
+
+                eprintln!("[vibebox] starting ssh to {ssh_user}@{ip}");
+                io_control_for_thread.request_terminal_restore();
+                io_control_for_thread.set_forward_output(false);
+                io_control_for_thread.set_forward_input(false);
+                ssh_connected_for_thread.store(true, Ordering::SeqCst);
+
+                let mut attempts = 0usize;
+                loop {
+                    attempts += 1;
+                    if !ssh_port_open(&ip) {
+                        eprintln!(
+                            "[vibebox] waiting for ssh port on {ip} (attempt {attempts}/{SSH_CONNECT_RETRIES})"
+                        );
+                        if attempts >= SSH_CONNECT_RETRIES {
+                            ssh_connected_for_thread.store(false, Ordering::SeqCst);
+                            eprintln!(
+                                "[vibebox] ssh port not ready after {SSH_CONNECT_RETRIES} attempts"
+                            );
+                            break;
+                        }
+                        thread::sleep(std::time::Duration::from_millis(
+                            SSH_CONNECT_DELAY_MS,
+                        ));
+                        continue;
+                    }
+
+                    eprintln!(
+                        "[vibebox] starting ssh to {ssh_user}@{ip} (attempt {attempts}/{SSH_CONNECT_RETRIES})"
+                    );
+                    let status = Command::new("ssh")
+                        .args([
+                            "-i",
+                            ssh_key_for_thread
+                                .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",
+                        ])
+                        .arg(format!("{ssh_user}@{ip}"))
+                        .stdin(Stdio::inherit())
+                        .stdout(Stdio::inherit())
+                        .stderr(Stdio::inherit())
+                        .status();
+
+                    match status {
+                        Ok(status) if status.success() => {
+                            ssh_connected_for_thread.store(false, Ordering::SeqCst);
+                            eprintln!("[vibebox] ssh exited with status: {status}");
+                            if let Some(tx) = input_tx_holder_for_thread.lock().unwrap().clone() {
+                                let _ =
+                                    tx.send(VmInput::Bytes(b"systemctl poweroff\n".to_vec()));
+                            }
+                            break;
+                        }
+                        Ok(status) if status.code() == Some(255) => {
+                            eprintln!("[vibebox] ssh connection failed: {status}");
+                            if attempts >= SSH_CONNECT_RETRIES {
+                                ssh_connected_for_thread.store(false, Ordering::SeqCst);
+                                eprintln!(
+                                    "[vibebox] ssh failed after {SSH_CONNECT_RETRIES} attempts"
+                                );
+                                break;
+                            }
+                            thread::sleep(std::time::Duration::from_millis(500));
+                            continue;
+                        }
+                        Ok(status) => {
+                            ssh_connected_for_thread.store(false, Ordering::SeqCst);
+                            eprintln!("[vibebox] ssh exited with status: {status}");
+                            if let Some(tx) = input_tx_holder_for_thread.lock().unwrap().clone() {
+                                let _ =
+                                    tx.send(VmInput::Bytes(b"systemctl poweroff\n".to_vec()));
+                            }
+                            break;
+                        }
+                        Err(err) => {
+                            ssh_connected_for_thread.store(false, Ordering::SeqCst);
+                            eprintln!("[vibebox] failed to start ssh: {err}");
+                            break;
+                        }
+                    }
+                }
+                break;
+            }
+
+            thread::sleep(std::time::Duration::from_millis(200));
+        }
+    });
+
+    io_ctx
+}

src/lib.rs

@@ -1,5 +1,6 @@
 pub mod session_manager;
 pub mod tui;
 pub mod vm;
+pub mod instance;
 
 pub use session_manager::{SessionError, SessionManager, SessionRecord};

src/provision.sh

@@ -13,11 +13,16 @@ apt-get install -y --no-install-recommends      \
         libssl-dev                              \
         curl                                    \
         git                                     \
-        ripgrep
+        ripgrep                                 \
+        openssh-server                          \
+        sudo
 
 # Set hostname to "vibe" so it's clear that you're inside the VM.
 hostnamectl set-hostname vibe
 
+# Enable SSH server so instances can use key-based auth.
+systemctl enable ssh
+
 # Set this env var so claude doesn't complain about running as root.'
 echo "export IS_SANDBOX=1" >> .bashrc
 

src/vm.rs

@@ -15,6 +15,7 @@ use std::{
     process::{Command, Stdio},
     sync::{
         Arc, Condvar, Mutex,
+        atomic::{AtomicBool, Ordering},
         mpsc::{self, Receiver, Sender},
     },
     thread,
@@ -43,7 +44,7 @@ const LOGIN_EXPECT_TIMEOUT: Duration = Duration::from_secs(120);
 const PROVISION_SCRIPT: &str = include_str!("provision.sh");
 
 #[derive(Clone)]
-enum LoginAction {
+pub(crate) enum LoginAction {
     Expect { text: String, timeout: Duration },
     Send(String),
     Script { path: PathBuf, index: usize },
@@ -51,14 +52,14 @@ enum LoginAction {
 use LoginAction::*;
 
 #[derive(Clone)]
-struct DirectoryShare {
+pub(crate) struct DirectoryShare {
     host: PathBuf,
     guest: PathBuf,
     read_only: bool,
 }
 
 impl DirectoryShare {
-    fn new(
+    pub(crate) fn new(
         host: PathBuf,
         mut guest: PathBuf,
         read_only: bool,
@@ -76,7 +77,7 @@ impl DirectoryShare {
         })
     }
 
-    fn from_mount_spec(spec: &str) -> Result<Self, Box<dyn std::error::Error>> {
+    pub(crate) fn from_mount_spec(spec: &str) -> Result<Self, Box<dyn std::error::Error>> {
         let parts: Vec<&str> = spec.split(':').collect();
         if parts.len() < 2 || parts.len() > 3 {
             return Err(format!("Invalid mount spec: {spec}").into());
@@ -115,23 +116,19 @@ impl DirectoryShare {
     }
 }
 
-pub fn run_cli() -> Result<(), Box<dyn std::error::Error>> {
-    let args = parse_cli()?;
-
-    if args.version() {
-        print_version();
-        return Ok(());
-    }
-
-    if args.help() {
-        print_help();
-        return Ok(());
-    }
-
-    run_with_args(args, spawn_vm_io)
+pub fn run_with_args<F>(args: CliArgs, io_handler: F) -> Result<(), Box<dyn std::error::Error>>
+where
+    F: FnOnce(Arc<OutputMonitor>, OwnedFd, OwnedFd) -> IoContext,
+{
+    run_with_args_and_extras(args, io_handler, Vec::new(), Vec::new())
 }
 
-pub fn run_with_args<F>(args: CliArgs, io_handler: F) -> Result<(), Box<dyn std::error::Error>>
+pub(crate) fn run_with_args_and_extras<F>(
+    args: CliArgs,
+    io_handler: F,
+    extra_login_actions: Vec<LoginAction>,
+    extra_directory_shares: Vec<DirectoryShare>,
+) -> Result<(), Box<dyn std::error::Error>>
 where
     F: FnOnce(Arc<OutputMonitor>, OwnedFd, OwnedFd) -> IoContext,
 {
@@ -233,6 +230,8 @@ where
         }
     }
 
+    directory_shares.extend(extra_directory_shares);
+
     for spec in &args.mounts {
         directory_shares.push(DirectoryShare::from_mount_spec(spec)?);
     }
@@ -241,6 +240,8 @@ where
         login_actions.push(motd_action);
     }
 
+    login_actions.extend(extra_login_actions);
+
     // Any user-provided login actions must come after our system ones
     login_actions.extend(args.login_actions);
 
@@ -461,18 +462,6 @@ fn motd_login_action(directory_shares: &[DirectoryShare]) -> Option<LoginAction>
     }
 
     let mut output = String::new();
-    output.push_str(
-        "
-░▒▓█▓▒░░▒▓█▓▒░▒▓█▓▒░▒▓███████▓▒░░▒▓████████▓▒░ 
-░▒▓█▓▒░░▒▓█▓▒░▒▓█▓▒░▒▓█▓▒░░▒▓█▓▒░▒▓█▓▒░        
- ░▒▓█▓▒▒▓█▓▒░░▒▓█▓▒░▒▓█▓▒░░▒▓█▓▒░▒▓█▓▒░        
- ░▒▓█▓▒▒▓█▓▒░░▒▓█▓▒░▒▓███████▓▒░░▒▓██████▓▒░   
-  ░▒▓█▓▓█▓▒░ ░▒▓█▓▒░▒▓█▓▒░░▒▓█▓▒░▒▓█▓▒░        
-  ░▒▓█▓▓█▓▒░ ░▒▓█▓▒░▒▓█▓▒░░▒▓█▓▒░▒▓█▓▒░        
-   ░▒▓██▓▒░  ░▒▓█▓▒░▒▓███████▓▒░░▒▓████████▓▒░
-
-",
-    );
     output.push_str(&format!(
         "{host_header:<host_width$}  {guest_header:<guest_width$}  {mode_header}\n",
         host_width = host_width
@@ -548,6 +537,49 @@ impl OutputMonitor {
     }
 }
 
+#[derive(Debug)]
+pub struct IoControl {
+    forward_input: AtomicBool,
+    forward_output: AtomicBool,
+    restore_terminal: AtomicBool,
+}
+
+impl IoControl {
+    pub fn new() -> Arc<Self> {
+        Arc::new(Self {
+            forward_input: AtomicBool::new(true),
+            forward_output: AtomicBool::new(true),
+            restore_terminal: AtomicBool::new(false),
+        })
+    }
+
+    pub fn set_forward_input(&self, enabled: bool) {
+        self.forward_input.store(enabled, Ordering::SeqCst);
+    }
+
+    pub fn set_forward_output(&self, enabled: bool) {
+        self.forward_output.store(enabled, Ordering::SeqCst);
+    }
+
+    pub fn request_terminal_restore(&self) {
+        self.restore_terminal.store(true, Ordering::SeqCst);
+    }
+
+    fn forward_input(&self) -> bool {
+        self.forward_input.load(Ordering::SeqCst)
+    }
+
+    fn forward_output(&self) -> bool {
+        self.forward_output.load(Ordering::SeqCst)
+    }
+
+    fn take_restore_terminal(&self) -> bool {
+        self.restore_terminal
+            .compare_exchange(true, false, Ordering::SeqCst, Ordering::SeqCst)
+            .is_ok()
+    }
+}
+
 fn ensure_base_image(
     base_raw: &Path,
     base_compressed: &Path,
@@ -667,14 +699,17 @@ pub fn create_pipe() -> (OwnedFd, OwnedFd) {
     (read_stream.into(), write_stream.into())
 }
 
-pub fn spawn_vm_io_with_line_handler<F>(
+pub fn spawn_vm_io_with_hooks<F, G>(
     output_monitor: Arc<OutputMonitor>,
     vm_output_fd: OwnedFd,
     vm_input_fd: OwnedFd,
+    io_control: Arc<IoControl>,
     mut on_line: F,
+    mut on_output: G,
 ) -> IoContext
 where
     F: FnMut(&str) -> bool + ::std::marker::Send + 'static,
+    G: FnMut(&[u8]) + ::std::marker::Send + 'static,
 {
     let (input_tx, input_rx): (Sender<VmInput>, Receiver<VmInput>) = mpsc::channel();
 
@@ -721,17 +756,36 @@ where
         }
     }
 
+    fn poll_wakeup_only(wakeup_fd: RawFd, timeout_ms: i32) -> bool {
+        let mut fds = [libc::pollfd {
+            fd: wakeup_fd,
+            events: libc::POLLIN,
+            revents: 0,
+        }];
+
+        let ret = unsafe { libc::poll(fds.as_mut_ptr(), 1, timeout_ms) };
+        ret > 0 && fds[0].revents & libc::POLLIN != 0
+    }
+
     // Copies from stdin to the VM; also polls wakeup_read to exit the thread when it's time to shutdown.
     let stdin_thread = thread::spawn({
         let input_tx = input_tx.clone();
         let raw_guard = raw_guard.clone();
         let wakeup_read = wakeup_read.try_clone().unwrap();
+        let io_control = io_control.clone();
 
         move || {
             let mut buf = [0u8; 1024];
             let mut pending_command: Vec<u8> = Vec::new();
             let mut command_mode = false;
             loop {
+                if !io_control.forward_input() {
+                    if poll_wakeup_only(wakeup_read.as_raw_fd(), 100) {
+                        break;
+                    }
+                    continue;
+                }
+
                 match poll_with_wakeup(libc::STDIN_FILENO, wakeup_read.as_raw_fd(), &mut buf) {
                     PollResult::Shutdown | PollResult::Error => break,
                     PollResult::Spurious => continue,
@@ -776,28 +830,36 @@ where
     let stdout_thread = thread::spawn({
         let raw_guard = raw_guard.clone();
         let wakeup_read = wakeup_read.try_clone().unwrap();
+        let io_control = io_control.clone();
 
         move || {
-            let mut stdout = std::io::stdout().lock();
             let mut buf = [0u8; 1024];
             loop {
+                if io_control.take_restore_terminal() {
+                    let mut guard = raw_guard.lock().unwrap();
+                    *guard = None;
+                }
                 match poll_with_wakeup(vm_output_fd.as_raw_fd(), wakeup_read.as_raw_fd(), &mut buf)
                 {
                     PollResult::Shutdown | PollResult::Error => break,
                     PollResult::Spurious => continue,
                     PollResult::Ready(bytes) => {
-                        // enable raw mode, if we haven't already
-                        if raw_guard.lock().unwrap().is_none()
-                            && let Ok(guard) = enable_raw_mode(libc::STDIN_FILENO)
-                        {
-                            *raw_guard.lock().unwrap() = Some(guard);
-                        }
+                        if io_control.forward_output() {
+                            // enable raw mode, if we haven't already
+                            if raw_guard.lock().unwrap().is_none()
+                                && let Ok(guard) = enable_raw_mode(libc::STDIN_FILENO)
+                            {
+                                *raw_guard.lock().unwrap() = Some(guard);
+                            }
 
-                        if stdout.write_all(bytes).is_err() {
-                            break;
+                            let mut stdout = std::io::stdout().lock();
+                            if stdout.write_all(bytes).is_err() {
+                                break;
+                            }
+                            let _ = stdout.flush();
                         }
-                        let _ = stdout.flush();
                         output_monitor.push(bytes);
+                        on_output(bytes);
                     }
                 }
             }
@@ -829,6 +891,25 @@ where
     }
 }
 
+pub fn spawn_vm_io_with_line_handler<F>(
+    output_monitor: Arc<OutputMonitor>,
+    vm_output_fd: OwnedFd,
+    vm_input_fd: OwnedFd,
+    on_line: F,
+) -> IoContext
+where
+    F: FnMut(&str) -> bool + ::std::marker::Send + 'static,
+{
+    spawn_vm_io_with_hooks(
+        output_monitor,
+        vm_output_fd,
+        vm_input_fd,
+        IoControl::new(),
+        on_line,
+        |_| {},
+    )
+}
+
 pub fn spawn_vm_io(
     output_monitor: Arc<OutputMonitor>,
     vm_output_fd: OwnedFd,