use crate::error::{Result, WebshotError};
use serde::{Deserialize, Serialize};
use std::path::{Path, PathBuf};
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct Config {
pub screenshots: Vec<ScreenshotConfig>,
#[serde(default)]
pub defaults: DefaultConfig,
}
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct ScreenshotConfig {
pub url: String,
pub output: PathBuf,
#[serde(default = "default_width")]
pub width: u32,
#[serde(default = "default_height")]
pub height: u32,
pub selector: Option<String>,
pub javascript: Option<String>,
pub wait_for: Option<String>,
#[serde(default = "default_timeout")]
pub timeout: u64,
#[serde(default)]
pub retina: bool,
pub quality: Option<u8>,
#[serde(default)]
pub wait: u64,
pub user_agent: Option<String>,
pub format: Option<String>,
#[serde(default)]
pub headers: std::collections::HashMap<String, String>,
#[serde(default)]
pub cookies: Vec<CookieConfig>,
pub auth: Option<AuthConfig>,
}
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct CookieConfig {
pub name: String,
pub value: String,
pub domain: Option<String>,
pub path: Option<String>,
pub secure: Option<bool>,
pub http_only: Option<bool>,
}
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct AuthConfig {
pub username: String,
pub password: String,
}
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct DefaultConfig {
#[serde(default = "default_width")]
pub width: u32,
#[serde(default = "default_height")]
pub height: u32,
#[serde(default = "default_timeout")]
pub timeout: u64,
pub user_agent: Option<String>,
pub output_dir: Option<PathBuf>,
#[serde(default)]
pub wait: u64,
#[serde(default)]
pub retina: bool,
pub quality: Option<u8>,
#[serde(default)]
pub headers: std::collections::HashMap<String, String>,
#[serde(default)]
pub cookies: Vec<CookieConfig>,
}
impl Default for DefaultConfig {
fn default() -> Self {
Self {
width: default_width(),
height: default_height(),
timeout: default_timeout(),
user_agent: None,
output_dir: None,
wait: 0,
retina: false,
quality: None,
headers: std::collections::HashMap::new(),
cookies: Vec::new(),
}
}
}
impl Config {
pub fn from_file<P: AsRef<Path>>(path: P) -> Result<Self> {
let content = std::fs::read_to_string(&path)?;
let mut config: Config = serde_yaml::from_str(&content)?;
for screenshot in &mut config.screenshots {
if screenshot.width == default_width() && config.defaults.width != default_width() {
screenshot.width = config.defaults.width;
}
if screenshot.height == default_height() && config.defaults.height != default_height() {
screenshot.height = config.defaults.height;
}
if screenshot.timeout == default_timeout() && config.defaults.timeout != default_timeout() {
screenshot.timeout = config.defaults.timeout;
}
if screenshot.user_agent.is_none() && config.defaults.user_agent.is_some() {
screenshot.user_agent = config.defaults.user_agent.clone();
}
if screenshot.quality.is_none() && config.defaults.quality.is_some() {
screenshot.quality = config.defaults.quality;
}
for (key, value) in &config.defaults.headers {
screenshot.headers.entry(key.clone()).or_insert_with(|| value.clone());
}
if screenshot.cookies.is_empty() && !config.defaults.cookies.is_empty() {
screenshot.cookies = config.defaults.cookies.clone();
}
if let Some(output_dir) = &config.defaults.output_dir {
if screenshot.output.is_relative() {
screenshot.output = output_dir.join(&screenshot.output);
}
}
}
Ok(config)
}
pub fn to_file<P: AsRef<Path>>(&self, path: P) -> Result<()> {
let content = serde_yaml::to_string(self)?;
std::fs::write(path, content)?;
Ok(())
}
pub fn validate(&self) -> Result<()> {
if self.screenshots.is_empty() {
return Err(WebshotError::config("No screenshots defined in configuration"));
}
for (i, screenshot) in self.screenshots.iter().enumerate() {
url::Url::parse(&screenshot.url)
.map_err(|e| WebshotError::config(format!("Invalid URL in screenshot {}: {}", i, e)))?;
if screenshot.width == 0 || screenshot.height == 0 {
return Err(WebshotError::InvalidViewport {
width: screenshot.width,
height: screenshot.height,
});
}
if let Some(quality) = screenshot.quality {
if !(1..=100).contains(&quality) {
return Err(WebshotError::config(format!(
"JPEG quality must be between 1-100, got: {}",
quality
)));
}
}
if screenshot.timeout == 0 {
return Err(WebshotError::config(format!(
"Timeout must be greater than 0, got: {}",
screenshot.timeout
)));
}
if screenshot.format.is_none() {
let extension = screenshot
.output
.extension()
.and_then(|ext| ext.to_str())
.map(|ext| ext.to_lowercase());
match extension.as_deref() {
Some("png") | Some("jpg") | Some("jpeg") | Some("pdf") => {}
Some(ext) => {
return Err(WebshotError::UnsupportedFormat {
format: ext.to_string(),
});
}
None => {
return Err(WebshotError::config(format!(
"Output file must have a valid extension: {}",
screenshot.output.display()
)));
}
}
}
}
Ok(())
}
}
fn default_width() -> u32 {
1280
}
fn default_height() -> u32 {
800
}
fn default_timeout() -> u64 {
30
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_config_serialization() {
let config = Config {
screenshots: vec![ScreenshotConfig {
url: "https://example.com".to_string(),
output: PathBuf::from("test.png"),
width: 1920,
height: 1080,
selector: Some(".header".to_string()),
javascript: None,
wait_for: None,
timeout: 30,
retina: false,
quality: None,
wait: 0,
user_agent: None,
format: None,
headers: std::collections::HashMap::new(),
cookies: Vec::new(),
auth: None,
}],
defaults: DefaultConfig::default(),
};
let yaml = serde_yaml::to_string(&config).unwrap();
let deserialized: Config = serde_yaml::from_str(&yaml).unwrap();
assert_eq!(config.screenshots.len(), deserialized.screenshots.len());
assert_eq!(config.screenshots[0].url, deserialized.screenshots[0].url);
}
#[test]
fn test_config_validation() {
let mut config = Config {
screenshots: vec![ScreenshotConfig {
url: "https://example.com".to_string(),
output: PathBuf::from("test.png"),
width: 1920,
height: 1080,
selector: None,
javascript: None,
wait_for: None,
timeout: 30,
retina: false,
quality: None,
wait: 0,
user_agent: None,
format: None,
headers: std::collections::HashMap::new(),
cookies: Vec::new(),
auth: None,
}],
defaults: DefaultConfig::default(),
};
assert!(config.validate().is_ok());
config.screenshots[0].url = "not-a-url".to_string();
assert!(config.validate().is_err());
config.screenshots[0].url = "https://example.com".to_string();
config.screenshots[0].width = 0;
assert!(config.validate().is_err());
}
}