diff --git a/Cargo.lock b/Cargo.lock index dc6a356..11c430d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5585,7 +5585,9 @@ version = "0.1.0" dependencies = [ "anyhow", "clap", + "iroh", "tokio", + "url", "wraith-core", ] diff --git a/crates/wraith/Cargo.toml b/crates/wraith/Cargo.toml index 9e57c36..bb52c54 100644 --- a/crates/wraith/Cargo.toml +++ b/crates/wraith/Cargo.toml @@ -7,8 +7,15 @@ edition = "2021" name = "wraith" path = "src/main.rs" +[features] +default = ["tls", "iroh"] +tls = ["wraith-core/tls"] +iroh = ["wraith-core/iroh", "dep:iroh", "dep:url"] + [dependencies] wraith-core = { path = "../wraith-core" } -clap = { version = "4", features = ["derive"] } +clap = { version = "4", features = ["derive", "env"] } tokio = { version = "1", features = ["full"] } -anyhow = "1" \ No newline at end of file +anyhow = "1" +iroh = { version = "0.34", optional = true } +url = { version = "2", optional = true } \ No newline at end of file diff --git a/crates/wraith/src/main.rs b/crates/wraith/src/main.rs index e71fdf5..eea1d15 100644 --- a/crates/wraith/src/main.rs +++ b/crates/wraith/src/main.rs @@ -1 +1,224 @@ -fn main() {} \ No newline at end of file +use std::net::SocketAddr; +use std::process; +use std::sync::Arc; + +use anyhow::{anyhow, Result}; +use clap::{Parser, Subcommand, ValueEnum}; +use wraith_core::auth::keys::KeySource; +use wraith_core::client::{ConnectOptions, TransportMode}; +use wraith_core::transport::TcpTransport; +#[cfg(feature = "tls")] +use wraith_core::transport::TlsTransport; +#[cfg(feature = "iroh")] +use wraith_core::transport::IrohTransport; +use wraith_core::transport::Transport; + +#[derive(Parser)] +#[command(name = "wraith", version, about = "Wraith SSH tunnel client")] +struct Cli { + #[command(subcommand)] + command: Commands, +} + +#[derive(Subcommand)] +enum Commands { + #[command(about = "Connect to a wraith server and start a SOCKS5 proxy / port forwarding session")] + Connect { + #[arg(long, help = "TCP/TLS server address (required for tcp/tls transport)", env = "WRAITH_SERVER")] + server: Option, + + #[arg(long, help = "iroh endpoint ID, base58-encoded (required for iroh transport)")] + peer: Option, + + #[arg(long, value_enum, default_value = "tcp", help = "Transport mode")] + transport: TransportModeArg, + + #[arg(long, help = "SSH private key path", env = "WRAITH_IDENTITY")] + identity: Option, + + #[arg(long, default_value = "127.0.0.1:1080", help = "SOCKS5 listen address")] + socks5: String, + + #[arg(long, action = clap::ArgAction::Append, help = "Port forward spec (repeatable, e.g. 5432:db:5432)")] + forward: Vec, + + #[arg(long, action = clap::ArgAction::Append, help = "Remote port forward spec (repeatable)")] + remote_forward: Vec, + + #[arg(long, help = "Upstream proxy URL (socks5:// or http://)")] + proxy: Option, + + #[arg(long, help = "iroh relay URL")] + iroh_relay: Option, + + #[arg(long, help = "SNI hostname for TLS")] + tls_server_name: Option, + + #[arg(long, help = "Accept self-signed TLS certs")] + insecure: bool, + }, +} + +#[derive(Clone, Debug, ValueEnum)] +enum TransportModeArg { + Tcp, + Tls, + Iroh, +} + +impl From for TransportMode { + fn from(val: TransportModeArg) -> Self { + match val { + TransportModeArg::Tcp => TransportMode::Tcp, + TransportModeArg::Tls => TransportMode::Tls, + TransportModeArg::Iroh => TransportMode::Iroh, + } + } +} + +#[tokio::main] +async fn main() { + if let Err(e) = run().await { + eprintln!("error: {e}"); + process::exit(1); + } +} + +async fn run() -> Result<()> { + let cli = Cli::parse(); + + match cli.command { + Commands::Connect { + server, + peer, + transport, + identity, + socks5, + forward, + remote_forward, + proxy, + iroh_relay, + tls_server_name, + insecure, + } => { + let identity_val = identity + .ok_or_else(|| anyhow!("--identity is required (or set WRAITH_IDENTITY env var)"))?; + let key_source = KeySource::File(identity_val.into()); + + let transport_mode: TransportMode = transport.into(); + + if proxy.is_some() && matches!(transport_mode, TransportMode::Tcp) { + eprintln!("warning: --proxy with --transport tcp is effectively a no-op (TCP transport is already a direct connection); use the SOCKS5 server instead"); + } + + let mut opts = ConnectOptions::new(key_source) + .transport_mode(transport_mode.clone()) + .socks5_addr(&socks5); + + if let Some(ref s) = server { + opts = opts.server(s); + } + if let Some(ref p) = peer { + opts = opts.peer(p); + } + for fwd in &forward { + opts = opts.forward(fwd); + } + for rfwd in &remote_forward { + opts = opts.remote_forward(rfwd); + } + if let Some(ref p) = proxy { + opts = opts.proxy(p); + } + if let Some(ref r) = iroh_relay { + opts = opts.iroh_relay(r); + } + if let Some(ref n) = tls_server_name { + opts = opts.tls_server_name(n); + } + if insecure { + opts = opts.insecure(true); + } + + opts.validate().map_err(|e| anyhow!("{e}"))?; + + match transport_mode { + TransportMode::Tcp => { + let addr: SocketAddr = server + .as_deref() + .ok_or_else(|| anyhow!("--server is required for tcp transport"))? + .parse() + .map_err(|e| anyhow!("invalid server address: {e}"))?; + let t = Arc::new(TcpTransport::new(addr)); + connect_and_run(opts, t).await + } + TransportMode::Tls => { + #[cfg(not(feature = "tls"))] + { + return Err(anyhow!("TLS transport is not available (wraith-core built without 'tls' feature)")); + } + #[cfg(feature = "tls")] + { + let addr: SocketAddr = server + .as_deref() + .ok_or_else(|| anyhow!("--server is required for tls transport"))? + .parse() + .map_err(|e| anyhow!("invalid server address: {e}"))?; + let mut t = TlsTransport::new(addr); + if let Some(ref n) = tls_server_name { + t = t.with_server_name(n); + } + t = t.with_insecure(insecure); + let t = Arc::new(t); + connect_and_run(opts, t).await + } + } + TransportMode::Iroh => { + #[cfg(not(feature = "iroh"))] + { + return Err(anyhow!("iroh transport is not available (wraith-core built without 'iroh' feature)")); + } + #[cfg(feature = "iroh")] + { + use iroh::{NodeId, RelayUrl}; + let node_id_str = peer + .as_deref() + .ok_or_else(|| anyhow!("--peer is required for iroh transport"))?; + let node_id: NodeId = node_id_str + .parse() + .map_err(|e| anyhow!("invalid iroh peer endpoint ID: {e}"))?; + let relay_url: Option = match iroh_relay.as_deref() { + Some(u) => Some( + u.parse() + .map_err(|e| anyhow!("invalid iroh relay URL: {e}"))?, + ), + None => None, + }; + let proxy_url: Option = match proxy.as_deref() { + Some(u) => Some( + u.parse() + .map_err(|e| anyhow!("invalid proxy URL: {e}"))?, + ), + None => None, + }; + let t = Arc::new( + IrohTransport::new(node_id, relay_url, proxy_url) + .await + .map_err(|e| anyhow!("failed to create iroh transport: {e}"))?, + ); + connect_and_run(opts, t).await + } + } + } + } + } +} + +async fn connect_and_run(opts: ConnectOptions, transport: Arc) -> Result<()> { + wraith_core::client::ClientSession::new(opts, transport) + .await + .map_err(|e| anyhow!("{e}"))? + .run() + .await + .map_err(|e| anyhow!("{e}")) +} \ No newline at end of file