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(())
}