linkshieldai 0.2.0

Rust SDK for the LinkShieldAI URL safety API.
Documentation
use std::path::PathBuf;
use std::time::Duration;

use clap::{Parser, Subcommand};
use linkshieldai::LinkShieldAI;

#[derive(Debug, Parser)]
#[command(
    name = "linkshieldai",
    about = "Command line client for the LinkShieldAI API."
)]
struct Cli {
    #[arg(long, help = "LinkShieldAI API key. Defaults to LINKSHIELDAI_API_KEY.")]
    api_key: Option<String>,

    #[arg(long, default_value = linkshieldai::DEFAULT_BASE_URL, help = "API base URL.")]
    base_url: String,

    #[arg(long, default_value_t = 10.0, help = "Request timeout in seconds.")]
    timeout: f64,

    #[arg(
        long,
        default_value_t = 2,
        help = "Retry attempts for transient failures."
    )]
    max_retries: usize,

    #[command(subcommand)]
    command: Command,
}

#[derive(Debug, Subcommand)]
enum Command {
    #[command(about = "Run the versioned URL-risk scan.")]
    Scan {
        url: String,
        #[arg(long, default_value = "standard", value_parser = ["standard", "detailed", "deep"])]
        mode: String,
    },

    #[command(about = "Run a basic URL safety check.")]
    Basic { url: String },

    #[command(about = "Run a detailed URL safety check.")]
    Detailed { url: String },

    #[command(about = "Run an NSFW URL check.")]
    Nsfw { url: String },

    #[command(about = "Run Chimera AI classification.")]
    Chimera { url: String },

    #[command(about = "Download a screenshot by file name or screenshot URL.")]
    Screenshot {
        file_name_or_url: String,
        #[arg(long, short)]
        output: PathBuf,
    },
}

#[tokio::main]
async fn main() {
    if let Err(error) = run().await {
        eprintln!(
            "{}",
            serde_json::to_string_pretty(&serde_json::json!({ "error": error.to_string() }))
                .unwrap()
        );
        std::process::exit(1);
    }
}

async fn run() -> linkshieldai::Result<()> {
    let cli = Cli::parse();
    let client = LinkShieldAI::builder()
        .api_key(cli.api_key.unwrap_or_default())
        .base_url(cli.base_url)
        .timeout(Duration::from_secs_f64(cli.timeout.max(0.0)))
        .max_retries(cli.max_retries)
        .build()?;

    match cli.command {
        Command::Scan { url, mode } => print_json(&client.scan(&url, &mode).await?)?,
        Command::Basic { url } => print_json(&client.basic_check(&url).await?)?,
        Command::Detailed { url } => print_json(&client.detailed_check(&url).await?)?,
        Command::Nsfw { url } => print_json(&client.nsfw_check(&url).await?)?,
        Command::Chimera { url } => print_json(&client.chimera(&url).await?)?,
        Command::Screenshot {
            file_name_or_url,
            output,
        } => {
            client
                .get_screenshot(&file_name_or_url, Some(&output))
                .await?;
            print_json(&serde_json::json!({ "saved": output }))?;
        }
    }

    Ok(())
}

fn print_json(value: &impl serde::Serialize) -> linkshieldai::Result<()> {
    println!(
        "{}",
        serde_json::to_string_pretty(value).map_err(|error| {
            linkshieldai::LinkShieldAIError::Response {
                message: error.to_string(),
                payload: None,
            }
        })?
    );
    Ok(())
}