youtube-legend-cli 0.4.0

//! Web-application-firewall classification and retry escalation.
//!
//! # Why status codes are not enough
//!
//! A `403` tells you a request was refused; it does not tell you *by
//! whom*, and the remediation differs sharply by vendor. Cloudflare
//! wants a `cf_clearance` cookie carried forward; `DataDome` wants its
//! own; Kasada invalidates on a header the others never send. A blind
//! retry against the wrong vendor burns the session and raises the risk
//! score.
//!
//! Classification here therefore runs on **response headers and
//! cookies**, which name the vendor directly, and never on the status
//! code alone. A `200` behind a Cloudflare interstitial is still a
//! Cloudflare challenge, and a `403` from an unprotected origin is just
//! a `403`.
//!
//! # Escalation
//!
//! [`EscalationPolicy`] counts consecutive failures against a single
//! vendor and returns [`EscalationVerdict::Abort`] once
//! `max_consecutive_failures` is reached. Switching vendor or
//! seeing a success resets the counter, because a different vendor is a
//! different problem.

use std::fmt;

/// Consecutive failures against the same vendor before the caller must
/// stop retrying. Chosen so that a transient challenge (which usually
/// clears within one or two attempts) is survivable, while a hard block
/// is not hammered.
///
/// Compiled default behind `net.waf.max_consecutive_failures`; read it
/// through `max_consecutive_failures`.
pub const DEFAULT_MAX_CONSECUTIVE_FAILURES: u32 = 5;

/// Failures tolerated before the circuit opens.
///
/// Resolves `net.waf.max_consecutive_failures`, falling back to
/// [`DEFAULT_MAX_CONSECUTIVE_FAILURES`]. A budget of zero would abort
/// before the first attempt, so the accepted range starts at one.
#[must_use]
pub fn max_consecutive_failures() -> u32 {
    crate::config::tuning_u32_in_range(
        "net.waf.max_consecutive_failures",
        DEFAULT_MAX_CONSECUTIVE_FAILURES,
        1,
        100,
    )
}

/// The protection vendors this crate can name.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum WafVendor {
    /// Cloudflare, including Turnstile and the managed challenge.
    Cloudflare,
    /// Akamai Bot Manager / Kona.
    Akamai,
    /// `DataDome`.
    DataDome,
    /// `PerimeterX` (now `HUMAN`).
    PerimeterX,
    /// Imperva, formerly Incapsula.
    Imperva,
    /// Kasada.
    Kasada,
    /// AWS WAF.
    AwsWaf,
}

impl WafVendor {
    /// Stable lowercase identifier, safe for logs and JSON.
    #[must_use]
    pub const fn as_str(self) -> &'static str {
        match self {
            Self::Cloudflare => "cloudflare",
            Self::Akamai => "akamai",
            Self::DataDome => "datadome",
            Self::PerimeterX => "perimeterx",
            Self::Imperva => "imperva",
            Self::Kasada => "kasada",
            Self::AwsWaf => "aws-waf",
        }
    }

    /// Parse the identifier [`Self::as_str`] emits, case-insensitively.
    ///
    /// Returns `None` for a token that names no vendor this crate can
    /// reason about, which is how a configured signature table refuses
    /// a typo instead of silently detecting the wrong vendor.
    #[must_use]
    pub fn from_token(token: &str) -> Option<Self> {
        match token.trim().to_ascii_lowercase().as_str() {
            "cloudflare" => Some(Self::Cloudflare),
            "akamai" => Some(Self::Akamai),
            "datadome" => Some(Self::DataDome),
            "perimeterx" => Some(Self::PerimeterX),
            "imperva" => Some(Self::Imperva),
            "kasada" => Some(Self::Kasada),
            "aws-waf" => Some(Self::AwsWaf),
            _ => None,
        }
    }
}

impl fmt::Display for WafVendor {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(self.as_str())
    }
}

/// Where a vendor gave itself away.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum WafSignal {
    /// A response header, by name.
    Header(String),
    /// A cookie, by name.
    Cookie(String),
}

impl WafSignal {
    /// The name that matched, without the header/cookie distinction.
    #[must_use]
    pub fn name(&self) -> &str {
        match self {
            Self::Header(n) | Self::Cookie(n) => n,
        }
    }

    /// `"header"` or `"cookie"`.
    #[must_use]
    pub const fn kind(&self) -> &'static str {
        match self {
            Self::Header(_) => "header",
            Self::Cookie(_) => "cookie",
        }
    }
}

impl fmt::Display for WafSignal {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}:{}", self.kind(), self.name())
    }
}

/// A vendor plus the single signal that identified it.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct WafDetection {
    /// The vendor in front of the origin.
    pub vendor: WafVendor,
    /// The header or cookie that named it.
    pub signal: WafSignal,
}

/// Exact response-header names that identify a vendor.
///
/// Compiled default behind `net.waf.header_signatures`.
const HEADER_SIGNATURES: &[(&str, WafVendor)] = &[
    ("cf-ray", WafVendor::Cloudflare),
    ("cf-cache-status", WafVendor::Cloudflare),
    ("cf-mitigated", WafVendor::Cloudflare),
    ("akamai-origin-hop", WafVendor::Akamai),
    ("akamai-grn", WafVendor::Akamai),
    ("x-datadome", WafVendor::DataDome),
    ("x-datadome-cid", WafVendor::DataDome),
    ("x-iinfo", WafVendor::Imperva),
    ("x-cdn", WafVendor::Imperva),
    ("x-kpsdk-ct", WafVendor::Kasada),
];

/// Header-name prefixes that identify a vendor. Checked after the exact
/// table, because an exact match is the stronger signal.
///
/// Compiled default behind `net.waf.header_prefix_signatures`.
const HEADER_PREFIX_SIGNATURES: &[(&str, WafVendor)] = &[
    ("x-akamai-", WafVendor::Akamai),
    ("x-px-", WafVendor::PerimeterX),
    ("x-amzn-waf-", WafVendor::AwsWaf),
];

/// Exact cookie names that identify a vendor.
///
/// Compiled default behind `net.waf.cookie_signatures`.
const COOKIE_SIGNATURES: &[(&str, WafVendor)] = &[
    ("cf_clearance", WafVendor::Cloudflare),
    ("__cf_bm", WafVendor::Cloudflare),
    ("__cflb", WafVendor::Cloudflare),
    ("ak_bmsc", WafVendor::Akamai),
    ("bm_sv", WafVendor::Akamai),
    ("_abck", WafVendor::Akamai),
    ("datadome", WafVendor::DataDome),
    ("_px3", WafVendor::PerimeterX),
    ("_pxhd", WafVendor::PerimeterX),
    ("_pxvid", WafVendor::PerimeterX),
    ("___utmvc", WafVendor::Imperva),
    ("x-kpsdk-ct", WafVendor::Kasada),
    ("aws-waf-token", WafVendor::AwsWaf),
];

/// Cookie-name prefixes that identify a vendor. Imperva session and
/// visitor cookies carry a per-site numeric suffix, so they can only be
/// matched by prefix.
///
/// Compiled default behind `net.waf.cookie_prefix_signatures`.
const COOKIE_PREFIX_SIGNATURES: &[(&str, WafVendor)] = &[
    ("incap_ses_", WafVendor::Imperva),
    ("visid_incap_", WafVendor::Imperva),
    ("nlbi_", WafVendor::Imperva),
];

/// Cookies that, once obtained, prove a challenge was already solved.
/// Discarding one of these forces a fresh challenge, which is both slow
/// and a fresh chance to be blocked.
///
/// Compiled default behind `net.waf.challenge_cookies`; read it through
/// [`challenge_cookies`].
pub const DEFAULT_CHALLENGE_COOKIES: &[&str] = &[
    "cf_clearance",
    "__cf_bm",
    "datadome",
    "_px3",
    "ak_bmsc",
    "_abck",
    "aws-waf-token",
];

/// Classifies the protection in front of a host from its response
/// headers and cookies.
///
/// `headers` is a list of `(name, value)` pairs; names are matched
/// case-insensitively, as HTTP requires. `cookie_names` is the set of
/// cookie names observed, whether from `Set-Cookie` or from the jar.
///
/// Returns the first match found, header table first. `None` means no
/// vendor was named — which is *not* the same as "no protection", only
/// "nothing identified itself".
///
/// ```
/// use youtube_legend_cli::net::{detect_waf, WafSignal, WafVendor};
///
/// let headers = [("CF-RAY".to_owned(), "8a1b2c3d4e5f".to_owned())];
/// let hit = detect_waf(&headers, &[]).expect("cloudflare must be named");
/// assert_eq!(hit.vendor, WafVendor::Cloudflare);
/// assert_eq!(hit.signal, WafSignal::Header("cf-ray".to_owned()));
/// ```
#[must_use]
pub fn detect(headers: &[(String, String)], cookie_names: &[String]) -> Option<WafDetection> {
    for (name, _) in headers {
        let lower = name.to_ascii_lowercase();
        if let Some((_, vendor)) = header_signatures().iter().find(|(sig, _)| *sig == lower) {
            return Some(WafDetection {
                vendor: *vendor,
                signal: WafSignal::Header(lower),
            });
        }
    }

    for (name, _) in headers {
        let lower = name.to_ascii_lowercase();
        if let Some((_, vendor)) = header_prefix_signatures()
            .iter()
            .find(|(prefix, _)| lower.starts_with(prefix.as_str()))
        {
            return Some(WafDetection {
                vendor: *vendor,
                signal: WafSignal::Header(lower),
            });
        }
    }

    for name in cookie_names {
        let lower = name.to_ascii_lowercase();
        if let Some((_, vendor)) = cookie_signatures().iter().find(|(sig, _)| *sig == lower) {
            return Some(WafDetection {
                vendor: *vendor,
                signal: WafSignal::Cookie(lower),
            });
        }
    }

    for name in cookie_names {
        let lower = name.to_ascii_lowercase();
        if let Some((_, vendor)) = cookie_prefix_signatures()
            .iter()
            .find(|(prefix, _)| lower.starts_with(prefix.as_str()))
        {
            return Some(WafDetection {
                vendor: *vendor,
                signal: WafSignal::Cookie(lower),
            });
        }
    }

    None
}

/// Resolve a signature table from configuration.
///
/// Each configured entry is `signature=vendor`, for example
/// `cf-ray=cloudflare`. An entry naming a vendor this crate does not
/// know is skipped with a warning rather than aborting the run.
fn resolve_signatures(key: &str, compiled: &[(&str, WafVendor)]) -> Vec<(String, WafVendor)> {
    let fallback = || -> Vec<(String, WafVendor)> {
        compiled
            .iter()
            .map(|(sig, vendor)| ((*sig).to_string(), *vendor))
            .collect()
    };
    let Some(entries) = crate::config::tuning_str_list(key) else {
        return fallback();
    };
    let mut out = Vec::with_capacity(entries.len());
    for entry in &entries {
        match entry.split_once('=') {
            Some((sig, vendor)) => match WafVendor::from_token(vendor.trim()) {
                Some(vendor) if !sig.trim().is_empty() => {
                    out.push((sig.trim().to_ascii_lowercase(), vendor));
                }
                _ => tracing::warn!(key, entry, "ignoring an entry with an unknown vendor"),
            },
            None => tracing::warn!(
                key,
                entry,
                "ignoring an entry that is not `signature=vendor`"
            ),
        }
    }
    if out.is_empty() {
        tracing::warn!(
            key,
            "no usable signature configured; keeping the compiled table"
        );
        return fallback();
    }
    out
}

