From a219376e9a4b948d1587553806e3f64daf59950c Mon Sep 17 00:00:00 2001 From: Roman Kvasnytskyi Date: Tue, 24 Sep 2024 15:19:47 +0200 Subject: [PATCH] feature: Basic functionality --- .github/workflows/ci.yml | 35 +++ .github/workflows/release.yml | 79 ++++++ Cargo.toml | 35 +++ README.md | 141 +++++++++- src/main.rs | 383 ++++++++++++++++++++++++++++ tests/integration_tests.rs | 468 ++++++++++++++++++++++++++++++++++ 6 files changed, 1140 insertions(+), 1 deletion(-) create mode 100644 .github/workflows/ci.yml create mode 100644 .github/workflows/release.yml create mode 100644 Cargo.toml create mode 100644 src/main.rs create mode 100644 tests/integration_tests.rs diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..c7a3431 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,35 @@ +name: CI + +on: + push: + branches: + - main + pull_request: + +jobs: + build-and-test: + runs-on: ${{ matrix.os }} + strategy: + matrix: + os: [ubuntu-latest, macos-latest, windows-latest] + + steps: + - uses: actions/checkout@v3 + + - name: Install Rust + uses: actions-rs/toolchain@v1 + with: + toolchain: stable + override: true + + - name: Build + run: cargo build --verbose + + - name: Run Tests + run: cargo test --verbose + + - name: Run Clippy + run: cargo clippy -- -D warnings + + - name: Run Formatter Check + run: cargo fmt -- --check diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml new file mode 100644 index 0000000..f714055 --- /dev/null +++ b/.github/workflows/release.yml @@ -0,0 +1,79 @@ +name: Release + +on: + push: + tags: + - "[0-9]+.[0-9]+.[0-9]+" + +permissions: + contents: write + +jobs: + build-and-upload: + name: Build and upload + runs-on: ${{ matrix.os }} + + strategy: + matrix: + # You can add more, for any target you'd like! + include: + - build: linux + os: ubuntu-latest + target: x86_64-unknown-linux-musl + + - build: macos + os: macos-latest + target: x86_64-apple-darwin + + + - build: windows + os: windows-latest + target: x86_64-pc-windows-msvc + + steps: + - name: Checkout + uses: actions/checkout@v3 + + - name: Get the release version from the tag + shell: bash + run: echo "VERSION=${GITHUB_REF#refs/tags/}" >> $GITHUB_ENV + + - name: Install Rust + # Or @nightly if you want + uses: dtolnay/rust-toolchain@stable + # Arguments to pass in + with: + # Make Rust compile to our target (defined in the matrix) + targets: ${{ matrix.target }} + + - name: Build + uses: actions-rs/cargo@v1 + with: + use-cross: true + command: build + args: --verbose --release --target ${{ matrix.target }} + + - name: Build archive + shell: bash + run: | + # Replace with the name of your binary + binary_name="fff" + dirname="$binary_name-${{ env.VERSION }}-${{ matrix.target }}" + mkdir "$dirname" + if [ "${{ matrix.os }}" = "windows-latest" ]; then + mv "target/${{ matrix.target }}/release/$binary_name.exe" "$dirname" + else + mv "target/${{ matrix.target }}/release/$binary_name" "$dirname" + fi + if [ "${{ matrix.os }}" = "windows-latest" ]; then + 7z a "$dirname.zip" "$dirname" + echo "ASSET=$dirname.zip" >> $GITHUB_ENV + else + tar -czf "$dirname.tar.gz" "$dirname" + echo "ASSET=$dirname.tar.gz" >> $GITHUB_ENV + fi + - name: Release + uses: softprops/action-gh-release@v1 + with: + files: | + ${{ env.ASSET }} \ No newline at end of file diff --git a/Cargo.toml b/Cargo.toml new file mode 100644 index 0000000..ef81fc1 --- /dev/null +++ b/Cargo.toml @@ -0,0 +1,35 @@ +[package] +name = "fff" +version = "1.0.0" +authors = ["Roman Kvasnytskyi "] +edition = "2021" + +[dependencies] +clap = { version = "4.5.18", features = ["derive"] } +futures = "0.3.30" +regex = "1.10.6" +reqwest = { version = "0.12.7", features = ["rustls-tls"] } +tokio = { version = "1.40.0", features = ["full"] } +hex = "0.4.3" +once_cell = "1.19.0" +twoway = "0.2" +colored = "2.1.0" +xxhash-rust = { version = "0.8.12", features = ["xxh3"] } +bytes = "1.7.2" + +[dev-dependencies] +assert_cmd = "2.0.16" +tempfile = "3.12.0" +predicates = "3.1.2" +httpmock = "0.6.8" +tokio = { version = "1.40.0", features = ["full"] } # For async tests + +[[bin]] +name = "fff" +path = "src/main.rs" +bench = true + +[profile.release] +lto = true +codegen-units = 1 +panic = "abort" diff --git a/README.md b/README.md index 04d7149..c4bdc68 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,141 @@ # fff -Rust reimplementation of tomnomnom fff + +fff is a high-performance, asynchronous command-line tool written in Rust for making HTTP requests to URLs provided via standard input (stdin). It's designed to be fast, efficient, and highly configurable, making it ideal for tasks like web scraping, testing, and automation. + + +# Features +- Asynchronous I/O using Tokio for high concurrency. +- Customizable concurrency limits. +- Support for HTTP methods, headers, and request bodies. +- Proxy support. +- Response filtering based on status codes, content, and more. +- Colorized output for easy readability. +- Fast hashing using xxHash for saving responses uniquely. +- Ignore HTML or empty responses. +- Save responses matching specific criteria. + +# Installation + +## Pre-built Binaries +Download the latest release for your platform from the [Releases page](https://github.com/Periecle/fff/releases). + +## Build from Source +To build fff from source, ensure you have Rust and Cargo installed. + + +```shell +# Clone the repository +git clone https://github.com/Periecle/fff.git +cd fff + +# Build the project in release mode and install it in your Path +cargo install --path . +``` + +# Usage + +## Basic Usage +Supply URLs via stdin, one per line: + +```shell +cat urls.txt | fff [OPTIONS] +``` + +## Options + +```shell +Usage: fff [OPTIONS] + +Request URLs provided on stdin fairly frickin' fast + +Options: + -b, --body Request body + -d, --delay Delay between issuing requests (ms) [default: 100] + -H, --header
Add a header to the request (can be specified multiple times) + --ignore-html Don't save HTML files; useful when looking for non-HTML files only + --ignore-empty Don't save empty files + -k, --keep-alive Use HTTP Keep-Alive + -m, --method HTTP method to use (default: GET, or POST if body is specified) [default: GET] + -M, --match Save responses that include in the body + -o, --output Directory to save responses in (will be created) [default: out] + -s, --save-status ... + Save responses with given status code (can be specified multiple times) + -S, --save Save all responses + -x, --proxy Use the provided HTTP proxy + -h, --help Print help information + -V, --version Print version information +``` + +# Examples + +## Basic Request + +Make request to each URL, do not save any responses +```shell +cat urls.txt | fff +``` + +## Custom HTTP Method and Body + +Make request to each URL using POST, with body and specific header. +```shell +echo "http://example.com/api" | fff -m POST -b '{"key":"value"}' -H "Content-Type: application/json" +``` + +## Using a Proxy + +Make request to each URL via specified proxy server. +```shell +cat urls.txt | fff -x http://proxyserver:8080 +``` + +## Saving Responses with Specific Status Codes + +Make request to each URL and save requests with status codes 200 and 300 into default directory "roots" +```shell +cat urls.txt | fff -s 200 -s 301 +``` + +## Ignoring HTML Responses +Can be useful if you want to fetch all non-html requests. +```shell +cat urls.txt | fff --ignore-html +``` + +## Matching Content in Responses +Matches only content that contains specified string. + +```shell +cat urls.txt | fff -M "Welcome to" +``` + +## Setting Concurrency and Delay + +For targets that have some rate-limits, or just sensitive to high amount of requests you can setup delay between requests in milliseconds. +```shell +cat urls.txt | fff -c 50 -d 500 +``` + +# Original Work +This tool was originally written by [tomnomnom in Go](https://github.com/tomnomnom/fff). + +Differences from the Original Go Tool +- Language: fff is written in Rust, leveraging Rust's safety and performance benefits. +- Asynchronous I/O: Utilizes Tokio for efficient asynchronous operations. +- Performance: Optimized for speed with features like ultrafast hashing using xxHash. +- Extensibility: Easier to extend and maintain due to Rust's powerful type system and package ecosystem. +- Enhanced Features: Additional options like ignoring HTML content, matching response bodies, and colorized output. +- Dependency Management: Uses Cargo for dependency management, simplifying the build process. + +# Contributing +Contributions are welcome! Please open an issue or submit a pull request on GitHub. + +# Fork the repository. +Create a new branch with your feature or bug fix. +Commit your changes with clear messages. +Push to your branch and open a pull request. +Ensure that all tests pass and adhere to the existing code style. + +# License +This project is licensed under the MIT License. + diff --git a/src/main.rs b/src/main.rs new file mode 100644 index 0000000..2af1594 --- /dev/null +++ b/src/main.rs @@ -0,0 +1,383 @@ +use bytes::Bytes; +use clap::Parser; +use colored::Colorize; +use futures::stream::FuturesUnordered; +use futures::StreamExt; +use once_cell::sync::Lazy; +use regex::Regex; +use reqwest::header::{HeaderMap, HeaderName, HeaderValue}; +use reqwest::{Client, Method, Proxy, StatusCode, Url, Version}; +use std::io::{self}; +use std::path::PathBuf; +use std::sync::Arc; +use std::time::Duration; +use tokio::fs as tokio_fs; +use tokio::io::{self as tokio_io, AsyncBufReadExt}; +use tokio::sync::Semaphore; +use tokio::time::sleep; +use xxhash_rust::xxh3::Xxh3; // Import bytes::Bytes + +/// Command-line arguments structure using `clap` +#[derive(Parser, Debug, Clone)] +#[command( + name = "fff", + about = "Request URLs provided on stdin fairly frickin' fast", + version = "1.0" +)] +struct Opts { + /// Request body + #[arg(short = 'b', long)] + body: Option, + + /// Delay between issuing requests (ms) + #[arg(short = 'd', long, default_value_t = 100)] + delay: u64, + + /// Add a header to the request (can be specified multiple times) + #[arg(short = 'H', long)] + header: Vec, + + /// Don't save HTML files; useful when looking for non-HTML files only + #[arg(long = "ignore-html")] + ignore_html: bool, + + /// Don't save empty files + #[arg(long = "ignore-empty")] + ignore_empty: bool, + + /// Use HTTP Keep-Alive + #[arg(short = 'k', long = "keep-alive", alias = "keep-alives")] + keep_alive: bool, + + /// HTTP method to use (default: GET, or POST if body is specified) + #[arg(short = 'm', long, default_value = "GET")] + method: String, + + /// Save responses that include in the body + #[arg(short = 'M', long)] + r#match: Option, + + /// Directory to save responses in (will be created) + #[arg(short = 'o', long, default_value = "out")] + output: PathBuf, + + /// Save responses with given status code (can be specified multiple times) + #[arg(short = 's', long = "save-status")] + save_status: Vec, + + /// Save all responses + #[arg(short = 'S', long = "save")] + save: bool, + + /// Use the provided HTTP proxy + #[arg(short = 'x', long = "proxy")] + proxy: Option, +} + +// Define the ResponseData struct to encapsulate response-related data +struct ResponseData { + method: Method, + raw_url: String, + response_body: Bytes, + resp_headers: HeaderMap, + resp_url: Url, + status: StatusCode, + version: Version, +} + +#[tokio::main] +async fn main() { + let opts = Arc::new(Opts::parse()); + let client = match new_client(&opts) { + Ok(c) => Arc::new(c), + Err(e) => { + eprintln!("{}", format!("Failed to create HTTP client: {}", e).red()); + std::process::exit(1); + } + }; + + let semaphore = Arc::new(Semaphore::new(100)); // Limit concurrency to 100 + let mut tasks = FuturesUnordered::new(); + + let stdin = tokio_io::stdin(); + let reader = tokio_io::BufReader::new(stdin); + let mut lines = reader.lines(); + + while let Some(line) = lines.next_line().await.unwrap_or_else(|e| { + eprintln!("{}", format!("Error reading line from stdin: {}", e).red()); + None + }) { + let url = line; + let permit = semaphore.clone().acquire_owned().await.unwrap(); + let client = Arc::clone(&client); + let opts = Arc::clone(&opts); + + tasks.push(tokio::spawn(async move { + if opts.delay > 0 { + sleep(Duration::from_millis(opts.delay)).await; + } + process_url(client, opts, url).await; + drop(permit); + })); + + while tasks.len() >= 100 { + tasks.next().await; + } + } + + while tasks.next().await.is_some() {} +} + +fn new_client(opts: &Opts) -> Result { + let mut builder = Client::builder() + .timeout(Duration::from_secs(10)) + .danger_accept_invalid_certs(true); + + if !opts.keep_alive { + builder = builder.pool_idle_timeout(Duration::from_secs(0)); + } + + if let Some(ref proxy_url) = opts.proxy { + builder = builder.proxy(Proxy::all(proxy_url)?); + } + + builder.build() +} + +async fn process_url(client: Arc, opts: Arc, raw_url: String) { + let mut method = opts.method.clone(); + let request_body = opts.body.clone(); + + if request_body.is_some() && method.eq_ignore_ascii_case("GET") { + method = "POST".to_string(); + } + + let url = match Url::parse(&raw_url) { + Ok(u) => u, + Err(_) => { + eprintln!("{}", format!("Invalid URL: {}", raw_url).red()); + return; + } + }; + + let method = method.parse::().unwrap_or(Method::GET); + + let mut req = client.request(method.clone(), url.clone()); + + // Add headers + if let Some(headers) = parse_headers(&opts.header) { + req = req.headers(headers); + } + + // Add body + if let Some(body) = request_body.clone() { + req = req.body(body); + } + + // Send the request + let resp = match req.send().await { + Ok(r) => r, + Err(e) => { + eprintln!("{}", format!("Request failed for {}: {}", raw_url, e).red()); + return; + } + }; + + // Extract response data + let status = resp.status(); + let version = resp.version(); + let resp_headers = resp.headers().clone(); + let resp_url = resp.url().clone(); + let response_body = match resp.bytes().await { + Ok(b) => b, + Err(e) => { + eprintln!( + "{}", + format!("Failed to read body for {}: {}", raw_url, e).red() + ); + return; + } + }; + + // Create ResponseData instance + let response_data = ResponseData { + method: method.clone(), + raw_url: raw_url.clone(), + response_body, + resp_headers, + resp_url, + status, + version, + }; + + let mut should_save = + opts.save || (!opts.save_status.is_empty() && opts.save_status.contains(&status.as_u16())); + + // Check if response is HTML + if opts.ignore_html && is_html(&response_data.response_body) { + should_save = false; + } + + // Check if response body is empty or whitespace + if opts.ignore_empty + && response_data + .response_body + .iter() + .all(|&b| b.is_ascii_whitespace()) + { + should_save = false; + } + + // Check if response body contains the match string + if let Some(ref m) = opts.r#match { + should_save = twoway::find_bytes(&response_data.response_body, m.as_bytes()).is_some(); + } + + if !should_save { + println!("{} {}", raw_url, colorize_status(status)); + return; + } + + if let Err(e) = save_response(&opts, &response_data).await { + eprintln!( + "{}", + format!("Failed to save response for {}: {}", raw_url, e).red() + ); + } else { + println!( + "{} {}", + raw_url, + format!("Saved ({})", status.as_u16()).green() + ); + } +} + +/// Function to colorize HTTP status codes +fn colorize_status(status: StatusCode) -> colored::ColoredString { + let status_code = status.as_u16(); + let status_str = status.as_str(); + + match status_code { + 200..=299 => status_str.green(), + 300..=399 => status_str.cyan(), + 400..=499 => status_str.yellow(), + 500..=599 => status_str.red(), + _ => status_str.normal(), + } +} + +fn parse_headers(headers: &[String]) -> Option { + let mut header_map = HeaderMap::new(); + for h in headers { + if let Some((name, value)) = h.split_once(':') { + let name = name.trim(); + let value = value.trim(); + if let (Ok(name), Ok(value)) = ( + HeaderName::from_bytes(name.as_bytes()), + HeaderValue::from_str(value), + ) { + header_map.append(name, value); + } + } + } + if header_map.is_empty() { + None + } else { + Some(header_map) + } +} + +fn is_html(body: &[u8]) -> bool { + body.windows(5).any(|w| w.eq_ignore_ascii_case(b" io::Result<()> { + let method = &response_data.method; + let raw_url = &response_data.raw_url; + let response_body = &response_data.response_body; + let resp_headers = &response_data.resp_headers; + let resp_url = &response_data.resp_url; + let status = response_data.status; + let version = response_data.version; + + let normalised_path = normalise_path(resp_url); + + let hash_input = format!( + "{}{}{}{}", + method, + raw_url, + opts.body.clone().unwrap_or_default(), + opts.header.join("") + ); + + // Use xxHash instead of SHA1 + let mut hasher = Xxh3::new(); + hasher.update(hash_input.as_bytes()); + let hash = hasher.digest(); + let hash_hex = format!("{:016x}", hash); + + let host = resp_url.host_str().unwrap_or("unknown"); + let output_dir = opts.output.join(host).join(normalised_path); + + tokio_fs::create_dir_all(&output_dir).await?; + + let body_filename = output_dir.join(format!("{}.body", hash_hex)); + tokio_fs::write(&body_filename, response_body).await?; + + let headers_filename = output_dir.join(format!("{}.headers", hash_hex)); + let mut buf = String::with_capacity(1024); + + // Request line + buf.push_str(&format!("{} {}\n\n", method, raw_url)); + + // Request headers + for h in &opts.header { + buf.push_str(&format!("> {}\n", h)); + } + buf.push('\n'); + + // Request body + if let Some(body) = &opts.body { + buf.push_str(body); + buf.push_str("\n\n"); + } + + // Status line + let version_str = match version { + Version::HTTP_09 => "0.9", + Version::HTTP_10 => "1.0", + Version::HTTP_11 => "1.1", + Version::HTTP_2 => "2", + Version::HTTP_3 => "3", + _ => "unknown", + }; + + buf.push_str(&format!( + "< HTTP/{} {} {}\n", + version_str, + status.as_u16(), + status.canonical_reason().unwrap_or("") + )); + + // Response headers + for (k, v) in resp_headers.iter() { + buf.push_str(&format!("< {}: {}\n", k, v.to_str().unwrap_or(""))); + } + + tokio_fs::write(&headers_filename, buf).await?; + + Ok(()) +} + +static PATH_NORMALISE_RE: Lazy = Lazy::new(|| Regex::new(r"[^a-zA-Z0-9/._-]+").unwrap()); + +fn normalise_path(url: &Url) -> String { + let path = url.path(); + let normalised = PATH_NORMALISE_RE.replace_all(path, "-").to_string(); + let normalised = normalised.trim_start_matches('/').to_string(); + if normalised.is_empty() { + "root".to_string() + } else { + normalised + } +} diff --git a/tests/integration_tests.rs b/tests/integration_tests.rs new file mode 100644 index 0000000..c4ee2ea --- /dev/null +++ b/tests/integration_tests.rs @@ -0,0 +1,468 @@ +// tests/integration_tests.rs + +use assert_cmd::Command; +use httpmock::prelude::*; +use predicates::prelude::*; +use regex::Regex; +use std::fs; +use std::time::Duration; +use tempfile::TempDir; + +// Include the normalise_path function to match the application's logic +fn normalise_path(url: &reqwest::Url) -> String { + let path = url.path(); + let re = Regex::new(r"[^a-zA-Z0-9/._-]+").unwrap(); + let normalised = re.replace_all(path, "-").to_string(); + // Remove leading slashes to ensure the path is relative + let normalised = normalised.trim_start_matches('/').to_string(); + // If the path is empty after trimming, use a default name + if normalised.is_empty() { + "root".to_string() + } else { + normalised + } +} + +#[tokio::test] +async fn test_basic_request() { + // Start a mock server + let server = MockServer::start_async().await; + + // Create a mock response + let body = "Hello, world!"; + + let _mock = server.mock(|when, then| { + when.method(GET).path("/"); + then.status(200) + .header("Content-Type", "text/plain") + .body(body); + }); + + // Use a temporary output directory + let temp_dir = TempDir::new().unwrap(); + + { + // Prepare the command + let mut cmd = Command::cargo_bin("fff").unwrap(); + + // Set the arguments + cmd.arg("-o").arg(temp_dir.path()).arg("-S"); // Save all responses + + // Provide the URL via stdin + cmd.write_stdin(format!("{}\n", server.url("/"))); + + // Run the command and capture output + cmd.assert() + .success() + .stdout(predicate::str::contains("Saved")); + } + + // Verify that the response body file was created + let host = server.address().ip().to_string(); + + // Compute the normalized path + let url = reqwest::Url::parse(&server.url("/")).unwrap(); + let normalised_path = normalise_path(&url); + + let expected_dir = temp_dir.path().join(host).join(normalised_path); + let entries = fs::read_dir(&expected_dir).expect("Expected directory not found"); + let mut found_body = false; + for entry in entries { + let entry = entry.expect("Failed to read directory entry"); + let path = entry.path(); + if path.extension().and_then(|s| s.to_str()) == Some("body") { + let content = fs::read_to_string(&path).expect("Failed to read body file"); + assert_eq!(content, body); + found_body = true; + break; + } + } + assert!(found_body, "Response body file not found"); +} + +#[tokio::test] +async fn test_post_request_with_body() { + // Start a mock server + let server = MockServer::start_async().await; + + // Create a mock response + let body = "Post response"; + + let _mock = server.mock(|when, then| { + when.method(POST).path("/post").body("test data"); + then.status(200).body(body); + }); + + // Use a temporary output directory + let temp_dir = TempDir::new().unwrap(); + + { + // Prepare the command + let mut cmd = Command::cargo_bin("fff").unwrap(); + + // Set the arguments + cmd.arg("-o") + .arg(temp_dir.path()) + .arg("-b") + .arg("test data") + .arg("-S"); + + // Provide the URL via stdin + cmd.write_stdin(format!("{}\n", server.url("/post"))); + + // Run the command and capture output + cmd.assert() + .success() + .stdout(predicate::str::contains("Saved")); + } + + // Verify that the response body file was created + let host = server.address().ip().to_string(); + + // Compute the normalized path + let url = reqwest::Url::parse(&server.url("/post")).unwrap(); + let normalised_path = normalise_path(&url); + + let expected_dir = temp_dir.path().join(host).join(normalised_path); + let entries = fs::read_dir(&expected_dir).expect("Expected directory not found"); + let mut found_body = false; + for entry in entries { + let entry = entry.expect("Failed to read directory entry"); + let path = entry.path(); + if path.extension().and_then(|s| s.to_str()) == Some("body") { + let content = fs::read_to_string(&path).expect("Failed to read body file"); + assert_eq!(content, body); + found_body = true; + break; + } + } + assert!(found_body, "Response body file not found"); +} + +#[tokio::test] +async fn test_match_option() { + // Start a mock server + let server = MockServer::start_async().await; + + // Create a mock response containing "special string" + let body = "This response contains a special string."; + + let _mock = server.mock(|when, then| { + when.method(GET).path("/"); + then.status(200).body(body); + }); + + // Use a temporary output directory + let temp_dir = TempDir::new().unwrap(); + + { + // Prepare the command + let mut cmd = Command::cargo_bin("fff").unwrap(); + + // Set the arguments + cmd.arg("-o") + .arg(temp_dir.path()) + .arg("-M") + .arg("special string"); + + // Provide the URL via stdin + cmd.write_stdin(format!("{}\n", server.url("/"))); + + // Run the command and capture output + cmd.assert() + .success() + .stdout(predicate::str::contains("Saved")); + } + + // Verify that the response body file was created + let host = server.address().ip().to_string(); + + // Compute the normalized path + let url = reqwest::Url::parse(&server.url("/")).unwrap(); + let normalised_path = normalise_path(&url); + + let expected_dir = temp_dir.path().join(host).join(normalised_path); + let entries = fs::read_dir(&expected_dir).expect("Expected directory not found"); + let mut found_body = false; + for entry in entries { + let entry = entry.expect("Failed to read directory entry"); + let path = entry.path(); + if path.extension().and_then(|s| s.to_str()) == Some("body") { + let content = fs::read_to_string(&path).expect("Failed to read body file"); + assert_eq!(content, body); + found_body = true; + break; + } + } + assert!(found_body, "Response body file not found"); +} + +#[tokio::test] +async fn test_save_status() { + // Start a mock server + let server = MockServer::start_async().await; + + let _mock = server.mock(|when, then| { + when.method(GET).path("/"); + then.status(404).body("Not Found"); + }); + + // Use a temporary output directory + let temp_dir = TempDir::new().unwrap(); + + { + // Prepare the command + let mut cmd = Command::cargo_bin("fff").unwrap(); + + // Set the arguments + cmd.arg("-o").arg(temp_dir.path()).arg("-s").arg("404"); // Save responses with status 404 + + // Provide the URL via stdin + cmd.write_stdin(format!("{}\n", server.url("/"))); + + // Run the command and capture output + cmd.assert() + .success() + .stdout(predicate::str::contains("Saved")); + } + + // Verify that the response body file was created + let host = server.address().ip().to_string(); + + // Compute the normalized path + let url = reqwest::Url::parse(&server.url("/")).unwrap(); + let normalised_path = normalise_path(&url); + + let expected_dir = temp_dir.path().join(host).join(normalised_path); + let entries = fs::read_dir(&expected_dir).expect("Expected directory not found"); + let mut found_body = false; + for entry in entries { + let entry = entry.expect("Failed to read directory entry"); + let path = entry.path(); + if path.extension().and_then(|s| s.to_str()) == Some("body") { + let content = fs::read_to_string(&path).expect("Failed to read body file"); + assert_eq!(content, "Not Found"); + found_body = true; + break; + } + } + assert!(found_body, "Response body file not found"); +} + +#[tokio::test] +async fn test_delay_between_requests() { + // Start a mock server + let server = MockServer::start_async().await; + + // Create mock responses + let _mock1 = server.mock(|when, then| { + when.method(GET).path("/1"); + then.status(200); + }); + + let _mock2 = server.mock(|when, then| { + when.method(GET).path("/2"); + then.status(200); + }); + + // Use a temporary output directory + let temp_dir = TempDir::new().unwrap(); + + // Prepare the command + let mut cmd = Command::cargo_bin("fff").unwrap(); + + // Set delay to 500ms + cmd.arg("-d") + .arg("500") + .arg("-S") // Save all responses + .arg("-o") + .arg(temp_dir.path()); // Output directory + + // Provide the URLs via stdin + cmd.write_stdin(format!("{}\n{}\n", server.url("/1"), server.url("/2"))); + + // Record the time before running + let start_time = std::time::Instant::now(); + + // Run the command and capture output + cmd.assert().success(); + + // Check that the time taken is at least 500ms + let elapsed = start_time.elapsed(); + assert!( + elapsed >= Duration::from_millis(500), + "Elapsed time was less than 500ms" + ); +} + +#[tokio::test] +async fn test_custom_headers() { + // Start a mock server + let server = MockServer::start_async().await; + + // Create a mock response + let _mock = server.mock(|when, then| { + when.method(GET) + .path("/") + .header("X-Test-Header", "HeaderValue"); + then.status(200); + }); + + // Prepare the command + let mut cmd = Command::cargo_bin("fff").unwrap(); + + // Set header + cmd.arg("-H").arg("X-Test-Header: HeaderValue"); + + // Provide the URL via stdin + cmd.write_stdin(format!("{}\n", server.url("/"))); + + // Run the command and capture output + cmd.assert().success(); + + // Verify that the mock was called + let hits = _mock.hits(); + assert_eq!( + hits, 1, + "The mock server did not receive the expected header" + ); +} + +#[tokio::test] +async fn test_ignore_html() { + // Start a mock server + let server = MockServer::start_async().await; + + // Create a mock HTML response + let html_body = "Test"; + + let _mock = server.mock(|when, then| { + when.method(GET).path("/"); + then.status(200) + .header("Content-Type", "text/html") + .body(html_body); + }); + + // Use a temporary output directory + let temp_dir = TempDir::new().unwrap(); + + { + // Prepare the command + let mut cmd = Command::cargo_bin("fff").unwrap(); + + // Set arguments + cmd.arg("-o") + .arg(temp_dir.path()) + .arg("--ignore-html") + .arg("-S"); // Save all responses + + // Provide the URL via stdin + cmd.write_stdin(format!("{}\n", server.url("/"))); + + // Run the command and capture output + cmd.assert() + .success() + .stdout(predicate::str::contains("200")); + } + + // Verify that the response body file was not created + let host = server.address().ip().to_string(); + + // Compute the normalized path + let url = reqwest::Url::parse(&server.url("/")).unwrap(); + let normalised_path = normalise_path(&url); + + let expected_dir = temp_dir.path().join(host).join(normalised_path); + let entries = fs::read_dir(&expected_dir); + assert!( + entries.is_err() || entries.unwrap().next().is_none(), + "Response body file should not be saved" + ); +} + +#[tokio::test] +async fn test_ignore_empty() { + // Start a mock server + let server = MockServer::start_async().await; + + let _mock = server.mock(|when, then| { + when.method(GET).path("/"); + then.status(200).body(""); + }); + + // Use a temporary output directory + let temp_dir = TempDir::new().unwrap(); + + { + // Prepare the command + let mut cmd = Command::cargo_bin("fff").unwrap(); + + // Set arguments + cmd.arg("-o") + .arg(temp_dir.path()) + .arg("--ignore-empty") + .arg("-S"); // Save all responses + + // Provide the URL via stdin + cmd.write_stdin(format!("{}\n", server.url("/"))); + + // Run the command and capture output + cmd.assert() + .success() + .stdout(predicate::str::contains("200")); + } + + // Verify that the response body file was not created + let host = server.address().ip().to_string(); + + // Compute the normalized path + let url = reqwest::Url::parse(&server.url("/")).unwrap(); + let normalised_path = normalise_path(&url); + + let expected_dir = temp_dir.path().join(host).join(normalised_path); + let entries = fs::read_dir(&expected_dir); + assert!( + entries.is_err() || entries.unwrap().next().is_none(), + "Response body file should not be saved" + ); +} + +#[tokio::test] +async fn test_proxy_option() { + // Start a mock server to act as the proxy + let proxy_server = MockServer::start_async().await; + + // Simulate a proxy by allowing any method and path + let _proxy_mock = proxy_server.mock(|when, then| { + when.any_request(); + then.status(200); + }); + + // Start another mock server to handle the actual request + let server = MockServer::start_async().await; + + // Create a mock response + let _mock = server.mock(|when, then| { + when.method(GET).path("/"); + then.status(200).body("Proxied response"); + }); + + // Prepare the command + let mut cmd = Command::cargo_bin("fff").unwrap(); + + // Set the proxy + cmd.arg("-x") + .arg(format!("http://{}", proxy_server.address())); + + // Provide the URL via stdin + cmd.write_stdin(format!("{}\n", server.url("/"))); + + // Run the command and capture output + cmd.assert() + .success() + .stdout(predicate::str::contains("200")); + + // Verify that the proxy server received the request + let hits = _proxy_mock.hits(); + assert!(hits > 0, "Proxy server was not used"); +}