use std::path::{Path, PathBuf};
use std::time::Duration;
use base64::Engine;
use chrono::Local;
use reqwest::{Client, Response, StatusCode, Url};
use serde::de::DeserializeOwned;
use serde_json::{Value, json};
use tracing::{debug, instrument};
use crate::account::Account;
use crate::error::{FlarerError, Result};
use crate::tool::Tool;
use crate::types::{CfEnvelope, JsonOptions, Output, SnapshotResult};
pub const DEFAULT_BASE_URL: &str = "https://api.cloudflare.com/client/v4";
#[derive(Debug, Clone)]
pub struct Flarer {
client: Client,
account: Account,
base_url: Url,
output_dir: PathBuf,
}
impl Flarer {
pub fn new(account: Account) -> Result<Self> {
Self::builder().account(account).build()
}
pub fn builder() -> FlarerBuilder {
FlarerBuilder::default()
}
pub fn http(&self) -> &Client {
&self.client
}
pub fn account(&self) -> &Account {
&self.account
}
#[instrument(skip(self), fields(account_id = %self.account.id))]
pub async fn verify(&self) -> Result<()> {
let url = self
.base_url
.join(&format!("accounts/{}/tokens/verify", self.account.id))?;
let resp = self
.client
.get(url)
.bearer_auth(self.account.token())
.send()
.await?;
let status = resp.status();
if status.is_success() {
debug!("token verified");
Ok(())
} else {
let body = resp.text().await.unwrap_or_default();
Err(FlarerError::Auth(format!("HTTP {}: {}", status, body)))
}
}
pub async fn run(&self, tool: Tool, url: &str) -> Result<Output> {
match tool {
Tool::Content | Tool::Markdown | Tool::Scrape | Tool::Links | Tool::Crawl => {
self.text(tool, url).await.map(Output::Json)
}
Tool::Screenshot | Tool::Pdf => self.binary(tool, url).await.map(Output::File),
Tool::Snapshot => self.snapshot(url).await,
Tool::Json => self
.json(url, JsonOptions::default())
.await
.map(Output::Json),
}
}
#[instrument(skip(self))]
pub async fn text(&self, tool: Tool, target_url: &str) -> Result<Value> {
let target = parse_url(target_url)?;
let endpoint = self.endpoint(tool)?;
let body = json!({
"url": target.as_str(),
"gotoOptions": { "waitUntil": "networkidle0" }
});
let resp = self.post(endpoint, &body).await?;
json_or_api_error(resp).await
}
#[instrument(skip(self))]
pub async fn binary(&self, tool: Tool, target_url: &str) -> Result<PathBuf> {
let target = parse_url(target_url)?;
let endpoint = self.endpoint(tool)?;
let body = json!({ "url": target.as_str() });
let resp = self.post(endpoint, &body).await?;
let status = resp.status();
let content_type = resp
.headers()
.get("content-type")
.and_then(|v| v.to_str().ok())
.unwrap_or("")
.to_string();
if !status.is_success() {
let body = resp.text().await.unwrap_or_default();
return Err(FlarerError::Api {
status: status.as_u16(),
body,
});
}
if !(content_type.contains("image")
|| content_type.contains("pdf")
|| content_type.contains("octet-stream"))
{
let body = resp.text().await.unwrap_or_default();
return Err(FlarerError::Unexpected(format!(
"expected binary response, got content-type {}: {}",
content_type, body
)));
}
let path = self.output_path(&target, tool);
let bytes = resp.bytes().await?;
ensure_parent(&path)?;
tokio::fs::write(&path, &bytes).await?;
Ok(path)
}
#[instrument(skip(self, options))]
pub async fn snapshot_with(
&self,
target_url: &str,
options: SnapshotOptions,
) -> Result<Output> {
let target = parse_url(target_url)?;
let endpoint = self.base_url.join(&format!(
"accounts/{}/browser-rendering/snapshot",
self.account.id
))?;
let body = json!({
"url": target.as_str(),
"setJavaScriptEnabled": options.javascript_enabled,
"screenshotOptions": { "fullPage": options.full_page },
"viewport": {
"width": options.viewport_width,
"height": options.viewport_height,
"deviceScaleFactor": options.device_scale_factor,
},
"gotoOptions": {
"waitUntil": options.wait_until,
"timeout": options.timeout_ms,
}
});
let resp = self.post_url(endpoint, &body).await?;
let envelope: CfEnvelope<SnapshotResult> = parse_envelope(resp).await?;
let result = envelope
.result
.ok_or_else(|| FlarerError::Unexpected("snapshot returned no result".into()))?;
let png_bytes = base64::engine::general_purpose::STANDARD.decode(&result.screenshot)?;
let png_path = self.output_path(&target, Tool::Snapshot);
ensure_parent(&png_path)?;
tokio::fs::write(&png_path, &png_bytes).await?;
Ok(Output::Snapshot {
png: png_path,
html: result.content,
})
}
pub async fn snapshot(&self, target_url: &str) -> Result<Output> {
self.snapshot_with(target_url, SnapshotOptions::default())
.await
}
#[instrument(skip(self, options))]
pub async fn json(&self, target_url: &str, options: JsonOptions) -> Result<Value> {
let target = parse_url(target_url)?;
let endpoint = self.endpoint(Tool::Json)?;
let mut body = json!({ "url": target.as_str() });
if let Some(p) = &options.prompt {
body["prompt"] = Value::String(p.clone());
}
if let Some(rf) = &options.response_format {
body["response_format"] = serde_json::to_value(rf)?;
}
let resp = self.post_url(endpoint, &body).await?;
json_or_api_error(resp).await
}
fn endpoint(&self, tool: Tool) -> Result<Url> {
Ok(self.base_url.join(&format!(
"accounts/{}/browser-rendering/{}",
self.account.id,
tool.as_path()
))?)
}
async fn post(&self, endpoint: Url, body: &Value) -> Result<Response> {
self.post_url(endpoint, body).await
}
async fn post_url(&self, endpoint: Url, body: &Value) -> Result<Response> {
debug!(%endpoint, "POST");
let resp = self
.client
.post(endpoint)
.bearer_auth(self.account.token())
.json(body)
.send()
.await?;
Ok(resp)
}
fn output_path(&self, target: &Url, tool: Tool) -> PathBuf {
let domain = target.host_str().unwrap_or("unknown");
let date_time = Local::now().format("%d%m%Y%H%M").to_string();
let ext = match tool {
Tool::Screenshot | Tool::Snapshot => "png",
Tool::Pdf => "pdf",
_ => "txt",
};
let filename = if matches!(tool, Tool::Snapshot) {
format!("{}_{}_snapshot.{}", domain, date_time, ext)
} else {
format!("{}_{}.{}", domain, date_time, ext)
};
self.output_dir.join(filename)
}
}
fn ensure_parent(path: &Path) -> Result<()> {
if let Some(parent) = path.parent() {
if !parent.as_os_str().is_empty() && !parent.exists() {
std::fs::create_dir_all(parent)?;
}
}
Ok(())
}
fn parse_url(s: &str) -> Result<Url> {
Url::parse(s).map_err(FlarerError::from)
}
async fn json_or_api_error(resp: Response) -> Result<Value> {
let status = resp.status();
let text = resp.text().await?;
if !status.is_success() {
return Err(FlarerError::Api {
status: status.as_u16(),
body: text,
});
}
serde_json::from_str(&text).map_err(FlarerError::from)
}
async fn parse_envelope<T: DeserializeOwned>(resp: Response) -> Result<CfEnvelope<T>> {
let status = resp.status();
let text = resp.text().await?;
if !status.is_success() {
return Err(FlarerError::Api {
status: status.as_u16(),
body: text,
});
}
let env: CfEnvelope<T> = serde_json::from_str(&text)?;
if !env.success {
return Err(FlarerError::Api {
status: StatusCode::OK.as_u16(),
body: serde_json::to_string(&env.errors).unwrap_or_default(),
});
}
Ok(env)
}
#[derive(Debug, Clone)]
pub struct SnapshotOptions {
pub javascript_enabled: bool,
pub full_page: bool,
pub viewport_width: u32,
pub viewport_height: u32,
pub device_scale_factor: f32,
pub wait_until: String,
pub timeout_ms: u32,
}
impl Default for SnapshotOptions {
fn default() -> Self {
Self {
javascript_enabled: true,
full_page: true,
viewport_width: 1440,
viewport_height: 900,
device_scale_factor: 1.0,
wait_until: "networkidle0".to_string(),
timeout_ms: 30_000,
}
}
}
#[derive(Debug, Default)]
pub struct FlarerBuilder {
account: Option<Account>,
base_url: Option<Url>,
output_dir: Option<PathBuf>,
timeout: Option<Duration>,
user_agent: Option<String>,
client: Option<Client>,
}
impl FlarerBuilder {
pub fn account(mut self, a: Account) -> Self {
self.account = Some(a);
self
}
pub fn base_url(mut self, url: Url) -> Self {
self.base_url = Some(url);
self
}
pub fn output_dir(mut self, dir: impl Into<PathBuf>) -> Self {
self.output_dir = Some(dir.into());
self
}
pub fn timeout(mut self, d: Duration) -> Self {
self.timeout = Some(d);
self
}
pub fn user_agent(mut self, ua: impl Into<String>) -> Self {
self.user_agent = Some(ua.into());
self
}
pub fn http_client(mut self, c: Client) -> Self {
self.client = Some(c);
self
}
pub fn build(self) -> Result<Flarer> {
let account = self
.account
.ok_or_else(|| FlarerError::Config("Account is required".into()))?;
let base_url = match self.base_url {
Some(u) => u,
None => Url::parse(DEFAULT_BASE_URL).expect("valid default URL"),
};
let base_url = if base_url.path().ends_with('/') {
base_url
} else {
let mut u = base_url.clone();
u.set_path(&format!("{}/", u.path()));
u
};
let output_dir = self.output_dir.unwrap_or_else(default_output_dir);
let client = match self.client {
Some(c) => c,
None => {
let mut b = Client::builder();
if let Some(t) = self.timeout {
b = b.timeout(t);
}
b = b.user_agent(
self.user_agent
.unwrap_or_else(|| format!("flarer/{}", env!("CARGO_PKG_VERSION"))),
);
b.build()?
}
};
Ok(Flarer {
client,
account,
base_url,
output_dir,
})
}
}
fn default_output_dir() -> PathBuf {
std::env::var_os("FLARER_OUTPUT_DIR")
.map(PathBuf::from)
.unwrap_or_else(|| PathBuf::from("outputs"))
}