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};
#[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, 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 (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;
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("")
),
}
}