/// Exact response-header names that identify a vendor.
///
/// Resolves `net.waf.header_signatures`.
#[must_use]
pub fn header_signatures() -> Vec<(String, WafVendor)> {
    resolve_signatures("net.waf.header_signatures", HEADER_SIGNATURES)
}

/// Header-name prefixes that identify a vendor.
///
/// Resolves `net.waf.header_prefix_signatures`.
#[must_use]
pub fn header_prefix_signatures() -> Vec<(String, WafVendor)> {
    resolve_signatures("net.waf.header_prefix_signatures", HEADER_PREFIX_SIGNATURES)
}

/// Exact cookie names that identify a vendor.
///
/// Resolves `net.waf.cookie_signatures`.
#[must_use]
pub fn cookie_signatures() -> Vec<(String, WafVendor)> {
    resolve_signatures("net.waf.cookie_signatures", COOKIE_SIGNATURES)
}

/// Cookie-name prefixes that identify a vendor.
///
/// Resolves `net.waf.cookie_prefix_signatures`.
#[must_use]
pub fn cookie_prefix_signatures() -> Vec<(String, WafVendor)> {
    resolve_signatures("net.waf.cookie_prefix_signatures", COOKIE_PREFIX_SIGNATURES)
}

/// Cookies that mark an interactive challenge and must survive a prune.
///
/// Resolves `net.waf.challenge_cookies`, falling back to
/// [`DEFAULT_CHALLENGE_COOKIES`].
#[must_use]
pub fn challenge_cookies() -> Vec<String> {
    crate::config::tuning_str_list_or("net.waf.challenge_cookies", DEFAULT_CHALLENGE_COOKIES)
}

/// Returns `true` when `name` is a challenge-clearance cookie that must
/// survive a jar prune or a session reset.
#[must_use]
pub fn is_challenge_cookie(name: &str) -> bool {
    let lower = name.to_ascii_lowercase();
    challenge_cookies().contains(&lower)
}

/// What the caller should do after a failure.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum EscalationVerdict {
    /// Another attempt is permitted.
    Retry,
    /// The consecutive-failure budget for this vendor is spent.
    Abort,
}

/// Counts consecutive failures per vendor and decides when to stop.
///
/// The counter is per *current* vendor: seeing a different vendor
/// restarts it, because a `DataDome` block after four Cloudflare blocks
/// is a new problem with a new budget.
///
/// ```
/// use youtube_legend_cli::net::{EscalationPolicy, EscalationVerdict, WafVendor};
///
/// let mut policy = EscalationPolicy::new(2);
/// assert_eq!(policy.record_failure(WafVendor::Cloudflare), EscalationVerdict::Retry);
/// assert_eq!(policy.record_failure(WafVendor::Cloudflare), EscalationVerdict::Abort);
/// policy.record_success();
/// assert_eq!(policy.record_failure(WafVendor::Cloudflare), EscalationVerdict::Retry);
/// ```
#[derive(Debug, Clone)]
pub struct EscalationPolicy {
    max_consecutive_failures: u32,
    current: Option<WafVendor>,
    consecutive: u32,
}

impl Default for EscalationPolicy {
    fn default() -> Self {
        Self::new(max_consecutive_failures())
    }
}

impl EscalationPolicy {
    /// Builds a policy with an explicit budget. A budget of `0` is
    /// raised to `1`, so that at least one attempt is always made.
    #[must_use]
    pub const fn new(max_consecutive_failures: u32) -> Self {
        Self {
            max_consecutive_failures: if max_consecutive_failures == 0 {
                1
            } else {
                max_consecutive_failures
            },
            current: None,
            consecutive: 0,
        }
    }

    /// The configured budget.
    #[must_use]
    pub const fn budget(&self) -> u32 {
        self.max_consecutive_failures
    }

    /// Failures recorded against the vendor currently being fought.
    #[must_use]
    pub const fn consecutive_failures(&self) -> u32 {
        self.consecutive
    }

    /// The vendor the counter currently applies to.
    #[must_use]
    pub const fn current_vendor(&self) -> Option<WafVendor> {
        self.current
    }

