First cut of a new rama CLI 'probe' command

rama-cli/src/cmd/mod.rs

@@ -4,4 +4,5 @@ pub mod echo;
 pub mod fp;
 pub mod http;
 pub mod ip;
+pub mod probe;
 pub mod proxy;

rama-cli/src/cmd/probe/mod.rs

@@ -0,0 +1,232 @@
+//! rama probe service
+
+use std::{
+    fs::File,
+    io::{BufRead, BufReader},
+    time::Duration,
+};
+
+use clap::Args;
+use rama::{
+    cli::args::RequestArgsBuilder,
+    error::{BoxError, OpaqueError},
+    http::{
+        client::HttpClient,
+        layer::{
+            decompression::DecompressionLayer,
+            follow_redirect::{policy::Limited, FollowRedirectLayer},
+            required_header::AddRequiredRequestHeadersLayer,
+            timeout::TimeoutLayer,
+        },
+        Request, Response, StatusCode,
+    },
+    rt::Executor,
+    service::{layer::MapResultLayer, Context, Layer, Service},
+    utils::graceful::{self, Shutdown, ShutdownGuard},
+};
+use tokio::sync::oneshot;
+use tracing::level_filters::LevelFilter;
+use tracing_subscriber::{fmt, layer::SubscriberExt, util::SubscriberInitExt, EnvFilter};
+
+use crate::error::ErrorWithExitCode;
+mod writer;
+
+// TODO Future features:
+// 1. Add -o to write the output to a file
+// 2. match strings
+// 3. host
+
+/// rama domain prober
+#[derive(Args, Debug, Clone)]
+pub struct CliCommandProbe {
+    #[arg(long)]
+    /// print debug info
+    debug: bool,
+
+    #[arg(short = 'l', long)]
+    /// A file containing domains to probe
+    list: Option<String>,
+
+    #[arg(short = 'v', long)]
+    /// print verbose output
+    verbose: bool,
+
+    // For the first cut I am imagining a usage example like:
+    // rama probe domains.txt
+    // where "domains.txt" is just a file containing simple domains like google.com
+    args: Option<Vec<String>>,
+}
+
+pub(crate) fn pretty_print_status_code(status_code: StatusCode, domain: String) {
+    match status_code.as_u16() {
+        200 => eprintln!("{}: \x1b[32m{}\x1b[0m", domain, status_code), // Green
+        301 => eprintln!("{}: \x1b[33m{}\x1b[0m", domain, status_code), // Yellow
+        400 => eprintln!("{}: \x1b[31m{}\x1b[0m", domain, status_code), // Red
+        _ => eprintln!("{}: {}", domain, status_code),                  // Default color
+    }
+}
+
+/// Run the rama probe command
+pub async fn run(cfg: CliCommandProbe) -> Result<(), BoxError> {
+    tracing_subscriber::registry()
+        .with(fmt::layer())
+        .with(
+            EnvFilter::builder()
+                .with_default_directive(
+                    if cfg.debug {
+                        if cfg.verbose {
+                            LevelFilter::TRACE
+                        } else {
+                            LevelFilter::DEBUG
+                        }
+                    } else {
+                        LevelFilter::ERROR
+                    }
+                    .into(),
+                )
+                .from_env_lossy(),
+        )
+        .init();
+    let (tx, rx) = oneshot::channel();
+    let (tx_final, rx_final) = oneshot::channel();
+
+    let shutdown = Shutdown::new(async move {
+        tokio::select! {
+            _ = graceful::default_signal() => {
+                let _ = tx_final.send(Ok(()));
+            }
+            result = rx => {
+                match result {
+                    Ok(result) => {
+                        let _ = tx_final.send(result);
+                    }
+                    Err(_) => {
+                        let _ = tx_final.send(Ok(()));
+                    }
+                }
+            }
+        }
+    });
+    shutdown.spawn_task_fn(move |guard| async move {
+        let result = run_inner(guard, cfg).await;
+        let _ = tx.send(result);
+    });
+
+    let _ = shutdown.shutdown_with_limit(Duration::from_secs(1)).await;
+
+    rx_final.await?
+}
+async fn run_inner(guard: ShutdownGuard, cfg: CliCommandProbe) -> Result<(), BoxError> {
+    if cfg.clone().list.is_some() {
+        let filename = &cfg.clone().list.unwrap().to_owned();
+        let file = File::open(filename)?;
+        let reader = BufReader::new(file);
+        for line in reader.lines() {
+            let line = line?;
+            let mut request_args_builder = RequestArgsBuilder::new();
+            request_args_builder.parse_arg(line.clone());
+            let request = request_args_builder.build()?;
+
+            let client = create_client(guard.clone(), cfg.clone()).await?;
+
+            let response = client.serve(Context::default(), request).await?;
+            let status = response.status();
+            pretty_print_status_code(status, line);
+            if status.is_client_error() {
+                return Err(ErrorWithExitCode::new(
+                    4,
+                    OpaqueError::from_display(format!("client http error, status: {status}")),
+                )
+                .into());
+            } else if status.is_server_error() {
+                return Err(ErrorWithExitCode::new(
+                    5,
+                    OpaqueError::from_display(format!("server http error, status: {status}")),
+                )
+                .into());
+            }
+        }
+    } else {
+        let domains = cfg.args.clone().unwrap();
+        if domains.len() == 1 && domains[0].contains(".txt") {
+            return Err(ErrorWithExitCode::new(
+                3,
+                OpaqueError::from_display(format!(
+    "Looks like you are trying to pass a txt file directly to 'rama probe'. Instead, use, 'rama probe -l' {}",
+    domains[0]
+)),
+            )
+            .into());
+        }
+        for domain in domains {
+            let mut request_args_builder = RequestArgsBuilder::new();
+            request_args_builder.parse_arg(domain.clone());
+            let request = request_args_builder.build()?;
+
+            let client = create_client(guard.clone(), cfg.clone()).await?;
+
+            let response = client.serve(Context::default(), request).await?;
+            let status = response.status();
+            pretty_print_status_code(status, domain);
+            if status.is_client_error() {
+                return Err(ErrorWithExitCode::new(
+                    4,
+                    OpaqueError::from_display(format!("client http error, status: {status}")),
+                )
+                .into());
+            } else if status.is_server_error() {
+                return Err(ErrorWithExitCode::new(
+                    5,
+                    OpaqueError::from_display(format!("server http error, status: {status}")),
+                )
+                .into());
+            }
+        }
+    }
+
+    Ok(())
+}
+async fn create_client<S>(
+    guard: ShutdownGuard,
+    cfg: CliCommandProbe,
+) -> Result<impl Service<S, Request, Response = Response, Error = BoxError>, BoxError>
+where
+    S: Send + Sync + 'static,
+{
+    // Pass None for both modes so we can just pretty print the status code
+    let (request_writer_mode, response_writer_mode) = (None, None);
+    let writer_kind = writer::WriterKind::Stdout;
+    let executor = Executor::graceful(guard);
+    let (request_writer, response_writer) = writer::create_traffic_writers(
+        &executor,
+        writer_kind,
+        false, // Do not write headers
+        request_writer_mode,
+        response_writer_mode,
+    )
+    .await?;
+
+    let client_builder = (
+        MapResultLayer::new(map_internal_client_error),
+        TimeoutLayer::new(Duration::from_secs(180)),
+        FollowRedirectLayer::with_policy(Limited::new(0)),
+        response_writer,
+        DecompressionLayer::new(),
+        AddRequiredRequestHeadersLayer::default(),
+        request_writer,
+    );
+    Ok(client_builder.layer(HttpClient::default()))
+}
+fn map_internal_client_error<E, Body>(
+    result: Result<Response<Body>, E>,
+) -> Result<Response, BoxError>
+where
+    E: Into<BoxError>,
+    Body: rama::http::dep::http_body::Body<Data = bytes::Bytes> + Send + Sync + 'static,
+    Body::Error: Into<BoxError>,
+{
+    match result {
+        Ok(response) => Ok(response.map(rama::http::Body::new)),
+        Err(err) => Err(err.into()),
+    }
+}

