mirror of
https://github.com/thegeneralist01/archivr
synced 2026-07-21 18:55:36 +02:00
feat: add user-configurable cookie rules (#20)
Adds per-instance cookie rules (admin-only) that are injected into every network touchpoint during capture. Storage: - New cookie_rules table in the auth DB (idempotent migration) - Rules have pattern_kind (global/wildcard/regex), optional url_pattern, and cookies_json (validated as string-only JSON object) Matching (resolve_cookies_for_url): - Global rules always apply - Wildcard: * and ? with full metacharacter escaping; matched against hostname via reqwest::Url when pattern has no ://, full URL otherwise - Regex: matched against the full URL - Later rules in ordinal order override earlier ones per cookie name All six network touchpoints receive resolved cookies: - http::probe_url_kind and http::download: Cookie request header - singlefile::save: Netscape cookie file -> --browser-cookies-file - ytdlp::fetch_metadata and ytdlp::download: Netscape cookie file -> --cookies - tweets::archive: semicolon credentials file -> --credentials-file (only when both ct0 and auth_token are present; otherwise falls back to ARCHIVR_TWITTER_CREDENTIALS_FILE) Security: - Cookie files written 0o600 (owner read/write only) - Exact parsed hostname used as cookie domain (no PSL stripping) - Files deleted unconditionally before any error propagates, including spawn failures (hold-result-then-delete pattern) - No cookie values in process args (no --add-header exposure) API: GET/POST /api/admin/cookie-rules, PATCH/DELETE /api/admin/cookie-rules/:uid Frontend: Cookies tab in Settings (admin only) with rule list, inline edit, pattern-type selector, client-side JSON validation CLI: CaptureConfig::default() - no behaviour change 254 tests passing (4 new cookie-rule handler tests)
This commit is contained in:
GitHub
parent
21b11c211f
commit
dae61e585d
17 changed files with 1013 additions and 103 deletions
|
|
@ -3,7 +3,7 @@ use chrono::Local;
|
|||
use uuid::Uuid;
|
||||
use serde_json::json;
|
||||
use std::{
|
||||
collections::HashSet,
|
||||
collections::{HashMap, HashSet},
|
||||
fs,
|
||||
path::{Path, PathBuf},
|
||||
};
|
||||
|
|
@ -69,6 +69,78 @@ impl PlatformMetadata {
|
|||
}
|
||||
}
|
||||
|
||||
/// Configuration passed to `perform_capture` to supply per-instance settings
|
||||
/// that live outside the archive (e.g. cookies stored in the auth DB).
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct CaptureConfig {
|
||||
pub cookie_rules: Vec<database::CookieRule>,
|
||||
}
|
||||
|
||||
/// Resolves which cookies apply to `url` by evaluating all rules in ordinal order.
|
||||
///
|
||||
/// Global rules (`url_pattern = None`) always apply.
|
||||
/// URL-specific rules apply when their pattern matches:
|
||||
/// - `wildcard`: `*` and `?` glob; matched against the URL hostname when the
|
||||
/// pattern does not contain `://`, or the full URL when it does.
|
||||
/// - `regex`: full regex matched against the full URL.
|
||||
///
|
||||
/// Later rules in ordinal order override earlier ones for the same cookie name.
|
||||
pub fn resolve_cookies_for_url(
|
||||
rules: &[database::CookieRule],
|
||||
url: &str,
|
||||
) -> HashMap<String, String> {
|
||||
let mut result = HashMap::new();
|
||||
for rule in rules {
|
||||
let applies = match rule.url_pattern.as_deref() {
|
||||
None => true,
|
||||
Some(pattern) => match rule.pattern_kind.as_str() {
|
||||
"wildcard" => wildcard_matches(pattern, url),
|
||||
"regex" => regex::Regex::new(pattern)
|
||||
.is_ok_and(|re| re.is_match(url)),
|
||||
_ => false,
|
||||
},
|
||||
};
|
||||
if applies {
|
||||
if let Ok(map) =
|
||||
serde_json::from_str::<HashMap<String, String>>(&rule.cookies_json)
|
||||
{
|
||||
result.extend(map);
|
||||
}
|
||||
}
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
/// Converts a wildcard pattern (`*` = any substring, `?` = any character) to a
|
||||
/// regex and tests it against the URL.
|
||||
///
|
||||
/// If the pattern contains `://` it is matched against the full URL; otherwise
|
||||
/// only against the hostname (via `reqwest::Url::parse`) so that patterns like
|
||||
/// `*.youtube.com` work naturally without matching URL path components.
|
||||
fn wildcard_matches(pattern: &str, url: &str) -> bool {
|
||||
use reqwest::Url as ReqwestUrl;
|
||||
let target: String = if pattern.contains("://") {
|
||||
url.to_string()
|
||||
} else {
|
||||
ReqwestUrl::parse(url)
|
||||
.ok()
|
||||
.and_then(|u| u.host_str().map(str::to_string))
|
||||
.unwrap_or_else(|| url.to_string())
|
||||
};
|
||||
let mut pat = String::with_capacity(pattern.len() * 2 + 2);
|
||||
pat.push('^');
|
||||
for ch in pattern.chars() {
|
||||
match ch {
|
||||
'*' => pat.push_str(".*"),
|
||||
'?' => pat.push('.'),
|
||||
c if "$.+[]{}()|^\\".contains(c) => { pat.push('\\'); pat.push(c); }
|
||||
c => pat.push(c),
|
||||
}
|
||||
}
|
||||
pat.push('$');
|
||||
regex::Regex::new(&pat).is_ok_and(|re| re.is_match(&target))
|
||||
}
|
||||
|
||||
fn generate_entry_title(source: Source, meta: &PlatformMetadata) -> String {
|
||||
match source {
|
||||
Source::YouTubeVideo => meta.title.clone().unwrap_or_else(|| "YouTube Video".to_string()),
|
||||
|
|
@ -786,6 +858,7 @@ pub fn perform_capture(
|
|||
locator: &str,
|
||||
archive_id: Option<&str>,
|
||||
quality: Option<&str>,
|
||||
config: &CaptureConfig,
|
||||
) -> Result<CaptureResult> {
|
||||
// Append a UUID so parallel captures starting in the same millisecond
|
||||
// never collide on the staging directory or file names.
|
||||
|
|
@ -801,6 +874,10 @@ pub fn perform_capture(
|
|||
|
||||
let mut source = determine_source(locator);
|
||||
|
||||
// Expand shorthands to the canonical URL for cookie matching.
|
||||
let canonical_url = expand_shorthand_to_url(locator, &source);
|
||||
let cookies = resolve_cookies_for_url(&config.cookie_rules, &canonical_url);
|
||||
|
||||
// Create the run record before probing so every attempt — including
|
||||
// probe failures — is visible in /runs with a proper status and error.
|
||||
let run = database::create_archive_run(&conn, user_id, 1)?;
|
||||
|
|
@ -808,7 +885,7 @@ pub fn perform_capture(
|
|||
// For generic http/https URLs, probe Content-Type to decide whether to
|
||||
// treat the URL as a raw file download or an HTML page for SingleFile.
|
||||
if source == Source::Url {
|
||||
match downloader::http::probe_url_kind(locator) {
|
||||
match downloader::http::probe_url_kind(locator, &cookies) {
|
||||
Ok(downloader::http::UrlKind::Html) => source = Source::WebPage,
|
||||
Ok(downloader::http::UrlKind::File) => {}
|
||||
Err(e) => {
|
||||
|
|
@ -870,7 +947,7 @@ pub fn perform_capture(
|
|||
|
||||
// Source: generic HTTP/S file URL
|
||||
if source == Source::Url {
|
||||
match downloader::http::download(locator, store_path, ×tamp) {
|
||||
match downloader::http::download(locator, store_path, ×tamp, &cookies) {
|
||||
Ok((hash, file_extension, title_hint)) => {
|
||||
let temp_file = store_path
|
||||
.join("temp")
|
||||
|
|
@ -929,7 +1006,7 @@ pub fn perform_capture(
|
|||
|
||||
// Source: web page — archive as a self-contained HTML snapshot via single-file-cli
|
||||
if source == Source::WebPage {
|
||||
match downloader::singlefile::save(locator, store_path, ×tamp) {
|
||||
match downloader::singlefile::save(locator, store_path, ×tamp, &cookies) {
|
||||
Ok(result) => {
|
||||
let file_extension = ".html".to_string();
|
||||
let temp_html = store_path
|
||||
|
|
@ -1083,11 +1160,16 @@ pub fn perform_capture(
|
|||
}
|
||||
};
|
||||
|
||||
// Tweet shorthands (tweet:ID) don't expand to a URL, so `canonical_url`
|
||||
// won't match x.com patterns. Always resolve cookies against x.com so
|
||||
// that wildcard/global rules containing `ct0`+`auth_token` are picked up.
|
||||
let tweet_cookies = resolve_cookies_for_url(&config.cookie_rules, "https://x.com/");
|
||||
match downloader::tweets::archive(
|
||||
locator,
|
||||
source == Source::TweetThread,
|
||||
store_path,
|
||||
×tamp,
|
||||
&tweet_cookies,
|
||||
) {
|
||||
Ok(_) => {
|
||||
let tweet_entry = record_tweet_entry(
|
||||
|
|
@ -1131,7 +1213,7 @@ pub fn perform_capture(
|
|||
| Source::Facebook
|
||||
| Source::TikTok
|
||||
| Source::Reddit
|
||||
| Source::Snapchat => downloader::ytdlp::fetch_metadata(&path),
|
||||
| Source::Snapchat => downloader::ytdlp::fetch_metadata(&path, &cookies),
|
||||
_ => None,
|
||||
};
|
||||
|
||||
|
|
@ -1154,7 +1236,7 @@ pub fn perform_capture(
|
|||
| Source::TikTok
|
||||
| Source::Reddit
|
||||
| Source::Snapchat => {
|
||||
match downloader::ytdlp::download(path.clone(), store_path, ×tamp, quality) {
|
||||
match downloader::ytdlp::download(path.clone(), store_path, ×tamp, quality, &cookies) {
|
||||
Ok(result) => result,
|
||||
Err(e) => {
|
||||
return Err(fail_run(
|
||||
|
|
@ -1168,7 +1250,7 @@ pub fn perform_capture(
|
|||
}
|
||||
Source::YouTubeMusicTrack => {
|
||||
// Music tracks are always audio-only regardless of the caller's quality hint.
|
||||
match downloader::ytdlp::download(path.clone(), store_path, ×tamp, Some("audio")) {
|
||||
match downloader::ytdlp::download(path.clone(), store_path, ×tamp, Some("audio"), &cookies) {
|
||||
Ok(result) => result,
|
||||
Err(e) => {
|
||||
return Err(fail_run(
|
||||
|
|
|
|||
|
|
@ -138,6 +138,16 @@ pub struct InstanceSettings {
|
|||
pub default_entry_visibility: u32,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
||||
pub struct CookieRule {
|
||||
pub rule_uid: String,
|
||||
pub url_pattern: Option<String>,
|
||||
pub pattern_kind: String,
|
||||
pub cookies_json: String,
|
||||
pub ordinal: i64,
|
||||
pub created_at: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, serde::Serialize)]
|
||||
pub struct CollectionRecord {
|
||||
pub id: i64,
|
||||
|
|
@ -464,6 +474,16 @@ pub fn initialize_auth_schema(conn: &Connection) -> Result<()> {
|
|||
created_at TEXT NOT NULL,
|
||||
last_login_at TEXT
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS cookie_rules (
|
||||
id INTEGER PRIMARY KEY,
|
||||
rule_uid TEXT NOT NULL UNIQUE,
|
||||
url_pattern TEXT,
|
||||
pattern_kind TEXT NOT NULL DEFAULT 'global',
|
||||
cookies_json TEXT NOT NULL DEFAULT '{}',
|
||||
ordinal INTEGER NOT NULL DEFAULT 0,
|
||||
created_at TEXT NOT NULL
|
||||
);
|
||||
"#,
|
||||
)?;
|
||||
// Add display_name column to users if not present (idempotent migration)
|
||||
|
|
@ -731,6 +751,85 @@ pub fn update_instance_settings(conn: &Connection, settings: &InstanceSettings)
|
|||
Ok(())
|
||||
}
|
||||
|
||||
pub fn list_cookie_rules(conn: &Connection) -> Result<Vec<CookieRule>> {
|
||||
let mut stmt = conn.prepare(
|
||||
"SELECT rule_uid, url_pattern, pattern_kind, cookies_json, ordinal, created_at
|
||||
FROM cookie_rules ORDER BY ordinal ASC, created_at ASC",
|
||||
)?;
|
||||
let records = stmt
|
||||
.query_map([], |row| {
|
||||
Ok(CookieRule {
|
||||
rule_uid: row.get(0)?,
|
||||
url_pattern: row.get(1)?,
|
||||
pattern_kind: row.get(2)?,
|
||||
cookies_json: row.get(3)?,
|
||||
ordinal: row.get(4)?,
|
||||
created_at: row.get(5)?,
|
||||
})
|
||||
})?
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
Ok(records)
|
||||
}
|
||||
|
||||
pub fn create_cookie_rule(
|
||||
conn: &Connection,
|
||||
url_pattern: Option<&str>,
|
||||
pattern_kind: &str,
|
||||
cookies_json: &str,
|
||||
) -> Result<CookieRule> {
|
||||
let rule_uid = format!("cr_{}", Uuid::new_v4().simple());
|
||||
let now = Utc::now().to_rfc3339();
|
||||
let ordinal: i64 = conn.query_row(
|
||||
"SELECT COALESCE(MAX(ordinal), -1) + 1 FROM cookie_rules",
|
||||
[],
|
||||
|r| r.get(0),
|
||||
)?;
|
||||
conn.execute(
|
||||
"INSERT INTO cookie_rules (rule_uid, url_pattern, pattern_kind, cookies_json, ordinal, created_at)
|
||||
VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
|
||||
params![rule_uid, url_pattern, pattern_kind, cookies_json, ordinal, now],
|
||||
)?;
|
||||
Ok(CookieRule {
|
||||
rule_uid,
|
||||
url_pattern: url_pattern.map(str::to_string),
|
||||
pattern_kind: pattern_kind.to_string(),
|
||||
cookies_json: cookies_json.to_string(),
|
||||
ordinal,
|
||||
created_at: now,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn update_cookie_rule(
|
||||
conn: &Connection,
|
||||
rule_uid: &str,
|
||||
url_pattern: Option<&str>,
|
||||
pattern_kind: &str,
|
||||
cookies_json: &str,
|
||||
ordinal: i64,
|
||||
) -> Result<()> {
|
||||
let rows = conn.execute(
|
||||
"UPDATE cookie_rules
|
||||
SET url_pattern = ?1, pattern_kind = ?2, cookies_json = ?3, ordinal = ?4
|
||||
WHERE rule_uid = ?5",
|
||||
params![url_pattern, pattern_kind, cookies_json, ordinal, rule_uid],
|
||||
)?;
|
||||
if rows == 0 {
|
||||
anyhow::bail!("cookie rule not found: {rule_uid}");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn delete_cookie_rule(conn: &Connection, rule_uid: &str) -> Result<()> {
|
||||
let rows = conn.execute(
|
||||
"DELETE FROM cookie_rules WHERE rule_uid = ?1",
|
||||
[rule_uid],
|
||||
)?;
|
||||
if rows == 0 {
|
||||
anyhow::bail!("cookie rule not found: {rule_uid}");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn get_user_password_hash(conn: &Connection, user_id: i64) -> Result<Option<String>> {
|
||||
conn.query_row(
|
||||
"SELECT password_hash FROM users WHERE id = ?1",
|
||||
|
|
|
|||
66
crates/archivr-core/src/downloader/cookies.rs
Normal file
66
crates/archivr-core/src/downloader/cookies.rs
Normal file
|
|
@ -0,0 +1,66 @@
|
|||
use std::{collections::HashMap, io::Write, path::Path};
|
||||
|
||||
/// Writes a Netscape-format HTTP cookie file suitable for yt-dlp `--cookies` and
|
||||
/// single-file `--browser-cookies-file`. The file is created with mode 0o600
|
||||
/// (owner read/write only) to protect cookie secrets.
|
||||
///
|
||||
/// `domain` is the exact hostname of the target URL (as returned by
|
||||
/// `domain_from_url`). No leading dot is added — cookies are scoped to the
|
||||
/// exact host only, not all subdomains. This is intentionally conservative:
|
||||
/// `www.youtube.com` cookies will not leak to `music.youtube.com`.
|
||||
pub fn write_netscape_cookie_file(
|
||||
cookies: &HashMap<String, String>,
|
||||
domain: &str,
|
||||
path: &Path,
|
||||
) -> std::io::Result<()> {
|
||||
#[cfg(unix)]
|
||||
let mut f = {
|
||||
use std::os::unix::fs::OpenOptionsExt;
|
||||
std::fs::OpenOptions::new()
|
||||
.write(true)
|
||||
.create(true)
|
||||
.truncate(true)
|
||||
.mode(0o600)
|
||||
.open(path)?
|
||||
};
|
||||
#[cfg(not(unix))]
|
||||
let mut f = std::fs::File::create(path)?;
|
||||
|
||||
writeln!(f, "# Netscape HTTP Cookie File")?;
|
||||
// Exact hostname — no leading dot, no subdomain wildcard.
|
||||
// The user controls scope via their URL patterns.
|
||||
let cookie_domain = domain.to_string();
|
||||
for (name, value) in cookies {
|
||||
// Fields: domain include_subdomains path secure expiry name value
|
||||
// include_subdomains is FALSE because there is no leading dot.
|
||||
writeln!(f, "{cookie_domain}\tFALSE\t/\tFALSE\t0\t{name}\t{value}")?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Extracts the exact hostname from a URL, suitable for use as the Netscape
|
||||
/// cookie domain.
|
||||
///
|
||||
/// Uses `reqwest::Url::parse` so IPv6 addresses, auth-in-URL, ports, and
|
||||
/// non-http schemes are handled correctly. Returns an empty string on parse
|
||||
/// failure or when the URL has no host component.
|
||||
///
|
||||
/// Examples:
|
||||
/// - `"https://www.youtube.com/watch?v=x"` → `"www.youtube.com"`
|
||||
/// - `"https://music.youtube.com/"` → `"music.youtube.com"`
|
||||
/// - `"https://twitter.com/"` → `"twitter.com"`
|
||||
pub fn domain_from_url(url: &str) -> String {
|
||||
reqwest::Url::parse(url)
|
||||
.ok()
|
||||
.and_then(|u| u.host_str().map(str::to_string))
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
/// Formats cookies as a single `Cookie:` header value ("name=value; name2=value2").
|
||||
pub fn cookies_to_header(cookies: &HashMap<String, String>) -> String {
|
||||
cookies
|
||||
.iter()
|
||||
.map(|(k, v)| format!("{k}={v}"))
|
||||
.collect::<Vec<_>>()
|
||||
.join("; ")
|
||||
}
|
||||
|
|
@ -1,6 +1,7 @@
|
|||
use anyhow::{Context, Result, bail};
|
||||
use std::path::Path;
|
||||
use std::{collections::HashMap, path::Path};
|
||||
|
||||
use crate::downloader::cookies::cookies_to_header;
|
||||
use crate::hash::hash_file;
|
||||
|
||||
/// Whether a URL resolves to an HTML document or a downloadable file.
|
||||
|
|
@ -16,25 +17,33 @@ pub enum UrlKind {
|
|||
/// Returns `Err` if the probe fails (network error, non-2xx/405 status).
|
||||
/// Redirects (3xx) are followed automatically by reqwest; only the final
|
||||
/// response status is checked.
|
||||
pub fn probe_url_kind(url: &str) -> Result<UrlKind> {
|
||||
pub fn probe_url_kind(url: &str, cookies: &HashMap<String, String>) -> Result<UrlKind> {
|
||||
let client = reqwest::blocking::Client::builder()
|
||||
.redirect(reqwest::redirect::Policy::limited(10))
|
||||
.user_agent("archivr/0.1")
|
||||
.build()
|
||||
.context("failed to build HTTP client")?;
|
||||
|
||||
let cookie_header = cookies_to_header(cookies);
|
||||
|
||||
// Prefer HEAD: no body transfer.
|
||||
let head = client
|
||||
.head(url)
|
||||
.send()
|
||||
.with_context(|| format!("failed to probe {url}"))?;
|
||||
let head = {
|
||||
let mut req = client.head(url);
|
||||
if !cookie_header.is_empty() {
|
||||
req = req.header(reqwest::header::COOKIE, &cookie_header);
|
||||
}
|
||||
req.send().with_context(|| format!("failed to probe {url}"))?
|
||||
};
|
||||
|
||||
if head.status() == reqwest::StatusCode::METHOD_NOT_ALLOWED {
|
||||
// Server rejected HEAD — do a GET but only inspect headers.
|
||||
let get = client
|
||||
.get(url)
|
||||
.send()
|
||||
.with_context(|| format!("failed to probe {url}"))?;
|
||||
let get = {
|
||||
let mut req = client.get(url);
|
||||
if !cookie_header.is_empty() {
|
||||
req = req.header(reqwest::header::COOKIE, &cookie_header);
|
||||
}
|
||||
req.send().with_context(|| format!("failed to probe {url}"))?
|
||||
};
|
||||
if !get.status().is_success() {
|
||||
bail!("HTTP {} probing {url}", get.status());
|
||||
}
|
||||
|
|
@ -74,17 +83,21 @@ pub fn probe_url_kind(url: &str) -> Result<UrlKind> {
|
|||
/// - The request fails or returns a non-2xx status.
|
||||
/// - The response Content-Type is `text/html` (caller should use a web-page archiver instead).
|
||||
/// - The body cannot be written to disk.
|
||||
pub fn download(url: &str, store_path: &Path, timestamp: &str) -> Result<(String, String, Option<String>)> {
|
||||
pub fn download(url: &str, store_path: &Path, timestamp: &str, cookies: &HashMap<String, String>) -> Result<(String, String, Option<String>)> {
|
||||
let client = reqwest::blocking::Client::builder()
|
||||
.redirect(reqwest::redirect::Policy::limited(10))
|
||||
.user_agent("archivr/0.1")
|
||||
.build()
|
||||
.context("failed to build HTTP client")?;
|
||||
|
||||
let response = client
|
||||
.get(url)
|
||||
.send()
|
||||
.with_context(|| format!("failed to fetch {url}"))?;
|
||||
let cookie_header = cookies_to_header(cookies);
|
||||
let response = {
|
||||
let mut req = client.get(url);
|
||||
if !cookie_header.is_empty() {
|
||||
req = req.header(reqwest::header::COOKIE, &cookie_header);
|
||||
}
|
||||
req.send().with_context(|| format!("failed to fetch {url}"))?
|
||||
};
|
||||
|
||||
if !response.status().is_success() {
|
||||
bail!("HTTP {} fetching {url}", response.status());
|
||||
|
|
@ -393,7 +406,7 @@ mod tests {
|
|||
#[test]
|
||||
fn probe_url_kind_fails_on_unreachable_host() {
|
||||
// 127.0.0.1:1 is guaranteed to refuse connections.
|
||||
let err = probe_url_kind("http://127.0.0.1:1/").unwrap_err();
|
||||
let err = probe_url_kind("http://127.0.0.1:1/", &HashMap::new()).unwrap_err();
|
||||
let msg = format!("{err:#}");
|
||||
assert!(
|
||||
msg.contains("probe") || msg.contains("connect") || msg.contains("refused"),
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
pub mod cookies;
|
||||
pub mod local;
|
||||
pub mod store;
|
||||
pub mod tweets;
|
||||
|
|
|
|||
|
|
@ -1,8 +1,9 @@
|
|||
use anyhow::{Context, Result, bail};
|
||||
use base64::engine::general_purpose::STANDARD as B64;
|
||||
use base64::Engine as _;
|
||||
use std::{env, io::Read, path::Path, process::Command};
|
||||
use std::{collections::HashMap, env, io::Read, path::Path, process::Command};
|
||||
|
||||
use crate::downloader::cookies::{domain_from_url, write_netscape_cookie_file};
|
||||
use crate::hash::hash_file;
|
||||
|
||||
/// Result of archiving a web page with single-file.
|
||||
|
|
@ -27,11 +28,11 @@ pub struct SaveResult {
|
|||
/// Reads two env vars:
|
||||
/// - `ARCHIVR_SINGLE_FILE`: path to the `single-file` binary (default: `"single-file"`).
|
||||
/// - `ARCHIVR_CHROME`: path to the Chromium/Chrome binary (default: `"chromium"`).
|
||||
pub fn save(url: &str, store_path: &Path, timestamp: &str) -> Result<SaveResult> {
|
||||
pub fn save(url: &str, store_path: &Path, timestamp: &str, cookies: &HashMap<String, String>) -> Result<SaveResult> {
|
||||
let single_file =
|
||||
env::var("ARCHIVR_SINGLE_FILE").unwrap_or_else(|_| "single-file".to_string());
|
||||
let chrome = env::var("ARCHIVR_CHROME").unwrap_or_else(|_| "chromium".to_string());
|
||||
save_with(url, store_path, timestamp, &single_file, &chrome)
|
||||
save_with(url, store_path, timestamp, &single_file, &chrome, cookies)
|
||||
}
|
||||
|
||||
/// Inner implementation; takes binary paths explicitly so tests can inject them
|
||||
|
|
@ -42,6 +43,7 @@ fn save_with(
|
|||
timestamp: &str,
|
||||
single_file: &str,
|
||||
chrome: &str,
|
||||
cookies: &HashMap<String, String>,
|
||||
) -> Result<SaveResult> {
|
||||
let temp_dir = store_path.join("temp").join(timestamp);
|
||||
std::fs::create_dir_all(&temp_dir).context("failed to create temp dir")?;
|
||||
|
|
@ -92,8 +94,20 @@ fn save_with(
|
|||
.collect();
|
||||
let browser_args = format!("[{}]", quoted.join(","));
|
||||
|
||||
let out = Command::new(single_file)
|
||||
.arg(url)
|
||||
// Write cookie file if cookies are provided.
|
||||
// Never pass cookie values in process args (ps exposure).
|
||||
let cookie_file: Option<std::path::PathBuf> = if !cookies.is_empty() {
|
||||
let cf = temp_dir.join("cookies.txt");
|
||||
let domain = domain_from_url(url);
|
||||
write_netscape_cookie_file(cookies, &domain, &cf)
|
||||
.context("failed to write single-file cookie file")?;
|
||||
Some(cf)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let mut cmd = Command::new(single_file);
|
||||
cmd.arg(url)
|
||||
.arg(&out_file)
|
||||
.arg(format!("--browser-executable-path={chrome}"))
|
||||
.arg("--browser-headless")
|
||||
|
|
@ -120,9 +134,21 @@ fn save_with(
|
|||
// Preserve fonts: defaults strip @font-face rules deemed "unused" or
|
||||
// "alternative" (unicode-range subsets), losing CDN-served fonts.
|
||||
.arg("--remove-unused-fonts=false")
|
||||
.arg("--remove-alternative-fonts=false")
|
||||
.arg("--remove-alternative-fonts=false");
|
||||
if let Some(cf) = &cookie_file {
|
||||
cmd.arg(format!("--browser-cookies-file={}", cf.display()));
|
||||
}
|
||||
let spawn_result = cmd
|
||||
.output()
|
||||
.with_context(|| format!("failed to spawn {single_file} process"))?;
|
||||
.with_context(|| format!("failed to spawn {single_file} process"));
|
||||
|
||||
// Delete cookie file unconditionally — including on spawn failure —
|
||||
// so secrets are never left in store/temp when the capture fails.
|
||||
if let Some(cf) = &cookie_file {
|
||||
let _ = std::fs::remove_file(cf);
|
||||
}
|
||||
|
||||
let out = spawn_result?;
|
||||
|
||||
if !out.status.success() {
|
||||
let stderr = String::from_utf8_lossy(&out.stderr);
|
||||
|
|
@ -296,6 +322,7 @@ mod tests {
|
|||
"test-ts",
|
||||
"/nonexistent/single-file",
|
||||
"chromium",
|
||||
&HashMap::new(),
|
||||
);
|
||||
let err = result.unwrap_err();
|
||||
let msg = format!("{err:#}");
|
||||
|
|
|
|||
|
|
@ -68,7 +68,13 @@ fn build_scraper_args(
|
|||
///
|
||||
/// Requires `ARCHIVR_TWITTER_CREDENTIALS_FILE` to be set. The scraper binary
|
||||
/// can be overridden via `ARCHIVR_TWEET_SCRAPER` and `ARCHIVR_TWEET_PYTHON`.
|
||||
pub fn archive(path: &str, thread: bool, store_path: &Path, timestamp: &str) -> Result<bool> {
|
||||
pub fn archive(
|
||||
path: &str,
|
||||
thread: bool,
|
||||
store_path: &Path,
|
||||
timestamp: &str,
|
||||
cookies: &HashMap<String, String>,
|
||||
) -> Result<bool> {
|
||||
let invocation_cwd = env::current_dir().context("Failed to read current working directory")?;
|
||||
// Output directory for Tweet JSON files.
|
||||
let output_dir = store_path.join("raw_tweets");
|
||||
|
|
@ -93,20 +99,59 @@ pub fn archive(path: &str, thread: bool, store_path: &Path, timestamp: &str) ->
|
|||
.unwrap_or_else(|| PathBuf::from("vendor/twitter/scrape_user_tweet_contents.py"));
|
||||
let scraper_path = absolutize_path_from_cwd(scraper_path, &invocation_cwd);
|
||||
|
||||
let credentials_file = if let Some(credentials_file) =
|
||||
env::var_os("ARCHIVR_TWITTER_CREDENTIALS_FILE")
|
||||
{
|
||||
absolutize_path_from_cwd(PathBuf::from(credentials_file), &invocation_cwd)
|
||||
// Credentials: only use cookie rules as Twitter credentials when the resolved
|
||||
// map contains both `ct0` AND `auth_token` — the two cookies the scraper
|
||||
// requires. A global cookie rule for an unrelated site must not suppress the
|
||||
// ARCHIVR_TWITTER_CREDENTIALS_FILE fallback.
|
||||
// Fall back to ARCHIVR_TWITTER_CREDENTIALS_FILE otherwise.
|
||||
// The temp file is written in the semicolon-delimited format the Python scraper
|
||||
// expects (`ct0=val;auth_token=val`) and deleted unconditionally after the
|
||||
// subprocess returns so secrets are never left on disk on failure.
|
||||
let temp_creds_path: Option<PathBuf>;
|
||||
let credentials_file: PathBuf;
|
||||
|
||||
let has_twitter_cookies =
|
||||
cookies.contains_key("ct0") && cookies.contains_key("auth_token");
|
||||
|
||||
if has_twitter_cookies {
|
||||
let cf = store_path.join("temp").join(timestamp).join("twitter-creds.txt");
|
||||
// Semicolon-separated, no spaces — matches what the scraper's split(";") expects.
|
||||
let creds_str = cookies
|
||||
.iter()
|
||||
.map(|(k, v)| format!("{k}={v}"))
|
||||
.collect::<Vec<_>>()
|
||||
.join(";");
|
||||
{
|
||||
use std::io::Write;
|
||||
#[cfg(unix)]
|
||||
let mut f = {
|
||||
use std::os::unix::fs::OpenOptionsExt;
|
||||
std::fs::OpenOptions::new()
|
||||
.write(true).create(true).truncate(true).mode(0o600)
|
||||
.open(&cf)
|
||||
.context("failed to write twitter credentials file")?
|
||||
};
|
||||
#[cfg(not(unix))]
|
||||
let mut f = std::fs::File::create(&cf)
|
||||
.context("failed to write twitter credentials file")?;
|
||||
f.write_all(creds_str.as_bytes())
|
||||
.context("failed to write twitter credentials file")?;
|
||||
}
|
||||
temp_creds_path = Some(cf.clone());
|
||||
credentials_file = cf;
|
||||
} else if let Some(env_path) = env::var_os("ARCHIVR_TWITTER_CREDENTIALS_FILE") {
|
||||
credentials_file = absolutize_path_from_cwd(PathBuf::from(env_path), &invocation_cwd);
|
||||
temp_creds_path = None;
|
||||
if !credentials_file.is_file() {
|
||||
bail!(
|
||||
"Twitter credentials file not found: {}",
|
||||
credentials_file.display()
|
||||
);
|
||||
}
|
||||
} else {
|
||||
bail!(
|
||||
"Twitter scraping requires ARCHIVR_TWITTER_CREDENTIALS_FILE to point to a cookies file."
|
||||
);
|
||||
};
|
||||
|
||||
if !credentials_file.is_file() {
|
||||
bail!(
|
||||
"Twitter credentials file not found: {}",
|
||||
credentials_file.display()
|
||||
"Twitter scraping requires either cookie rules for x.com/twitter.com \
|
||||
or ARCHIVR_TWITTER_CREDENTIALS_FILE to be set."
|
||||
);
|
||||
}
|
||||
|
||||
|
|
@ -116,12 +161,21 @@ pub fn archive(path: &str, thread: bool, store_path: &Path, timestamp: &str) ->
|
|||
cmd.arg(arg);
|
||||
}
|
||||
|
||||
let output = cmd.output().with_context(|| {
|
||||
// Hold the Result so we can delete the credentials file before propagating
|
||||
// any error — including spawn failures.
|
||||
let spawn_result = cmd.output().with_context(|| {
|
||||
format!(
|
||||
"Failed to spawn tweet scraper at {}",
|
||||
scraper_path.display()
|
||||
)
|
||||
})?;
|
||||
});
|
||||
|
||||
// Remove temp credentials file unconditionally — secrets must not persist.
|
||||
if let Some(cf) = &temp_creds_path {
|
||||
let _ = fs::remove_file(cf);
|
||||
}
|
||||
|
||||
let output = spawn_result?;
|
||||
|
||||
if !output.status.success() {
|
||||
let stderr = String::from_utf8_lossy(&output.stderr);
|
||||
|
|
@ -465,7 +519,7 @@ mod tests {
|
|||
fs::write(&credentials, "ct0=test;auth_token=test").unwrap();
|
||||
set_test_env("ARCHIVR_TWITTER_CREDENTIALS_FILE", &credentials);
|
||||
|
||||
let archived = archive("tweet:123", false, &store_path, "ts").unwrap();
|
||||
let archived = archive("tweet:123", false, &store_path, "ts", &HashMap::new()).unwrap();
|
||||
|
||||
assert!(!archived);
|
||||
|
||||
|
|
@ -538,7 +592,7 @@ EOF
|
|||
set_test_env("ARCHIVR_TWEET_SCRAPER", &script);
|
||||
set_test_env("ARCHIVR_TWEET_PYTHON", "/bin/sh");
|
||||
|
||||
let archived = archive("tweet:123", false, &store_path, "ts").unwrap();
|
||||
let archived = archive("tweet:123", false, &store_path, "ts", &HashMap::new()).unwrap();
|
||||
let tweet_file = output_dir.join("tweet-123.json");
|
||||
let contents = fs::read_to_string(&tweet_file).unwrap();
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,13 @@
|
|||
use anyhow::{bail, Context, Result};
|
||||
use std::{env, path::{Path, PathBuf}, process::Command};
|
||||
use std::{
|
||||
collections::HashMap,
|
||||
env,
|
||||
path::{Path, PathBuf},
|
||||
process::Command,
|
||||
};
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::downloader::cookies::{domain_from_url, write_netscape_cookie_file};
|
||||
use crate::hash::hash_file;
|
||||
|
||||
/// Returns the yt-dlp `-f` format selector for `quality`.
|
||||
|
|
@ -96,6 +103,7 @@ pub fn download(
|
|||
store_path: &Path,
|
||||
timestamp: &String,
|
||||
quality: Option<&str>,
|
||||
cookies: &HashMap<String, String>,
|
||||
) -> Result<(String, String)> {
|
||||
println!("Downloading with yt-dlp: {path}");
|
||||
|
||||
|
|
@ -105,6 +113,18 @@ pub fn download(
|
|||
let temp_dir = store_path.join("temp").join(timestamp);
|
||||
std::fs::create_dir_all(&temp_dir)?;
|
||||
|
||||
// Write a restrictive-permissions cookie file if cookies are provided.
|
||||
// Never pass cookie values in process args (ps exposure).
|
||||
let cookie_file: Option<PathBuf> = if !cookies.is_empty() {
|
||||
let cf_path = temp_dir.join("cookies.txt");
|
||||
let domain = domain_from_url(&path);
|
||||
write_netscape_cookie_file(cookies, &domain, &cf_path)
|
||||
.context("failed to write yt-dlp cookie file")?;
|
||||
Some(cf_path)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
// %(ext)s lets yt-dlp write the correct extension for the chosen format.
|
||||
let out_template = temp_dir.join(format!("{timestamp}.%(ext)s"));
|
||||
|
||||
|
|
@ -123,12 +143,21 @@ pub fn download(
|
|||
// Force the video container to mp4 so we always have a known extension.
|
||||
cmd.arg("--merge-output-format").arg("mp4");
|
||||
}
|
||||
if let Some(cf) = &cookie_file {
|
||||
cmd.arg("--cookies").arg(cf);
|
||||
}
|
||||
let out = cmd
|
||||
.arg("-o")
|
||||
.arg(&out_template)
|
||||
.output()
|
||||
.with_context(|| format!("failed to spawn {ytdlp} process"))?;
|
||||
.with_context(|| format!("failed to spawn {ytdlp} process"));
|
||||
|
||||
// Remove cookie file immediately regardless of outcome.
|
||||
if let Some(cf) = &cookie_file {
|
||||
let _ = std::fs::remove_file(cf);
|
||||
}
|
||||
|
||||
let out = out?;
|
||||
if !out.status.success() {
|
||||
let stderr = String::from_utf8_lossy(&out.stderr);
|
||||
bail!("yt-dlp failed: {stderr}");
|
||||
|
|
@ -166,19 +195,39 @@ fn find_downloaded_file(temp_dir: &Path, timestamp: &str) -> Result<PathBuf> {
|
|||
/// This is a simulate call — it does NOT download any media.
|
||||
/// On failure (non-zero exit or no stdout), prints the captured stderr
|
||||
/// to stderr (for debugging) then returns `None` so callers can proceed.
|
||||
pub fn fetch_metadata(path: &str) -> Option<String> {
|
||||
pub fn fetch_metadata(path: &str, cookies: &HashMap<String, String>) -> Option<String> {
|
||||
let ytdlp = std::env::var("ARCHIVR_YT_DLP").unwrap_or_else(|_| "yt-dlp".to_string());
|
||||
|
||||
let out = std::process::Command::new(&ytdlp)
|
||||
.arg("--dump-json")
|
||||
// Write a temp cookie file if needed; UUID-named to avoid collisions.
|
||||
let cookie_file: Option<PathBuf> = if !cookies.is_empty() {
|
||||
let domain = domain_from_url(path);
|
||||
let p = std::env::temp_dir()
|
||||
.join(format!("archivr-cookies-{}.txt", Uuid::new_v4().simple()));
|
||||
write_netscape_cookie_file(cookies, &domain, &p).ok()?;
|
||||
Some(p)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let mut cmd = std::process::Command::new(&ytdlp);
|
||||
cmd.arg("--dump-json")
|
||||
// Same rationale as download(): only called for single-item sources;
|
||||
// prevents --dump-json from emitting one JSON object per playlist item
|
||||
// when the URL contains a list= parameter.
|
||||
.arg("--no-playlist")
|
||||
.arg(path)
|
||||
.output()
|
||||
.ok()?;
|
||||
.arg("--no-playlist");
|
||||
if let Some(cf) = &cookie_file {
|
||||
cmd.arg("--cookies").arg(cf);
|
||||
}
|
||||
cmd.arg(path);
|
||||
|
||||
let out = cmd.output().ok();
|
||||
|
||||
// Remove cookie file regardless of outcome.
|
||||
if let Some(cf) = &cookie_file {
|
||||
let _ = std::fs::remove_file(cf);
|
||||
}
|
||||
|
||||
let out = out?;
|
||||
if !out.status.success() {
|
||||
let stderr = String::from_utf8_lossy(&out.stderr);
|
||||
eprintln!(
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue