astraguard 1.1.7

//! The main [`AstraGuardClient`] and its license-validation methods.

use std::collections::HashMap;
use std::sync::{Arc, Mutex};

use serde_json::json;

use crate::types::{AstraGuardError, LicenseDetails, LicenseResponse, Result, UpdateCheckResponse};

/// Response-token freshness window (ms). A signed response older than this is
/// rejected even if its HMAC and nonce match - matches the server's window.
const RESPONSE_FRESHNESS_MS: i64 = 120_000;

/// The primary entry point for the AstraGuard SDK.
///
/// Create one instance per application and reuse it. The client owns a
/// connection pool and an optional background heartbeat task.
///
/// # Example
/// ```rust,no_run
/// # use astraguard::AstraGuardClient;
/// # #[tokio::main] async fn main() -> Result<(), Box<dyn std::error::Error>> {
/// let client = AstraGuardClient::new("https://api.astraguard.io", "your-product-id")?;
/// let details = client.verify_license("XXXX-XXXX-XXXX-XXXX").await?;
/// if details.raw.valid {
///     println!("License OK, expires: {:?}", details.expires_at);
/// }
/// # Ok(())
/// # }
/// ```
pub struct AstraGuardClient {
    http: reqwest::Client,
    api_url: String,
    product_id: String,
    /// HMAC-SHA256 key for response authentication. Fail-closed when set.
    auth_key: Option<Vec<u8>>,
    auto_enforce_security: bool,
    cloud_config: Arc<Mutex<HashMap<String, String>>>,
    last_license_key: Arc<Mutex<Option<String>>>,
}

impl AstraGuardClient {
    /// Create a new client.
    ///
    /// # Errors
    /// Returns [`AstraGuardError::Build`] if the HTTP client cannot be
    /// initialised (e.g. the TLS backend is unavailable).
    pub fn new(api_url: impl Into<String>, product_id: impl Into<String>) -> Result<Self> {
        let http = reqwest::Client::builder()
            .timeout(std::time::Duration::from_secs(10))
            .user_agent(concat!("AstraGuardSDK-Rust/", env!("CARGO_PKG_VERSION")))
            .build()
            .map_err(|e| AstraGuardError::Build(e.to_string()))?;

        Ok(Self {
            http,
            api_url: api_url.into().trim_end_matches('/').to_string(),
            product_id: product_id.into(),
            auth_key: None,
            auto_enforce_security: false,
            cloud_config: Arc::new(Mutex::new(HashMap::new())),
            last_license_key: Arc::new(Mutex::new(None)),
        })
    }

    /// Set a Base64-encoded HMAC-SHA256 key to authenticate server responses.
    ///
    /// When set, every `/validate` and `/activate` response must carry a valid
    /// response token (`rt`) whose HMAC matches over `nonce|rts|valid|productId`,
    /// whose echoed nonce (`rn`) equals the fresh per-request nonce this client
    /// generated, and whose timestamp (`rts`) is within the freshness window.
    /// This binds every response to the specific request that produced it, so a
    /// captured genuine response cannot be replayed via a mock/MITM server for a
    /// different or fake request. Retrieve the key from the AstraGuard dashboard.
    ///
    /// # Errors
    /// Returns [`AstraGuardError::Build`] if `base64_key` is not valid Base64 -
    /// rather than silently disabling verification, a bad key is surfaced loudly
    /// so a typo can never downgrade the client to "no verification".
    pub fn with_response_auth(mut self, base64_key: &str) -> Result<Self> {
        use base64::Engine;
        let decoded = base64::engine::general_purpose::STANDARD
            .decode(base64_key)
            .map_err(|e| {
                AstraGuardError::Build(format!("invalid response-auth key (not base64): {e}"))
            })?;
        if decoded.is_empty() {
            return Err(AstraGuardError::Build(
                "invalid response-auth key (decoded to empty)".to_string(),
            ));
        }
        self.auth_key = Some(decoded);
        Ok(self)
    }

    /// Enable or disable automatic security checks (anti-debug / anti-VM).
    pub fn set_auto_enforce_security(&mut self, enabled: bool) {
        self.auto_enforce_security = enabled;
    }

    /// Returns the stable hardware ID used to bind this machine's licenses.
    /// Deterministic across calls for the same machine.
    #[must_use]
    pub fn get_machine_id() -> String {
        crate::hwid::get()
    }