rama-cli/src/cmd/probe/writer.rs

@@ -0,0 +1,42 @@
+use rama::{
+    error::BoxError,
+    http::layer::traffic_writer::{
+        BidirectionalMessage, BidirectionalWriter, RequestWriterLayer, ResponseWriterLayer,
+        WriterMode,
+    },
+    rt::Executor,
+};
+use tokio::{io::stdout, sync::mpsc::Sender};
+
+#[derive(Debug, Clone)]
+pub(super) enum WriterKind {
+    Stdout,
+}
+
+pub(super) async fn create_traffic_writers(
+    executor: &Executor,
+    kind: WriterKind,
+    all: bool,
+    request_mode: Option<WriterMode>,
+    response_mode: Option<WriterMode>,
+) -> Result<
+    (
+        RequestWriterLayer<BidirectionalWriter<Sender<BidirectionalMessage>>>,
+        ResponseWriterLayer<BidirectionalWriter<Sender<BidirectionalMessage>>>,
+    ),
+    BoxError,
+> {
+    let writer = match kind {
+        WriterKind::Stdout => stdout(),
+    };
+    let bidirectional_writer = if all {
+        BidirectionalWriter::new(executor, writer, 32, request_mode, response_mode)
+    } else {
+        BidirectionalWriter::last(executor, writer, request_mode, response_mode)
+    };
+
+    Ok((
+        RequestWriterLayer::new(bidirectional_writer.clone()),
+        ResponseWriterLayer::new(bidirectional_writer),
+    ))
+}

rama-cli/src/main.rs

@@ -45,7 +45,7 @@ use clap::{Parser, Subcommand};
 use rama::error::BoxError;
 
 pub mod cmd;
-use cmd::{echo, fp, http, ip, proxy};
+use cmd::{echo, fp, http, ip, probe, proxy};
 
 pub mod error;
 
@@ -62,6 +62,7 @@ struct Cli {
 #[allow(clippy::large_enum_variant)]
 enum CliCommands {
     Http(http::CliCommandHttp),
+    Probe(probe::CliCommandProbe),
     Proxy(proxy::CliCommandProxy),
     Echo(echo::CliCommandEcho),
     Ip(ip::CliCommandIp),
@@ -75,6 +76,7 @@ async fn main() -> Result<(), BoxError> {
     #[allow(clippy::exit)]
     match match cli.cmds {
         CliCommands::Http(cfg) => http::run(cfg).await,
+        CliCommands::Probe(cfg) => probe::run(cfg).await,
         CliCommands::Proxy(cfg) => proxy::run(cfg).await,
         CliCommands::Echo(cfg) => echo::run(cfg).await,
         CliCommands::Ip(cfg) => ip::run(cfg).await,