use std::collections::BTreeMap;
use std::path::PathBuf;
use std::time::Duration;
use anyhow::{Context, Result, bail};
use clap::{ArgAction, Parser};
use tokio::sync::mpsc;
use torzy::{DEFAULT_COMMAND, LaunchConfig, TorzyEvent};
#[cfg(feature = "tui")]
use std::io::{self, IsTerminal, Write};
#[cfg(feature = "tui")]
use crossterm::{
ExecutableCommand, QueueableCommand, cursor,
style::{self, Stylize},
terminal::{self, ClearType},
};
#[derive(Parser, Debug)]
#[command(
name = "torzy",
version,
about = "Launch isolated workers through separate Tor SOCKS routes"
)]
struct Cli {
#[arg(short = 'n', long, default_value_t = 5)]
envs: usize,
#[arg(
short = 'c',
long,
value_name = "COMMAND",
help = "Worker command. Defaults to a local ready message. Supports {TOR_ROUTE}, {TOR_WORKER_ID}, {TORZY_ROOT}, {TORZY_WORKSPACE}."
)]
command: Option<String>,
#[arg(
short = 'r',
long,
value_name = "ROUTE",
value_delimiter = ',',
help = "Comma-separated SOCKS routes; one per worker. Omit to auto-generate from --base-socks-port."
)]
tor_routes: Vec<String>,
#[arg(long, default_value_t = 9050)]
base_socks_port: u16,
#[arg(short = 'j', long, help = "Max workers running concurrently")]
max_concurrent: Option<usize>,
#[arg(short = 't', long, default_value_t = 120)]
timeout_seconds: u64,
#[arg(
long,
default_value_t = 0,
help = "Memory limit per worker in MB (0 = unlimited)"
)]
memory_limit_mb: u64,
#[arg(
long,
value_name = "DIR",
help = "Root directory for worker and tor data"
)]
root_dir: Option<PathBuf>,
#[arg(
long,
value_name = "DIR",
help = "Run worker commands from this directory instead of each isolated workspace"
)]
working_directory: Option<PathBuf>,
#[arg(long, default_value_t = false)]
keep_workspaces: bool,
#[arg(
long = "no-manage-tor",
action = ArgAction::SetFalse,
default_value_t = true,
help = "Do not start/stop tor; use existing --tor-routes instead"
)]
manage_tor: bool,
#[arg(long, default_value = "tor")]
tor_binary: String,
#[arg(long, default_value_t = 90)]
tor_timeout: u64,
#[arg(long, value_name = "DIR", help = "Persistent tor data directory")]
tor_data_dir: Option<PathBuf>,
#[arg(
long = "no-proxy-env",
action = ArgAction::SetFalse,
default_value_t = true,
help = "Do not export ALL_PROXY/HTTPS_PROXY/HTTP_PROXY/CARGO_HTTP_PROXY to workers"
)]
proxy_env: bool,
#[arg(
long = "env",
value_name = "KEY=VALUE",
help = "Extra environment variable for workers"
)]
extra_env: Vec<String>,
#[arg(long, default_value_t = false)]
dry_run: bool,
#[arg(
long,
default_value_t = false,
help = "Show a sanitized terminal dashboard"
)]
tui: bool,
#[arg(long, value_name = "FILE", help = "Write run summary JSON")]
summary_json: Option<PathBuf>,
#[arg(long, default_value_t = false, help = "Pretty-print summary JSON")]
pretty_json: bool,
}
#[tokio::main]
async fn main() -> Result<()> {
let cli = Cli::parse();
let extra_env = parse_extra_env(cli.extra_env)?;
let config = LaunchConfig {
envs: cli.envs,
command: cli.command.unwrap_or_else(|| DEFAULT_COMMAND.to_string()),
tor_routes: cli.tor_routes,
base_socks_port: cli.base_socks_port,
timeout: Duration::from_secs(cli.timeout_seconds.max(1)),
max_concurrent: cli.max_concurrent,
memory_limit_mb: cli.memory_limit_mb,
root_dir: cli.root_dir,
keep_workspaces: cli.keep_workspaces,
manage_tor: cli.manage_tor,
tor_binary: cli.tor_binary,
tor_timeout: Duration::from_secs(cli.tor_timeout.max(1)),
tor_data_dir: cli.tor_data_dir,
dry_run: cli.dry_run,
proxy_env: cli.proxy_env,
working_directory: cli.working_directory,
extra_env,
};
let summary = if cli.tui {
run_tui(config).await?
} else {
run_plain(config).await?
};
if let Some(path) = cli.summary_json {
let payload = if cli.pretty_json {
serde_json::to_string_pretty(&summary)?
} else {
serde_json::to_string(&summary)?
};
tokio::fs::write(&path, payload)
.await
.context("writing summary JSON")?;
println!("Summary written.");
}
println!(
"Done. {}/{} workers succeeded.",
summary.succeeded, summary.total_envs
);
if summary.failed > 0 || summary.timed_out > 0 || summary.killed > 0 {
std::process::exit(1);
}
Ok(())
}
fn parse_extra_env(values: Vec<String>) -> Result<BTreeMap<String, String>> {
let mut env = BTreeMap::new();
for value in values {
let Some((key, val)) = value.split_once('=') else {
bail!("--env must be KEY=VALUE, got: {value}");
};
if key.trim().is_empty() {
bail!("--env key cannot be empty");
}
env.insert(key.to_string(), val.to_string());
}
Ok(env)
}
fn print_event(event: &TorzyEvent) {
match event {
TorzyEvent::TorStarted { id } => println!("[tor-{id}] started"),
TorzyEvent::TorBootstrap { id, pct } => println!("[tor-{id}] bootstrap {pct}%"),
TorzyEvent::TorReady { id } => println!("[tor-{id}] ready"),
TorzyEvent::TorFailed { id, detail } => println!("[tor-{id}] failed: {detail}"),
TorzyEvent::TorStopped { id } => println!("[tor-{id}] stopped"),
TorzyEvent::WorkersReady => println!("[torzy] workers ready"),
TorzyEvent::WorkerStarted { id } => println!("[worker-{id}] started"),
TorzyEvent::WorkerOutput { id, stream, line } => {
println!("[worker-{id} {}] {line}", stream.as_str());
}
TorzyEvent::WorkerFinished {
id,
status,
exit_code,
detail,
} => println!(
"[worker-{id}] finished status={} exit={:?} {}",
status.as_str(),
exit_code,
detail.as_deref().unwrap_or("")
),
}
}
async fn run_plain(config: LaunchConfig) -> Result<torzy::RunSummary> {
let (tx, mut rx) = mpsc::unbounded_channel();
let printer = tokio::spawn(async move {
while let Some(event) = rx.recv().await {
print_event(&event);
}
});
let summary = torzy::run(config, Some(tx)).await?;
let _ = printer.await;
Ok(summary)
}
#[cfg(not(feature = "tui"))]
async fn run_tui(_config: LaunchConfig) -> Result<torzy::RunSummary> {
bail!("this binary was built without the tui feature")
}
#[cfg(feature = "tui")]
async fn run_tui(config: LaunchConfig) -> Result<torzy::RunSummary> {
let mut stdout = io::stdout();
if !stdout.is_terminal() {
return run_plain(config).await;
}
let _guard = TerminalGuard::enter(&mut stdout)?;
let total = config.envs.max(1);
let (tx, mut rx) = mpsc::unbounded_channel();
let runner = tokio::spawn(async move { torzy::run(config, Some(tx)).await });
let mut state = TuiState::new(total);
let mut tick = tokio::time::interval(Duration::from_millis(120));
loop {
tokio::select! {
event = rx.recv() => {
if let Some(event) = event {
state.apply(event);
} else if runner.is_finished() {
break;
}
}
_ = tick.tick() => {}
}
draw_tui(&mut stdout, &state)?;
if runner.is_finished() && state.completed() >= total {
break;
}
}
draw_tui(&mut stdout, &state)?;
runner.await.context("joining torzy runner")?
}
#[cfg(feature = "tui")]
struct TerminalGuard;
#[cfg(feature = "tui")]
impl TerminalGuard {
fn enter(stdout: &mut io::Stdout) -> Result<Self> {
terminal::enable_raw_mode().context("enabling raw mode")?;
stdout
.execute(terminal::EnterAlternateScreen)
.context("entering alternate screen")?;
stdout.execute(cursor::Hide).context("hiding cursor")?;
Ok(Self)
}
}
#[cfg(feature = "tui")]
impl Drop for TerminalGuard {
fn drop(&mut self) {
let mut stdout = io::stdout();
let _ = stdout.execute(cursor::Show);
let _ = stdout.execute(terminal::LeaveAlternateScreen);
let _ = terminal::disable_raw_mode();
}
}
#[cfg(feature = "tui")]
#[derive(Clone)]
struct TuiWorker {
tor_pct: u8,
status: String,
exit_code: Option<i32>,
detail: Option<String>,
}
#[cfg(feature = "tui")]
struct TuiState {
workers: Vec<TuiWorker>,
logs: Vec<String>,
workers_ready: bool,
}
#[cfg(feature = "tui")]
impl TuiState {
fn new(total: usize) -> Self {
Self {
workers: vec![
TuiWorker {
tor_pct: 0,
status: "queued".to_string(),
exit_code: None,
detail: None,
};
total
],
logs: Vec::new(),
workers_ready: false,
}
}
fn apply(&mut self, event: TorzyEvent) {
match event {
TorzyEvent::TorStarted { id } => self.set_status(id, "tor starting"),
TorzyEvent::TorBootstrap { id, pct } => {
if let Some(worker) = self.workers.get_mut(id) {
worker.tor_pct = pct;
worker.status = "tor bootstrap".to_string();
}
}
TorzyEvent::TorReady { id } => {
if let Some(worker) = self.workers.get_mut(id) {
worker.tor_pct = 100;
worker.status = "tor ready".to_string();
}
}
TorzyEvent::TorFailed { id, detail } => {
if let Some(worker) = self.workers.get_mut(id) {
worker.status = "tor failed".to_string();
worker.detail = Some(detail);
}
}
TorzyEvent::TorStopped { id } => {
self.push_log(format!("tor-{id} stopped"));
}
TorzyEvent::WorkersReady => {
self.workers_ready = true;
self.push_log("workers ready".to_string());
}
TorzyEvent::WorkerStarted { id } => self.set_status(id, "running"),
TorzyEvent::WorkerOutput { id, stream, line } => {
self.push_log(format!("worker-{id} {}: {line}", stream.as_str()));
}
TorzyEvent::WorkerFinished {
id,
status,
exit_code,
detail,
} => {
if let Some(worker) = self.workers.get_mut(id) {
worker.status = status.as_str().to_string();
worker.exit_code = exit_code;
worker.detail = detail;
}
}
}
}
fn set_status(&mut self, id: usize, status: &str) {
if let Some(worker) = self.workers.get_mut(id) {
worker.status = status.to_string();
}
}
fn push_log(&mut self, line: String) {
if self.logs.len() >= 8 {
self.logs.remove(0);
}
self.logs.push(line);
}
fn completed(&self) -> usize {
self.workers
.iter()
.filter(|worker| {
matches!(
worker.status.as_str(),
"succeeded" | "failed" | "timed_out" | "killed" | "tor failed"
)
})
.count()
}
fn succeeded(&self) -> usize {
self.workers
.iter()
.filter(|worker| worker.status == "succeeded")
.count()
}
}
#[cfg(feature = "tui")]
fn draw_tui(stdout: &mut io::Stdout, state: &TuiState) -> Result<()> {
let (cols, rows) = terminal::size().unwrap_or((100, 30));
stdout.queue(cursor::MoveTo(0, 0))?;
stdout.queue(terminal::Clear(ClearType::All))?;
let title = format!(
"torzy | workers {} | complete {} | ok {}",
state.workers.len(),
state.completed(),
state.succeeded()
);
stdout.queue(style::PrintStyledContent(title.bold()))?;
stdout.queue(cursor::MoveTo(0, 1))?;
stdout.queue(style::Print(if state.workers_ready {
"routes ready; workers launched"
} else {
"preparing tor routes"
}))?;
let table_top = 3;
stdout.queue(cursor::MoveTo(0, table_top))?;
stdout.queue(style::Print("worker tor status detail"))?;
let max_worker_rows = rows.saturating_sub(14).max(3) as usize;
for (row, worker) in state.workers.iter().take(max_worker_rows).enumerate() {
let y = table_top + 1 + row as u16;
stdout.queue(cursor::MoveTo(0, y))?;
let detail = worker.detail.as_deref().unwrap_or("");
let line = format!(
"{:<8} {:>3}% {:<12} {}",
row, worker.tor_pct, worker.status, detail
);
stdout.queue(style::Print(truncate(&line, cols as usize)))?;
}
let log_top = rows.saturating_sub(10).max(table_top + 5);
stdout.queue(cursor::MoveTo(0, log_top))?;
stdout.queue(style::PrintStyledContent(
"recent sanitized output".underlined(),
))?;
for (idx, line) in state.logs.iter().rev().take(8).rev().enumerate() {
stdout.queue(cursor::MoveTo(0, log_top + 1 + idx as u16))?;
stdout.queue(style::Print(truncate(line, cols as usize)))?;
}
stdout.flush()?;
Ok(())
}
#[cfg(feature = "tui")]
fn truncate(value: &str, width: usize) -> String {
if value.chars().count() <= width {
return value.to_string();
}
value
.chars()
.take(width.saturating_sub(1))
.collect::<String>()
+ "~"
}