    /// Records one failure attributed to `vendor` and returns whether a
    /// further attempt is permitted.
    pub fn record_failure(&mut self, vendor: WafVendor) -> EscalationVerdict {
        if self.current != Some(vendor) {
            self.current = Some(vendor);
            self.consecutive = 0;
        }
        self.consecutive = self.consecutive.saturating_add(1);
        if self.consecutive >= self.max_consecutive_failures {
            EscalationVerdict::Abort
        } else {
            EscalationVerdict::Retry
        }
    }

    /// Records a failure that no vendor claimed. Unattributed failures
    /// are counted under the vendor already being fought when there is
    /// one, and otherwise start their own unattributed run.
    pub fn record_unclassified_failure(&mut self) -> EscalationVerdict {
        self.consecutive = self.consecutive.saturating_add(1);
        if self.consecutive >= self.max_consecutive_failures {
            EscalationVerdict::Abort
        } else {
            EscalationVerdict::Retry
        }
    }

    /// Clears the counter after a successful request.
    pub fn record_success(&mut self) {
        self.current = None;
        self.consecutive = 0;
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn h(pairs: &[(&str, &str)]) -> Vec<(String, String)> {
        pairs
            .iter()
            .map(|(k, v)| ((*k).to_owned(), (*v).to_owned()))
            .collect()
    }

    fn c(names: &[&str]) -> Vec<String> {
        names.iter().map(|n| (*n).to_owned()).collect()
    }

    #[test]
    fn cloudflare_is_detected_by_header() {
        for name in ["cf-ray", "CF-Ray", "cf-cache-status", "cf-mitigated"] {
            let hit = detect(&h(&[(name, "x")]), &[]).expect("must detect");
            assert_eq!(hit.vendor, WafVendor::Cloudflare, "header {name}");
            assert_eq!(hit.signal.kind(), "header");
        }
    }

    #[test]
    fn cloudflare_is_detected_by_cookie() {
        for name in ["cf_clearance", "__cf_bm", "__cflb"] {
            let hit = detect(&[], &c(&[name])).expect("must detect");
            assert_eq!(hit.vendor, WafVendor::Cloudflare, "cookie {name}");
            assert_eq!(hit.signal, WafSignal::Cookie(name.to_owned()));
        }
    }

    #[test]
    fn akamai_is_detected_by_header_prefix_and_cookie() {
        assert_eq!(
            detect(&h(&[("X-Akamai-Transformed", "9")]), &[])
                .expect("must detect")
                .vendor,
            WafVendor::Akamai
        );
        assert_eq!(
            detect(&h(&[("akamai-origin-hop", "2")]), &[])
                .expect("must detect")
                .vendor,
            WafVendor::Akamai
        );
        assert_eq!(
            detect(&[], &c(&["ak_bmsc"])).expect("must detect").vendor,
            WafVendor::Akamai
        );
    }

    #[test]
    fn datadome_is_detected_by_cookie() {
        let hit = detect(&[], &c(&["datadome"])).expect("must detect");
        assert_eq!(hit.vendor, WafVendor::DataDome);
        assert_eq!(hit.signal, WafSignal::Cookie("datadome".to_owned()));
    }

    #[test]
    fn perimeterx_is_detected_by_cookies_and_header_prefix() {
        for name in ["_px3", "_pxhd", "_pxvid"] {
            assert_eq!(
                detect(&[], &c(&[name])).expect("must detect").vendor,
                WafVendor::PerimeterX,
                "cookie {name}"
            );
        }
        assert_eq!(
            detect(&h(&[("x-px-block", "1")]), &[])
                .expect("must detect")
                .vendor,
            WafVendor::PerimeterX
        );
    }

    #[test]
    fn imperva_is_detected_by_suffixed_cookies() {
        for name in ["incap_ses_1234_5678", "visid_incap_5678", "___utmvc"] {
            assert_eq!(
                detect(&[], &c(&[name])).expect("must detect").vendor,
                WafVendor::Imperva,
                "cookie {name}"
            );
        }
    }

    #[test]
    fn kasada_is_detected_by_cookie() {
        assert_eq!(
            detect(&[], &c(&["x-kpsdk-ct"]))
                .expect("must detect")
                .vendor,
            WafVendor::Kasada
        );
    }

    #[test]
    fn aws_waf_is_detected_by_header_prefix() {
        assert_eq!(
            detect(&h(&[("x-amzn-waf-action", "block")]), &[])
                .expect("must detect")
                .vendor,
            WafVendor::AwsWaf
        );
    }

    #[test]
    fn a_bare_403_names_no_vendor() {
        // The whole point: status code is not an input, and ordinary
        // headers must not be misread as a signature.
        let ordinary = h(&[
            ("content-type", "text/html"),
            ("server", "nginx"),
            ("x-frame-options", "DENY"),
        ]);
        assert_eq!(detect(&ordinary, &c(&["session_id", "csrftoken"])), None);
    }

    #[test]
    fn header_signal_wins_over_cookie_signal() {
        let hit = detect(&h(&[("cf-ray", "abc")]), &c(&["datadome"])).expect("must detect");
        assert_eq!(hit.vendor, WafVendor::Cloudflare);
        assert_eq!(hit.signal.kind(), "header");
    }

    #[test]
    fn challenge_cookies_are_recognised() {
        assert!(is_challenge_cookie("cf_clearance"));
        assert!(is_challenge_cookie("CF_CLEARANCE"));
        assert!(is_challenge_cookie("datadome"));
        assert!(!is_challenge_cookie("session_id"));
    }

    #[test]
    fn escalation_aborts_at_the_budget() {
        let mut policy = EscalationPolicy::new(3);
        assert_eq!(policy.budget(), 3);
        assert_eq!(
            policy.record_failure(WafVendor::Cloudflare),
            EscalationVerdict::Retry
        );
        assert_eq!(
            policy.record_failure(WafVendor::Cloudflare),
            EscalationVerdict::Retry
        );
        assert_eq!(
            policy.record_failure(WafVendor::Cloudflare),
            EscalationVerdict::Abort
        );
        assert_eq!(policy.consecutive_failures(), 3);
    }

    #[test]
    fn switching_vendor_restarts_the_budget() {
        let mut policy = EscalationPolicy::new(2);
        assert_eq!(
            policy.record_failure(WafVendor::Cloudflare),
            EscalationVerdict::Retry
        );
        assert_eq!(
            policy.record_failure(WafVendor::DataDome),
            EscalationVerdict::Retry
        );
        assert_eq!(policy.current_vendor(), Some(WafVendor::DataDome));
        assert_eq!(policy.consecutive_failures(), 1);
    }

    #[test]
    fn success_clears_the_counter() {
        let mut policy = EscalationPolicy::new(2);
        let _ = policy.record_failure(WafVendor::Kasada);
        policy.record_success();
        assert_eq!(policy.consecutive_failures(), 0);
        assert_eq!(policy.current_vendor(), None);
        assert_eq!(
            policy.record_failure(WafVendor::Kasada),
            EscalationVerdict::Retry
        );
    }

    #[test]
    fn zero_budget_still_allows_one_attempt_before_aborting() {
        let mut policy = EscalationPolicy::new(0);
        assert_eq!(policy.budget(), 1);
        assert_eq!(
            policy.record_failure(WafVendor::Imperva),
            EscalationVerdict::Abort
        );
    }

    #[test]
    fn default_budget_matches_the_named_constant() {
        assert_eq!(
            EscalationPolicy::default().budget(),
            max_consecutive_failures()
        );
    }

    #[test]
    fn unclassified_failures_are_counted() {
        let mut policy = EscalationPolicy::new(2);
        assert_eq!(
            policy.record_unclassified_failure(),
            EscalationVerdict::Retry
        );
        assert_eq!(
            policy.record_unclassified_failure(),
            EscalationVerdict::Abort
        );
    }
}