alkdev/alknet
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Implement wraith serve CLI subcommand with clap

Browse Source 
Add serve subcommand with all flags matching server.md CLI interface:
--key, --authorized-keys, --cert-authority, --transport, --listen,
--tls-cert, --tls-key, --acme-domain, --stealth, --proxy,
--iroh-relay, --max-connections-per-ip, --max-auth-attempts.

--key is required, --transport defaults to tcp, --listen defaults to
0.0.0.0:22. --stealth validates TLS transport. --acme-domain requires
acme feature flag. --transport iroh prints endpoint ID on startup.
Key inputs accept file paths. Errors reported to stderr with non-zero
exit code. Also adds acme feature flag and rustls-pemfile/rustls-pki-types
dependencies for TLS cert loading.
This commit is contained in:
glm-5.1
2026-06-02 12:28:37 +00:00
parent c7b8c5c5e0
commit 62d57dd477
3 changed files with 444 additions and 110 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
Cargo.lock generated
View File

@@ -3863,6 +3863,15 @@ dependencies = [
"x509-parser 0.16.0",
]
[[package]]
name = "rustls-pemfile"
version = "2.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "dce314e5fee3f39953d46bb63bb8a46d40c2f8fb7cc5a3b6cab2bde9721d6e50"
dependencies = [
"rustls-pki-types",
]
[[package]]
name = "rustls-pki-types"
version = "1.14.1"
@@ -5586,6 +5595,10 @@ dependencies = [
"anyhow",
"clap",
"iroh",
"rustls",
"rustls-acme",
"rustls-pemfile",
"rustls-pki-types",
"tokio",
"url",
"wraith-core",

7
crates/wraith/Cargo.toml
View File

@@ -9,8 +9,9 @@ path = "src/main.rs"
[features]
default = ["tls", "iroh"]
tls = ["wraith-core/tls"]
tls = ["wraith-core/tls", "dep:rustls-pemfile", "dep:rustls-pki-types"]
iroh = ["wraith-core/iroh", "dep:iroh", "dep:url"]
acme = ["wraith-core/acme", "dep:rustls-acme", "dep:rustls", "tls"]
[dependencies]
wraith-core = { path = "../wraith-core" }
@@ -19,3 +20,7 @@ tokio = { version = "1", features = ["full"] }
anyhow = "1"
iroh = { version = "0.34", optional = true }
url = { version = "2", optional = true }
rustls-acme = { version = "0.12", optional = true }
rustls = { version = "0.23", optional = true, features = ["aws_lc_rs"] }
rustls-pemfile = { version = "2", optional = true }
rustls-pki-types = { version = "1", optional = true }

344
crates/wraith/src/main.rs
View File

@@ -6,15 +6,16 @@ use anyhow::{anyhow, Result};
use clap::{Parser, Subcommand, ValueEnum};
use wraith_core::auth::keys::KeySource;
use wraith_core::client::{ConnectOptions, TransportMode};
use wraith_core::server::{ServeOptions, ServeTransportMode, Server};
#[cfg(feature = "iroh")]
use wraith_core::transport::IrohTransport;
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")]
#[command(name = "wraith", version, about = "Wraith SSH tunnel tool")]
struct Cli {
#[command(subcommand)]
command: Commands,
@@ -22,12 +23,21 @@ struct Cli {
#[derive(Subcommand)]
enum Commands {
#[command(about = "Connect to a wraith server and start a SOCKS5 proxy / port forwarding session")]
#[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")]
#[arg(
long,
help = "TCP/TLS server address (required for tcp/tls transport)",
env = "WRAITH_SERVER"
)]
server: Option<String>,
#[arg(long, help = "iroh endpoint ID, base58-encoded (required for iroh transport)")]
#[arg(
long,
help = "iroh endpoint ID, base58-encoded (required for iroh transport)"
)]
peer: Option<String>,
#[arg(long, value_enum, default_value = "tcp", help = "Transport mode")]
@@ -57,6 +67,68 @@ enum Commands {
#[arg(long, help = "Accept self-signed TLS certs")]
insecure: bool,
},
#[command(about = "Start the wraith server (accept SSH connections)")]
Serve {
#[arg(long, help = "SSH host key path (required)")]
key: String,
#[arg(long, help = "Authorized keys file path")]
authorized_keys: Option<String>,
#[arg(long, help = "CA public key for certificate authority auth")]
cert_authority: Option<String>,
#[arg(
long,
value_enum,
default_value = "tcp",
help = "Transport mode (tcp, tls, iroh)"
)]
transport: ServeTransportModeArg,
#[arg(
long,
default_value = "0.0.0.0:22",
help = "Listen address for TCP/TLS"
)]
listen: String,
#[arg(long, help = "TLS certificate path (manual)")]
tls_cert: Option<String>,
#[arg(long, help = "TLS private key path (manual)")]
tls_key: Option<String>,
#[arg(long, help = "ACME auto-cert domain")]
acme_domain: Option<String>,
#[arg(
long,
help = "Serve fake nginx 404 to non-SSH connections (requires --transport tls)"
)]
stealth: bool,
#[arg(long, help = "Outbound proxy URL (socks5:// or http://)")]
proxy: Option<String>,
#[arg(long, help = "iroh relay server URL")]
iroh_relay: Option<String>,
#[arg(
long,
default_value_t = 0,
help = "Max concurrent connections per IP (0 = unlimited)"
)]
max_connections_per_ip: usize,
#[arg(
long,
default_value_t = 10,
help = "Max auth failures before disconnect"
)]
max_auth_attempts: usize,
},
}
#[derive(Clone, Debug, ValueEnum)]
@@ -76,6 +148,23 @@ impl From<TransportModeArg> for TransportMode {
}
}
#[derive(Clone, Debug, ValueEnum)]
enum ServeTransportModeArg {
Tcp,
Tls,
Iroh,
}
impl From<ServeTransportModeArg> for ServeTransportMode {
fn from(val: ServeTransportModeArg) -> Self {
match val {
ServeTransportModeArg::Tcp => ServeTransportMode::Tcp,
ServeTransportModeArg::Tls => ServeTransportMode::Tls,
ServeTransportModeArg::Iroh => ServeTransportMode::Iroh,
}
}
}
#[tokio::main]
async fn main() {
if let Err(e) = run().await {
@@ -101,6 +190,70 @@ async fn run() -> Result<()> {
tls_server_name,
insecure,
} => {
run_connect(
server,
peer,
transport,
identity,
socks5,
forward,
remote_forward,
proxy,
iroh_relay,
tls_server_name,
insecure,
)
.await
}
Commands::Serve {
key,
authorized_keys,
cert_authority,
transport,
listen,
tls_cert,
tls_key,
acme_domain,
stealth,
proxy,
iroh_relay,
max_connections_per_ip,
max_auth_attempts,
} => {
run_serve(
key,
authorized_keys,
cert_authority,
transport,
listen,
tls_cert,
tls_key,
acme_domain,
stealth,
proxy,
iroh_relay,
max_connections_per_ip,
max_auth_attempts,
)
.await
}
}
}
#[allow(clippy::too_many_arguments)]
async fn run_connect(
server: Option<String>,
peer: Option<String>,
transport: TransportModeArg,
identity: Option<String>,
socks5: String,
forward: Vec<String>,
remote_forward: Vec<String>,
proxy: Option<String>,
iroh_relay: Option<String>,
tls_server_name: Option<String>,
insecure: bool,
) -> Result<()> {
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());
@@ -155,7 +308,9 @@ async fn run() -> Result<()> {
TransportMode::Tls => {
#[cfg(not(feature = "tls"))]
{
return Err(anyhow!("TLS transport is not available (wraith-core built without 'tls' feature)"));
Err(anyhow!(
"TLS transport is not available (wraith-core built without 'tls' feature)"
))
}
#[cfg(feature = "tls")]
{
@@ -176,7 +331,9 @@ async fn run() -> Result<()> {
TransportMode::Iroh => {
#[cfg(not(feature = "iroh"))]
{
return Err(anyhow!("iroh transport is not available (wraith-core built without 'iroh' feature)"));
Err(anyhow!(
"iroh transport is not available (wraith-core built without 'iroh' feature)"
))
}
#[cfg(feature = "iroh")]
{
@@ -195,10 +352,7 @@ async fn run() -> Result<()> {
None => None,
};
let proxy_url: Option<url::Url> = match proxy.as_deref() {
Some(u) => Some(
u.parse()
.map_err(|e| anyhow!("invalid proxy URL: {e}"))?,
),
Some(u) => Some(u.parse().map_err(|e| anyhow!("invalid proxy URL: {e}"))?),
None => None,
};
let t = Arc::new(
@@ -211,8 +365,6 @@ async fn run() -> Result<()> {
}
}
}
}
}
async fn connect_and_run<T: Transport>(opts: ConnectOptions, transport: Arc<T>) -> Result<()> {
wraith_core::client::ClientSession::new(opts, transport)
@@ -222,3 +374,167 @@ async fn connect_and_run<T: Transport>(opts: ConnectOptions, transport: Arc<T>)
.await
.map_err(|e| anyhow!("{e}"))
}
#[allow(clippy::too_many_arguments)]
async fn run_serve(
key: String,
authorized_keys: Option<String>,
cert_authority: Option<String>,
transport: ServeTransportModeArg,
listen: String,
tls_cert: Option<String>,
tls_key: Option<String>,
acme_domain: Option<String>,
stealth: bool,
proxy: Option<String>,
iroh_relay: Option<String>,
max_connections_per_ip: usize,
max_auth_attempts: usize,
) -> Result<()> {
let transport_mode: ServeTransportMode = transport.into();
if acme_domain.is_some() {
#[cfg(not(feature = "acme"))]
{
return Err(anyhow!(
"ACME support is not available (wraith built without 'acme' feature)"
));
}
}
if stealth && transport_mode != ServeTransportMode::Tls {
return Err(anyhow!(
"stealth mode requires TLS transport (--transport tls)"
));
}
let mut opts = ServeOptions::new(KeySource::File(key.into()))
.transport_mode(transport_mode.clone())
.listen_addr(&listen)
.stealth(stealth)
.max_connections_per_ip(max_connections_per_ip)
.max_auth_attempts(max_auth_attempts);
if let Some(ref path) = authorized_keys {
opts = opts.authorized_keys(KeySource::File(path.into()));
}
if let Some(ref path) = cert_authority {
opts = opts.cert_authority(KeySource::File(path.into()));
}
if let Some(ref path) = tls_cert {
opts = opts.tls_cert(path);
}
if let Some(ref path) = tls_key {
opts = opts.tls_key(path);
}
if let Some(ref domain) = acme_domain {
opts = opts.acme_domain(domain);
}
if let Some(ref url) = proxy {
opts = opts.proxy(url);
}
if let Some(ref url) = iroh_relay {
opts = opts.iroh_relay(url);
}
opts.validate().map_err(|e| anyhow!("{e}"))?;
let server = Server::new(opts).map_err(|e| anyhow!("{e}"))?;
match transport_mode {
ServeTransportMode::Tcp => {
let addr: SocketAddr = listen
.parse()
.map_err(|e| anyhow!("invalid listen address: {e}"))?;
let acceptor = wraith_core::transport::TcpAcceptor::bind(addr)
.await
.map_err(|e| anyhow!("bind failed: {e}"))?;
server.run(acceptor, None).await.map_err(|e| anyhow!("{e}"))
}
ServeTransportMode::Tls => {
#[cfg(not(feature = "tls"))]
{
Err(anyhow!(
"TLS transport is not available (wraith-core built without 'tls' feature)"
))
}
#[cfg(feature = "acme")]
{
if let Some(ref domain) = acme_domain {
let addr: SocketAddr = listen
.parse()
.map_err(|e| anyhow!("invalid listen address: {e}"))?;
let provider = Arc::new(
wraith_core::transport::AcmeCertProvider::domain(domain)
.with_production_directory(),
);
let acceptor =
wraith_core::transport::AcmeTlsAcceptor::bind_acme(addr, provider)
.await
.map_err(|e| anyhow!("ACME bind failed: {e}"))?;
return server.run(acceptor, None).await.map_err(|e| anyhow!("{e}"));
}
}
#[cfg(feature = "tls")]
{
use rustls_pki_types::{CertificateDer, PrivateKeyDer};
let addr: SocketAddr = listen
.parse()
.map_err(|e| anyhow!("invalid listen address: {e}"))?;
let cert_path = tls_cert.ok_or_else(|| {
anyhow!("--tls-cert is required for TLS transport (or use --acme-domain)")
})?;
let key_path = tls_key.ok_or_else(|| {
anyhow!("--tls-key is required for TLS transport (or use --acme-domain)")
})?;
let cert_data = std::fs::read(&cert_path)
.map_err(|e| anyhow!("failed to read TLS cert '{}': {e}", cert_path))?;
let key_data = std::fs::read(&key_path)
.map_err(|e| anyhow!("failed to read TLS key '{}': {e}", key_path))?;
let certs: Vec<CertificateDer<'static>> =
rustls_pemfile::certs(&mut &cert_data[..])
.collect::<Result<Vec<_>, _>>()
.map_err(|e| anyhow!("failed to parse TLS certificates: {e}"))?;
let key: PrivateKeyDer<'static> = rustls_pemfile::private_key(&mut &key_data[..])
.map_err(|e| anyhow!("failed to parse TLS private key: {e}"))?
.ok_or_else(|| anyhow!("no private key found in {}", key_path))?;
let acceptor = wraith_core::transport::TlsAcceptor::bind(addr, certs, key, None)
.await
.map_err(|e| anyhow!("TLS bind failed: {e}"))?;
server.run(acceptor, None).await.map_err(|e| anyhow!("{e}"))
}
}
ServeTransportMode::Iroh => {
#[cfg(not(feature = "iroh"))]
{
Err(anyhow!(
"iroh transport is not available (wraith-core built without 'iroh' feature)"
))
}
#[cfg(feature = "iroh")]
{
use iroh::RelayUrl;
let relay_url: Option<RelayUrl> = 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<url::Url> = match proxy.as_deref() {
Some(u) => Some(u.parse().map_err(|e| anyhow!("invalid proxy URL: {e}"))?),
None => None,
};
let acceptor = wraith_core::transport::IrohAcceptor::bind(relay_url, proxy_url)
.await
.map_err(|e| anyhow!("iroh bind failed: {e}"))?;
let endpoint_id = acceptor.endpoint_id();
eprintln!("iroh endpoint ID: {endpoint_id}");
server
.run(acceptor, Some(&endpoint_id))
.await
.map_err(|e| anyhow!("{e}"))
}
}
}
}