    /// Returns the value of a cloud-config variable fetched from the server,
    /// or `None` if the key is absent.
    pub fn cloud_config(&self, key: &str) -> Option<String> {
        self.cloud_config
            .lock()
            .unwrap_or_else(|e| e.into_inner())
            .get(key)
            .cloned()
    }

    // ── Validation ────────────────────────────────────────────────────────────

    /// Validate a license key against the AstraGuard API.
    ///
    /// On network failure the SDK falls back to the local offline cache.
    ///
    /// # Errors
    /// - [`AstraGuardError::LicenseRejected`]  -  license is invalid/expired.
    /// - [`AstraGuardError::Network`]  -  network unavailable and no cache.
    /// - [`AstraGuardError::SignatureMismatch`]  -  HMAC check failed.
    #[must_use = "check `details.raw.valid` to gate access"]
    pub async fn verify_license(&self, license_key: &str) -> Result<LicenseDetails> {
        let hwid = crate::hwid::get();

        if self.auto_enforce_security {
            crate::security::anti_debug::start_background_monitor();
            if crate::security::anti_vm::is_virtual_machine() {
                return Err(AstraGuardError::LicenseRejected(
                    "virtual machine detected".to_string(),
                ));
            }
        }

        let resp = self.call_validate(license_key, &hwid).await;

        match resp {
            Ok(lic) => {
                if lic.valid {
                    crate::cache::save(license_key, &hwid, &self.product_id, &lic);
                    *self
                        .last_license_key
                        .lock()
                        .unwrap_or_else(|e| e.into_inner()) = Some(license_key.to_string());
                    let mut cc = self.cloud_config.lock().unwrap_or_else(|e| e.into_inner());
                    *cc = lic.variables.clone();
                    Ok(parse_details(lic, license_key, false))
                } else {
                    let reason = lic
                        .reason
                        .clone()
                        .unwrap_or_else(|| "invalid_license".to_string());
                    Err(AstraGuardError::LicenseRejected(reason))
                }
            }
            Err(AstraGuardError::Network(_)) | Err(AstraGuardError::ServerError { .. }) => {
                if let Some(cached) = crate::cache::load(license_key, &hwid, &self.product_id) {
                    return Ok(parse_details(cached, license_key, true));
                }
                resp.map(|_| unreachable!())
            }
            Err(e) => Err(e),
        }
    }

    /// Like [`verify_license`] but invokes `on_invalid` on failure instead of
    /// returning an error.
    pub async fn verify_or_callback<F>(&self, license_key: &str, on_invalid: Option<F>)
    where
        F: Fn(),
    {
        match self.verify_license(license_key).await {
            Ok(_) => {}
            Err(e) => {
                eprintln!("[AstraGuard] license check failed: {e}");
                if let Some(f) = on_invalid.as_ref() {
                    f();
                }
            }
        }
    }

    // ── Activate ─────────────────────────────────────────────────────────────

    /// Activate a license key for the first time (binds the current HWID).
    ///
    /// Use this on the customer first run. For subsequent launches use
    /// [`verify_license`] instead.
    pub async fn activate(&self, license_key: &str) -> Result<LicenseResponse> {
        let hwid = crate::hwid::get();
        let nonce = generate_nonce();
        let body = json!({
            "key": license_key,
            "hwid": hwid,
            "productId": self.product_id,
            "nonce": nonce,
            "debugDetected": crate::security::anti_debug::is_attached(),
        });

        let resp = self
            .http
            .post(format!("{}/activate", self.api_url))
            .json(&body)
            .send()
            .await?;

        let status = resp.status().as_u16();
        let text = resp.text().await?;

        if !(200..300).contains(&(status as usize)) {
            return Err(AstraGuardError::ServerError { status, body: text });
        }

        let mut lic: LicenseResponse = serde_json::from_str(&text)?;
        // /activate responses carry `success`, not `valid` (that field is a
        // /validate-only concept). Without this, `lic.valid` would silently
        // stay `false` on every successful activation - a real customer
        // integration checking `resp.valid` after activate() would always
        // see `false` and treat their own valid license as rejected.
        if let Some(success) = lic.success {
            lic.valid = success;
        }
        // Same anti-replay verification as validate(): activate() is the
        // first-run HWID-binding call, so it must be equally MITM-resistant.
        self.verify_response(&lic, &nonce)?;
        Ok(lic)
    }

    // ── Heartbeat ────────────────────────────────────────────────────────────

