use std::path::Path;
use std::time::Duration;
use reqwest::{Client as HttpClient, StatusCode};
use serde_json::Value;
use tokio::time::sleep;
use url::Url;
use crate::error::{LinkShieldAIError, Result};
use crate::parse::{
ensure_payload, error_message_from_payload, parse_basic, parse_chimera, parse_detailed,
parse_nsfw,
};
use crate::types::{
BasicCheckResult, ChimeraResult, DetailedCheckResult, NsfwCheckResult, ScanResult,
};
pub const DEFAULT_BASE_URL: &str = "https://api.linkshieldai.com";
#[derive(Debug, Clone)]
pub struct LinkShieldAI {
api_key: String,
base_url: String,
timeout: Duration,
max_retries: usize,
backoff_factor: Duration,
http_client: HttpClient,
}
#[derive(Debug, Clone)]
pub struct LinkShieldAIBuilder {
api_key: Option<String>,
base_url: String,
timeout: Duration,
max_retries: usize,
backoff_factor: Duration,
http_client: Option<HttpClient>,
}
impl LinkShieldAI {
pub fn new(api_key: impl Into<String>) -> Result<Self> {
Self::builder().api_key(api_key).build()
}
pub fn builder() -> LinkShieldAIBuilder {
LinkShieldAIBuilder {
api_key: None,
base_url: DEFAULT_BASE_URL.to_string(),
timeout: Duration::from_secs(10),
max_retries: 2,
backoff_factor: Duration::from_millis(500),
http_client: None,
}
}
pub async fn basic_check(&self, url: &str) -> Result<BasicCheckResult> {
let payload = self.get_json("/", &[("url", url)]).await?;
Ok(parse_basic(payload))
}
pub async fn detailed_check(&self, url: &str) -> Result<DetailedCheckResult> {
let payload = self.get_json("/classify_link", &[("url", url)]).await?;
Ok(parse_detailed(payload))
}
pub async fn nsfw_check(&self, url: &str) -> Result<NsfwCheckResult> {
let payload = self.get_json("/nsfw/site", &[("url", url)]).await?;
Ok(parse_nsfw(payload))
}
pub async fn chimera(&self, url: &str) -> Result<ChimeraResult> {
let payload = self.get_json("/chimera", &[("url", url)]).await?;
Ok(parse_chimera(payload))
}
pub async fn scan(&self, url: &str, mode: &str) -> Result<ScanResult> {
if !matches!(mode, "standard" | "detailed" | "deep") {
return Err(LinkShieldAIError::InvalidInput(
"scan mode must be standard, detailed, or deep".to_string(),
));
}
let endpoint = format!("{}/v1/scan", self.base_url);
for attempt in 0..=self.max_retries {
let response = self
.http_client
.post(&endpoint)
.bearer_auth(&self.api_key)
.json(&serde_json::json!({ "url": url, "mode": mode }))
.timeout(self.timeout)
.send()
.await;
match response {
Ok(response)
if attempt < self.max_retries && is_retryable_status(response.status()) =>
{
let delay = retry_delay(attempt, self.backoff_factor, retry_after(&response));
sleep(delay).await;
}
Ok(response) => {
let status = response.status();
let retry = retry_after(&response);
let payload: Value = response.json().await?;
if status.is_success() {
return serde_json::from_value(payload).map_err(|error| {
LinkShieldAIError::Response {
message: error.to_string(),
payload: None,
}
});
}
if status == StatusCode::TOO_MANY_REQUESTS {
return Err(LinkShieldAIError::RateLimit { retry_after: retry });
}
let object = payload.as_object().cloned();
let message = object
.as_ref()
.map(|value| error_message_from_payload(value, "LinkShieldAI scan failed."))
.unwrap_or_else(|| "LinkShieldAI scan failed.".to_string());
return Err(LinkShieldAIError::Status {
status_code: status.as_u16(),
message,
payload: Some(payload),
});
}
Err(error) if attempt < self.max_retries => {
sleep(retry_delay(attempt, self.backoff_factor, None)).await;
drop(error);
}
Err(error) => return Err(LinkShieldAIError::Connection(error)),
}
}
Err(LinkShieldAIError::InvalidInput(
"request failed after retries".to_string(),
))
}
pub async fn is_malicious(&self, url: &str) -> Result<bool> {
Ok(self.basic_check(url).await?.is_malicious)
}
pub async fn is_nsfw(&self, url: &str) -> Result<bool> {
Ok(self.nsfw_check(url).await?.is_nsfw)
}
pub async fn get_screenshot(
&self,
file_name_or_url: &str,
output_path: Option<&Path>,
) -> Result<Vec<u8>> {
let file_name = screenshot_file_name(file_name_or_url)?;
let response = self
.request("/screenshot/", &[("file", file_name.as_str())], true)
.await?;
let bytes = response.bytes().await?.to_vec();
if let Some(path) = output_path {
tokio::fs::write(path, &bytes).await?;
}
Ok(bytes)
}
async fn get_json(
&self,
path: &str,
query: &[(&str, &str)],
) -> Result<serde_json::Map<String, Value>> {
let response = self.request(path, query, false).await?;
let status = response.status();
let final_retry_after = retry_after(&response);
let content_type = response
.headers()
.get(reqwest::header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok())
.unwrap_or("unknown")
.to_string();
let text = response.text().await?;
let parsed: Value = serde_json::from_str(&text).map_err(|_| {
let preview = text.trim().replace('\n', " ");
let preview = preview.chars().take(300).collect::<String>();
let message = if preview.is_empty() {
"LinkShieldAI API returned invalid JSON.".to_string()
} else {
format!("LinkShieldAI API returned invalid JSON. Content-Type: {content_type}. Body preview: {preview}")
};
LinkShieldAIError::Response { message, payload: None }
})?;
if status.is_success() {
return ensure_payload(parsed);
}
let payload = parsed.as_object().cloned();
if status == StatusCode::TOO_MANY_REQUESTS {
return Err(LinkShieldAIError::RateLimit {
retry_after: final_retry_after,
});
}
let message = payload
.as_ref()
.map(|payload| {
error_message_from_payload(
payload,
&format!("LinkShieldAI API returned HTTP {}.", status.as_u16()),
)
})
.unwrap_or_else(|| format!("LinkShieldAI API returned HTTP {}.", status.as_u16()));
Err(LinkShieldAIError::Status {
status_code: status.as_u16(),
message,
payload: payload.map(Value::Object),
})
}
async fn request(
&self,
path: &str,
query: &[(&str, &str)],
screenshot_path: bool,
) -> Result<reqwest::Response> {
let url = if screenshot_path {
let file_name = query
.iter()
.find(|(key, _)| *key == "file")
.map(|(_, value)| *value)
.unwrap_or_default();
format!("{}{}{}", self.base_url, path, file_name)
} else {
let mut url = Url::parse(&(self.base_url.clone() + path))?;
{
let mut pairs = url.query_pairs_mut();
for (key, value) in query {
pairs.append_pair(key, value);
}
}
url.to_string()
};
let mut last_error: Option<reqwest::Error> = None;
for attempt in 0..=self.max_retries {
let response = self
.http_client
.get(&url)
.bearer_auth(&self.api_key)
.timeout(self.timeout)
.send()
.await;
match response {
Ok(response) => {
if attempt < self.max_retries && is_retryable_status(response.status()) {
let delay =
retry_delay(attempt, self.backoff_factor, retry_after(&response));
sleep(delay).await;
continue;
}
return Ok(response);
}
Err(error) => {
if attempt < self.max_retries {
last_error = Some(error);
sleep(retry_delay(attempt, self.backoff_factor, None)).await;
continue;
}
return Err(LinkShieldAIError::Connection(error));
}
}
}
match last_error {
Some(error) => Err(LinkShieldAIError::Connection(error)),
None => Err(LinkShieldAIError::InvalidInput(
"request failed after retries".to_string(),
)),
}
}
}
impl LinkShieldAIBuilder {
pub fn api_key(mut self, api_key: impl Into<String>) -> Self {
let api_key = api_key.into();
self.api_key = if api_key.is_empty() {
None
} else {
Some(api_key)
};
self
}
pub fn base_url(mut self, base_url: impl Into<String>) -> Self {
self.base_url = base_url.into();
self
}
pub fn timeout(mut self, timeout: Duration) -> Self {
self.timeout = timeout;
self
}
pub fn max_retries(mut self, max_retries: usize) -> Self {
self.max_retries = max_retries;
self
}
pub fn backoff_factor(mut self, backoff_factor: Duration) -> Self {
self.backoff_factor = backoff_factor;
self
}
pub fn http_client(mut self, http_client: HttpClient) -> Self {
self.http_client = Some(http_client);
self
}
pub fn build(self) -> Result<LinkShieldAI> {
let api_key = self
.api_key
.or_else(|| std::env::var("LINKSHIELDAI_API_KEY").ok())
.filter(|value| !value.is_empty())
.ok_or(LinkShieldAIError::Authentication)?;
let base_url = self.base_url.trim_end_matches('/').to_string();
Ok(LinkShieldAI {
api_key,
base_url,
timeout: self.timeout,
max_retries: self.max_retries,
backoff_factor: self.backoff_factor,
http_client: self.http_client.unwrap_or_default(),
})
}
}
fn is_retryable_status(status: StatusCode) -> bool {
matches!(
status,
StatusCode::TOO_MANY_REQUESTS
| StatusCode::BAD_GATEWAY
| StatusCode::SERVICE_UNAVAILABLE
| StatusCode::GATEWAY_TIMEOUT
)
}
fn retry_after(response: &reqwest::Response) -> Option<Duration> {
response
.headers()
.get(reqwest::header::RETRY_AFTER)
.and_then(|value| value.to_str().ok())
.and_then(|value| value.parse::<f64>().ok())
.map(|seconds| Duration::from_secs_f64(seconds.max(0.0)))
}
fn retry_delay(
attempt: usize,
backoff_factor: Duration,
retry_after: Option<Duration>,
) -> Duration {
retry_after.unwrap_or_else(|| backoff_factor.mul_f64(2_f64.powi(attempt as i32)))
}
fn screenshot_file_name(file_name_or_url: &str) -> Result<String> {
if let Ok(url) = Url::parse(file_name_or_url) {
if let Some(segment) = url
.path_segments()
.and_then(|mut segments| segments.next_back())
{
if !segment.is_empty() {
return Ok(segment.to_string());
}
}
}
let candidate = Path::new(file_name_or_url)
.file_name()
.and_then(|value| value.to_str())
.unwrap_or_default();
if candidate.is_empty() {
return Err(LinkShieldAIError::InvalidInput(
"A screenshot file name or screenshot URL is required.".to_string(),
));
}
Ok(candidate.to_string())
}