    /// Start a background task that re-validates the license every
    /// `interval_secs` seconds. Calls `on_invalid` when the license lapses.
    pub fn start_heartbeat<F>(&self, interval_secs: u64, on_invalid: Option<F>)
    where
        F: Fn() + Send + 'static,
    {
        let http = self.http.clone();
        let api_url = self.api_url.clone();
        let product_id = self.product_id.clone();
        let auth_key = self.auth_key.clone();
        let last_key = Arc::clone(&self.last_license_key);

        tokio::spawn(async move {
            let mut interval = tokio::time::interval(std::time::Duration::from_secs(interval_secs));
            loop {
                interval.tick().await;

                let key = {
                    let guard = last_key.lock().unwrap_or_else(|e| e.into_inner());
                    guard.clone()
                };
                let Some(license_key) = key else { continue };

                let hwid = crate::hwid::get();
                let nonce = generate_nonce();
                let body = json!({
                    "key": license_key,
                    "hwid": hwid,
                    "productId": product_id,
                    "nonce": nonce,
                });

                let Ok(resp) = http
                    .post(format!("{}/validate", api_url))
                    .json(&body)
                    .send()
                    .await
                else {
                    continue;
                };

                let Ok(text) = resp.text().await else {
                    continue;
                };

                let Ok(lic) = serde_json::from_str::<LicenseResponse>(&text) else {
                    continue;
                };

                // Reject replayed/forged heartbeat responses the same way the
                // foreground path does - a MITM that pins "valid" forever is
                // exactly what the heartbeat exists to catch.
                if let Some(ref key_bytes) = auth_key {
                    if !verify_response_token(key_bytes, &lic, &nonce, &product_id) {
                        if let Some(f) = on_invalid.as_ref() {
                            f();
                        }
                        continue;
                    }
                }

                if !lic.valid {
                    if let Some(f) = on_invalid.as_ref() {
                        f();
                    }
                }
            }
        });
    }

    // ── Update check ─────────────────────────────────────────────────────────

    /// Check whether a newer release is available for this product.
    ///
    /// `current_version` should be the running application version string
    /// (e.g. `"1.2.3"`).
    pub async fn check_for_updates(&self, current_version: &str) -> Result<UpdateCheckResponse> {
        let encoded = percent_encode(current_version);
        let url = format!(
            "{}/products/{}/check-update?version={}",
            self.api_url, self.product_id, encoded
        );
        let resp = self.http.get(&url).send().await?;
        let status = resp.status().as_u16();
        let text = resp.text().await?;
        if !(200..300).contains(&(status as usize)) {
            return Err(AstraGuardError::ServerError { status, body: text });
        }
        Ok(serde_json::from_str(&text)?)
    }

    // ── Internal ─────────────────────────────────────────────────────────────

    async fn call_validate(&self, license_key: &str, hwid: &str) -> Result<LicenseResponse> {
        let nonce = generate_nonce();
        let body = json!({
            "key": license_key,
            "hwid": hwid,
            "productId": self.product_id,
            "nonce": nonce,
            "debugDetected": crate::security::anti_debug::is_attached(),
        });

        let resp = self
            .http
            .post(format!("{}/validate", self.api_url))
            .json(&body)
            .send()
            .await?;

        let status = resp.status().as_u16();
        let text = resp.text().await?;

        if !(200..300).contains(&(status as usize)) {
            return Err(AstraGuardError::ServerError { status, body: text });
        }

        let lic: LicenseResponse = serde_json::from_str(&text)?;
        self.verify_response(&lic, &nonce)?;
        Ok(lic)
    }

    /// Verifies a server response's anti-replay token against the nonce this
    /// client sent for this specific request.
    ///
    /// Fails closed, not open: without an auth key there is no way to tell a
    /// genuine server reply from one served by a network-level attacker
    /// (hosts-file redirect + a fake server that always answers valid: true)
    /// - that bypass needs zero reverse engineering, it just has to run. Call
    /// `with_response_auth()` with the key from Dashboard -> Products ->
    /// Response Key before validating.
    fn verify_response(&self, lic: &LicenseResponse, expected_nonce: &str) -> Result<()> {
        let Some(ref key_bytes) = self.auth_key else {
            return Err(AstraGuardError::ResponseAuthNotConfigured);
        };
        if !verify_response_token(key_bytes, lic, expected_nonce, &self.product_id) {
            return Err(AstraGuardError::SignatureMismatch);
        }
        Ok(())
    }
}

// ── Anti-replay response-token verification ─────────────────────────────────────

/// Generates a fresh 256-bit random nonce (64 hex chars) for one request,
/// using the OS CSPRNG. Each request gets a distinct nonce so the server's
/// signed response is bound to exactly this request and cannot be replayed.
fn generate_nonce() -> String {
    use rand::RngCore;
    let mut buf = [0u8; 32];
    rand::rngs::OsRng.fill_bytes(&mut buf);
    hex::encode(buf)
}

/// Returns `true` iff the response carries a valid anti-replay token: the
/// HMAC-SHA256 of `nonce|rts|valid|productId` matches `rt`, the echoed nonce
/// `rn` exactly equals the one this client sent, and `rts` is fresh.
///
/// Fails **closed**: any missing field, stale timestamp, nonce mismatch, or
/// HMAC mismatch returns `false`. This is what makes a captured-and-replayed
/// response (the mock-server MITM attack) fail even though its HMAC was once
/// genuine - the replayed `rn` won't match the new request's nonce, and/or
/// `rts` will be stale.
fn verify_response_token(
    key: &[u8],
    lic: &LicenseResponse,
    expected_nonce: &str,
    product_id: &str,
) -> bool {
    use hmac::{Hmac, Mac};
    use sha2::Sha256;

    let (Some(rt), Some(rts), Some(rn)) = (lic.rt.as_ref(), lic.rts, lic.rn.as_ref()) else {
        return false; // key configured but response is unsigned -> reject
    };

    // Nonce echo must match the exact nonce we generated for THIS request.
    if rn != expected_nonce {
        return false;
    }

    // Freshness: reject a signature older/newer than the window (replay of an
    // old capture, or a clock-skewed forgery).
    let now_ms = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .map(|d| d.as_millis() as i64)
        .unwrap_or(0);
    if (now_ms - rts).abs() > RESPONSE_FRESHNESS_MS {
        return false;
    }

    // HMAC over the exact server payload: nonce|rts|valid|productId.
    let payload = format!(
        "{}|{}|{}|{}",
        expected_nonce,
        rts,
        if lic.valid { "1" } else { "0" },
        product_id
    );
    let Ok(expected_rt) = base64::Engine::decode(&base64::engine::general_purpose::STANDARD, rt)
    else {
        return false;
    };
    type HmacSha256 = Hmac<Sha256>;
    let Ok(mut mac) = HmacSha256::new_from_slice(key) else {
        return false;
    };
    mac.update(payload.as_bytes());
    // verify_slice is constant-time.
    mac.verify_slice(&expected_rt).is_ok()
}

// ── Helpers ───────────────────────────────────────────────────────────────────

fn parse_details(lic: LicenseResponse, key: &str, is_offline: bool) -> LicenseDetails {
    use std::time::{Duration, UNIX_EPOCH};

    let expires_at = lic.expires_at.as_deref().and_then(|s| {
        crate::cache::chrono_like_parse_pub(s).map(|secs| UNIX_EPOCH + Duration::from_secs(secs))
    });

    LicenseDetails {
        raw: lic,
        license_key: key.to_string(),
        expires_at,
        is_offline,
    }
}

fn percent_encode(s: &str) -> String {
    let mut out = String::with_capacity(s.len());
    for b in s.bytes() {
        if b.is_ascii_alphanumeric() || b == b'-' || b == b'_' || b == b'.' || b == b'~' {
            out.push(b as char);
        } else {
            out.push('%');
            out.push(
                char::from_digit((b >> 4) as u32, 16)
                    .unwrap_or('0')
                    .to_ascii_uppercase(),
            );
            out.push(
                char::from_digit((b & 0xF) as u32, 16)
                    .unwrap_or('0')
                    .to_ascii_uppercase(),
            );
        }
    }
    out
}

#[cfg(test)]
mod tests {
    use super::*;
    use hmac::{Hmac, Mac};
    use sha2::Sha256;

    const KEY: &[u8] = b"test-response-auth-key-32-bytes!";
    const PID: &str = "prod-123";

    // Mints a server-style response token exactly as the server does:
    // base64(HMAC-SHA256(key, "nonce|rts|valid|productId")).
    fn mint(nonce: &str, rts: i64, valid: bool) -> String {
        let payload = format!(
            "{}|{}|{}|{}",
            nonce,
            rts,
            if valid { "1" } else { "0" },
            PID
        );
        let mut mac = Hmac::<Sha256>::new_from_slice(KEY).unwrap();
        mac.update(payload.as_bytes());
        base64::Engine::encode(
            &base64::engine::general_purpose::STANDARD,
            mac.finalize().into_bytes(),
        )
    }

    fn now_ms() -> i64 {
        std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_millis() as i64
    }

    fn resp(
        rt: Option<String>,
        rts: Option<i64>,
        rn: Option<String>,
        valid: bool,
    ) -> LicenseResponse {
        LicenseResponse {
            valid,
            rt,
            rts,
            rn,
            ..Default::default()
        }
    }

    #[test]
    fn accepts_a_fresh_genuine_token() {
        let nonce = "a".repeat(64);
        let rts = now_ms();
        let rt = mint(&nonce, rts, true);
        let lic = resp(Some(rt), Some(rts), Some(nonce.clone()), true);
        assert!(verify_response_token(KEY, &lic, &nonce, PID));
    }

    #[test]
    fn rejects_replay_with_a_different_request_nonce() {
        // Attacker captured a genuine response for `captured_nonce` and replays
        // it against a NEW request whose nonce is `our_nonce`. The echoed nonce
        // won't match ours -> rejected, even though the HMAC was once genuine.
        let captured_nonce = "a".repeat(64);
        let rts = now_ms();
        let rt = mint(&captured_nonce, rts, true);
        let lic = resp(Some(rt), Some(rts), Some(captured_nonce), true);
        let our_nonce = "b".repeat(64);
        assert!(!verify_response_token(KEY, &lic, &our_nonce, PID));
    }

    #[test]
    fn rejects_a_stale_token() {
        let nonce = "c".repeat(64);
        let rts = now_ms() - (RESPONSE_FRESHNESS_MS + 5_000); // outside the window
        let rt = mint(&nonce, rts, true);
        let lic = resp(Some(rt), Some(rts), Some(nonce.clone()), true);
        assert!(!verify_response_token(KEY, &lic, &nonce, PID));
    }

    #[test]
    fn rejects_a_flipped_valid_flag() {
        // Server signed valid=false; attacker flips the JSON `valid` to true but
        // can't recompute the HMAC without the key -> rejected.
        let nonce = "d".repeat(64);
        let rts = now_ms();
        let rt = mint(&nonce, rts, false); // signed as INVALID
        let lic = resp(Some(rt), Some(rts), Some(nonce.clone()), true); // body claims VALID
        assert!(!verify_response_token(KEY, &lic, &nonce, PID));
    }

    #[test]
    fn rejects_a_missing_token() {
        let nonce = "e".repeat(64);
        let lic = resp(None, None, None, true);
        assert!(!verify_response_token(KEY, &lic, &nonce, PID));
    }

    #[test]
    fn nonce_is_64_hex_chars_and_unique() {
        let a = generate_nonce();
        let b = generate_nonce();
        assert_eq!(a.len(), 64);
        assert!(a.chars().all(|c| c.is_ascii_hexdigit()));
        assert_ne!(a, b, "each request must get a distinct nonce");
    }

    /// Real `/activate` response body captured from production (2026-07-11) -
    /// note there is NO top-level `valid` field, only `success`. Before the
    /// fix, deserializing this directly left `lic.valid == false` even though
    /// the activation succeeded, so any real integration checking `resp.valid`
    /// after `activate()` always saw a false rejection. This test locks in
    /// the real wire shape so a future refactor can't silently reintroduce it.
    #[test]
    fn activate_response_shape_has_no_valid_field_only_success() {
        let raw = r#"{"success":true,"message":"Already activated","license":{"id":"984dac15-17c5-4399-a7ff-5338c6c60462","machineId":"7dd3d45a3b7d7ee39f9e8e0c99419c2e","expiresAt":"","features":[],"issuedAt":"","revokedAt":"","signature":"sig","userId":"admin","deviceLimit":1},"signature":"sig","rt":"rt","rts":1783768883093,"rn":"nonce"}"#;
        let lic: LicenseResponse =
            serde_json::from_str(raw).expect("must deserialize despite missing `valid`");
        // Deserializing alone (bypassing activate()'s success->valid copy) must
        // default to false, not panic or silently succeed - this documents why
        // activate() has to do the copy itself rather than relying on serde.
        assert!(!lic.valid);
        assert_eq!(lic.success, Some(true));
    }
}