proofmode 0.9.0

pub mod types;

pub use types::{PlayIntegrityVerifyRequest, PlayIntegrityVerifyResponse};

use chrono::{Duration, Utc};
use jsonwebtoken::{encode, Algorithm, EncodingKey, Header};
use reqwest::Client;
use sha2::{Digest, Sha256};
use std::sync::Arc;
use tokio::sync::RwLock;

use super::error::{Result, SignError};
use types::*;

/// Trait for storing Play Integrity verification results and replay protection.
/// proofsign-rust implements this with PostgreSQL; proofmode only defines the interface.
#[async_trait::async_trait]
pub trait IntegrityStore: Send + Sync {
    /// Check if a token hash has been used before.
    async fn is_token_used(&self, token_hash: &str) -> Result<bool>;

    /// Store a verification result and token hash.
    async fn store_verification(
        &self,
        device_id: &str,
        package_name: &str,
        nonce: &str,
        token_hash: &str,
        verdict: &serde_json::Value,
        public_key: Option<&str>,
    ) -> Result<()>;

    /// Get the device's registered public key and current assertion counter.
    async fn get_device_signing_info(&self, device_id: &str) -> Result<Option<(String, i64)>>;

    /// Update the assertion counter for a device.
    async fn update_device_counter(&self, device_id: &str, new_counter: i64) -> Result<()>;
}

struct CachedToken {
    token: String,
    expires_at: chrono::DateTime<chrono::Utc>,
}

/// Play Integrity verifier that calls the Google Play Integrity API
/// to verify tokens and validate nonces/packages.
pub struct PlayIntegrityVerifier {
    client: Client,
    service_account: Arc<ServiceAccountKey>,
    cached_token: Arc<RwLock<Option<CachedToken>>>,
}

impl PlayIntegrityVerifier {
    /// Create a new verifier from a Google service account key JSON string or file path.
    pub fn new(service_account_key: &str) -> Result<Self> {
        let client = Client::new();

        let key_content = if service_account_key.trim().starts_with('{') {
            service_account_key.to_string()
        } else {
            std::fs::read_to_string(service_account_key).map_err(|e| {
                SignError::ConfigError(format!(
                    "Failed to read service account key from {}: {}",
                    service_account_key, e
                ))
            })?
        };

        let service_account: ServiceAccountKey =
            serde_json::from_str(&key_content).map_err(|e| {
                SignError::ConfigError(format!("Failed to parse service account key: {}", e))
            })?;

        Ok(Self {
            client,
            service_account: Arc::new(service_account),
            cached_token: Arc::new(RwLock::new(None)),
        })
    }

    /// Create from the GOOGLE_SERVICE_ACCOUNT_KEY environment variable.
    pub fn from_env() -> Result<Self> {
        let key_json = std::env::var("GOOGLE_SERVICE_ACCOUNT_KEY").map_err(|_| {
            SignError::ConfigError(
                "GOOGLE_SERVICE_ACCOUNT_KEY environment variable is required but not set"
                    .to_string(),
            )
        })?;

        Self::new(&key_json)
    }

    async fn get_access_token(&self) -> Result<String> {
        // Check cached token
        {
            let cache = self.cached_token.read().await;
            if let Some(ref cached) = *cache {
                if cached.expires_at > Utc::now() + Duration::minutes(5) {
                    return Ok(cached.token.clone());
                }
            }
        }

        // Generate new token
        let now = Utc::now();
        let claims = JwtClaims {
            iss: self.service_account.client_email.clone(),
            scope: "https://www.googleapis.com/auth/playintegrity".to_string(),
            aud: self.service_account.token_uri.clone(),
            iat: now.timestamp(),
            exp: (now + Duration::hours(1)).timestamp(),
        };

        let header = Header::new(Algorithm::RS256);
        let key = EncodingKey::from_rsa_pem(self.service_account.private_key.as_bytes())
            .map_err(|e| SignError::ConfigError(format!("Failed to parse private key: {}", e)))?;

        let jwt = encode(&header, &claims, &key)
            .map_err(|e| SignError::PlayIntegrity(format!("Failed to create JWT: {}", e)))?;

        let params = [
            ("grant_type", "urn:ietf:params:oauth:grant-type:jwt-bearer"),
            ("assertion", &jwt),
        ];

        let response = self
            .client
            .post(&self.service_account.token_uri)
            .form(&params)
            .send()
            .await
            .map_err(|e| SignError::PlayIntegrity(format!("Failed to get access token: {}", e)))?;

        if !response.status().is_success() {
            let status = response.status();
            let body = response.text().await.unwrap_or_default();
            return Err(SignError::PlayIntegrity(format!(
                "Token request failed with {}: {}",
                status, body
            )));
        }

        let token_response: GoogleTokenResponse = response.json().await.map_err(|e| {
            SignError::PlayIntegrity(format!("Failed to parse token response: {}", e))
        })?;

        // Cache the token
        let expires_at = Utc::now() + Duration::seconds(token_response.expires_in);
        {
            let mut cache = self.cached_token.write().await;
            *cache = Some(CachedToken {
                token: token_response.access_token.clone(),
                expires_at,
            });
        }

        Ok(token_response.access_token)
    }

    async fn decode_integrity_token(
        &self,
        token: &str,
        package_name: &str,
    ) -> Result<PlayIntegrityTokenPayload> {
        let access_token = self.get_access_token().await?;

        let url = format!(
            "https://playintegrity.googleapis.com/v1/{}:decodeIntegrityToken",
            package_name
        );

        let response = self
            .client
            .post(&url)
            .bearer_auth(&access_token)
            .json(&serde_json::json!({
                "integrityToken": token
            }))
            .send()
            .await
            .map_err(|e| {
                SignError::PlayIntegrity(format!("Failed to decode integrity token: {}", e))
            })?;

        if !response.status().is_success() {
            let status = response.status();
            let body = response.text().await.unwrap_or_default();
            return Err(SignError::PlayIntegrity(format!(
                "Play Integrity API returned {}: {}",
                status, body
            )));
        }

        let response_json: serde_json::Value = response.json().await.map_err(|e| {
            SignError::PlayIntegrity(format!("Failed to parse API response: {}", e))
        })?;

        let token_payload = response_json.get("tokenPayloadExternal").ok_or_else(|| {
            SignError::PlayIntegrity("Missing token payload in response".to_string())
        })?;

        let token_payload: PlayIntegrityTokenPayload =
            serde_json::from_value(token_payload.clone()).map_err(|e| {
                SignError::PlayIntegrity(format!("Failed to decode token payload: {}", e))
            })?;

        Ok(token_payload)
    }

    /// Verify a Play Integrity token against the Google API.
    /// Validates nonce, package name, and prevents token replay via the store.
    pub async fn verify_token(
        &self,
        request: &PlayIntegrityVerifyRequest,
        store: &dyn IntegrityStore,
    ) -> Result<PlayIntegrityVerifyResponse> {
        // 1. Hash the token for replay detection
        let mut hasher = Sha256::new();
        hasher.update(&request.integrity_token);
        let token_hash = hex::encode(hasher.finalize());

        // 2. Check replay
        if store.is_token_used(&token_hash).await? {
            return Err(SignError::TokenReplay(format!(
                "Token with hash {} has already been used",
                token_hash
            )));
        }

        // 3. Verify with Google API
        let token_payload = self
            .decode_integrity_token(&request.integrity_token, &request.package_name)
            .await?;

        // 4. Validate nonce
        if let Some(request_details) = &token_payload.request_details {
            if let Some(token_nonce) = &request_details.nonce {
                if !compare_nonces(token_nonce, &request.nonce) {
                    return Err(SignError::PlayIntegrity(
                        "Nonce in token does not match request nonce".to_string(),
                    ));
                }
            } else {
                return Err(SignError::PlayIntegrity(
                    "Missing nonce in token".to_string(),
                ));
            }
        } else {
            return Err(SignError::PlayIntegrity(
                "Missing request details in token".to_string(),
            ));
        }

        // 5. Validate package name
        if let Some(app_integrity) = &token_payload.app_integrity {
            if let Some(pkg_name) = &app_integrity.package_name {
                if pkg_name != &request.package_name {
                    return Err(SignError::PlayIntegrity(format!(
                        "Package name in token ({}) does not match request ({})",
                        pkg_name, request.package_name
                    )));
                }
            } else {
                return Err(SignError::PlayIntegrity(
                    "Missing package name in token".to_string(),
                ));
            }
        } else {
            return Err(SignError::PlayIntegrity(
                "Missing app integrity in token".to_string(),
            ));
        }

        // 6. Store verification result
        let verdict_json = serde_json::to_value(&token_payload)
            .map_err(|e| SignError::PlayIntegrity(format!("Failed to serialize verdict: {}", e)))?;

        store
            .store_verification(
                &request.device_id,
                &request.package_name,
                &request.nonce,
                &token_hash,
                &verdict_json,
                request.public_key.as_deref(),
            )
            .await?;

        Ok(PlayIntegrityVerifyResponse {
            success: true,
            message: "Play Integrity token verified successfully".to_string(),
            device_integrity: serde_json::to_value(&token_payload.device_integrity).ok(),
            app_integrity: serde_json::to_value(&token_payload.app_integrity).ok(),
        })
    }
}

/// Compare two nonces that may be in different base64 encodings
/// (standard vs URL-safe). Decodes both and compares raw bytes.
pub fn compare_nonces(nonce1: &str, nonce2: &str) -> bool {
    use base64::{engine::general_purpose, Engine};

    let bytes1 = general_purpose::STANDARD
        .decode(nonce1)
        .or_else(|_| general_purpose::URL_SAFE_NO_PAD.decode(nonce1))
        .or_else(|_| general_purpose::URL_SAFE.decode(nonce1));

    let bytes2 = general_purpose::STANDARD
        .decode(nonce2)
        .or_else(|_| general_purpose::URL_SAFE_NO_PAD.decode(nonce2))
        .or_else(|_| general_purpose::URL_SAFE.decode(nonce2));

    match (bytes1, bytes2) {
        (Ok(b1), Ok(b2)) => b1 == b2,
        _ => nonce1 == nonce2,
    }
}

/// Verify an Android device's cryptographic request signature.
/// This verifies that a request was signed by the device's registered private key.
pub fn verify_android_device_signature(
    public_key_base64: &str,
    device_id: &str,
    counter: i64,
    timestamp: i64,
    claim: &str,
    request_signature_base64: &str,
) -> Result<()> {
    use base64::{engine::general_purpose::STANDARD, Engine};
    use p256::ecdsa::{signature::Verifier, Signature, VerifyingKey};
    use p256::pkcs8::DecodePublicKey;

    // Verify timestamp is recent (within 5 minutes)
    let now_ms = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap()
        .as_millis() as i64;
    let five_minutes_ms: i64 = 5 * 60 * 1000;
    if (now_ms - timestamp).abs() > five_minutes_ms {
        return Err(SignError::PlayIntegrity(
            "Request timestamp too old".to_string(),
        ));
    }

    // Reconstruct the data that was signed: "deviceId|counter|timestamp|claim"
    let data_to_verify = format!("{}|{}|{}|{}", device_id, counter, timestamp, claim);

    // Decode the public key (SPKI format from Android Keystore)
    let public_key_der = STANDARD
        .decode(public_key_base64)
        .map_err(|e| SignError::InvalidKey(format!("Failed to decode public key: {}", e)))?;

    let verifying_key = VerifyingKey::from_public_key_der(&public_key_der)
        .map_err(|e| SignError::InvalidKey(format!("Failed to parse public key: {}", e)))?;

    // Decode the signature
    let signature_bytes = STANDARD
        .decode(request_signature_base64)
        .map_err(|e| SignError::InvalidSignature(format!("Invalid signature encoding: {}", e)))?;

    // Parse as DER-encoded ECDSA signature (Android uses DER format)
    let signature = Signature::from_der(&signature_bytes)
        .map_err(|e| SignError::InvalidSignature(format!("Invalid signature format: {}", e)))?;

    // Verify the signature
    verifying_key
        .verify(data_to_verify.as_bytes(), &signature)
        .map_err(|e| {
            SignError::VerificationFailed(format!("Signature verification failed: {}", e))
        })?;

    Ok(())
}

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

    #[test]
    fn test_compare_nonces_same() {
        assert!(compare_nonces("aGVsbG8=", "aGVsbG8="));
    }

    #[test]
    fn test_compare_nonces_different_encoding() {
        // Standard base64 vs URL-safe
        use base64::{engine::general_purpose, Engine};
        let data = vec![0xFF, 0xFE, 0xFD];
        let standard = general_purpose::STANDARD.encode(&data);
        let url_safe = general_purpose::URL_SAFE.encode(&data);
        assert!(compare_nonces(&standard, &url_safe));
    }

    #[test]
    fn test_compare_nonces_different() {
        assert!(!compare_nonces("aGVsbG8=", "d29ybGQ="));
    }

    #[test]
    fn test_compare_nonces_url_safe_no_pad() {
        use base64::{engine::general_purpose, Engine};
        let data = vec![0xFF, 0xFE, 0xFD, 0xFC];
        let standard = general_purpose::STANDARD.encode(&data);
        let no_pad = general_purpose::URL_SAFE_NO_PAD.encode(&data);
        assert!(compare_nonces(&standard, &no_pad));
    }

    #[test]
    fn test_compare_nonces_plain_string_match() {
        // Non-base64 strings fall back to string comparison
        assert!(compare_nonces("plain", "plain"));
        assert!(!compare_nonces("plain", "other"));
    }

    fn test_signing_key() -> (p256::ecdsa::SigningKey, String) {
        use base64::{engine::general_purpose::STANDARD, Engine};
        use p256::ecdsa::SigningKey;
        use p256::pkcs8::EncodePublicKey;

        // Deterministic test key (any valid 32-byte scalar)
        let secret = [
            0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e,
            0x0f, 0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18, 0x19, 0x1a, 0x1b, 0x1c,
            0x1d, 0x1e, 0x1f, 0x20,
        ];
        let signing_key = SigningKey::from_slice(&secret).unwrap();
        let verifying_key = signing_key.verifying_key();
        let public_key_der = verifying_key.to_public_key_der().unwrap();
        let public_key_base64 = STANDARD.encode(public_key_der.as_bytes());
        (signing_key, public_key_base64)
    }

    #[test]
    fn test_verify_android_device_signature_valid() {
        use base64::{engine::general_purpose::STANDARD, Engine};
        use p256::ecdsa::{signature::Signer, Signature};

        let (signing_key, public_key_base64) = test_signing_key();

        let device_id = "test-device-001";
        let counter: i64 = 42;
        let now_ms = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_millis() as i64;
        let claim = "test-claim-hash";

        let data_to_sign = format!("{}|{}|{}|{}", device_id, counter, now_ms, claim);
        let signature: Signature = signing_key.sign(data_to_sign.as_bytes());
        let signature_base64 = STANDARD.encode(signature.to_der());

        let result = verify_android_device_signature(
            &public_key_base64,
            device_id,
            counter,
            now_ms,
            claim,
            &signature_base64,
        );
        assert!(result.is_ok());
    }

    #[test]
    fn test_verify_android_device_signature_wrong_claim() {
        use base64::{engine::general_purpose::STANDARD, Engine};
        use p256::ecdsa::{signature::Signer, Signature};

        let (signing_key, public_key_base64) = test_signing_key();

        let device_id = "test-device-001";
        let counter: i64 = 42;
        let now_ms = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_millis() as i64;

        let data_to_sign = format!("{}|{}|{}|{}", device_id, counter, now_ms, "original-claim");
        let signature: Signature = signing_key.sign(data_to_sign.as_bytes());
        let signature_base64 = STANDARD.encode(signature.to_der());

        let result = verify_android_device_signature(
            &public_key_base64,
            device_id,
            counter,
            now_ms,
            "tampered-claim",
            &signature_base64,
        );
        assert!(result.is_err());
    }

    #[test]
    fn test_verify_android_device_signature_expired_timestamp() {
        use base64::{engine::general_purpose::STANDARD, Engine};
        use p256::ecdsa::{signature::Signer, Signature};

        let (signing_key, public_key_base64) = test_signing_key();

        let device_id = "test-device-001";
        let counter: i64 = 42;
        // 10 minutes ago -- outside the 5-minute window
        let old_timestamp = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_millis() as i64
            - 10 * 60 * 1000;
        let claim = "test-claim";

        let data_to_sign = format!("{}|{}|{}|{}", device_id, counter, old_timestamp, claim);
        let signature: Signature = signing_key.sign(data_to_sign.as_bytes());
        let signature_base64 = STANDARD.encode(signature.to_der());

        let result = verify_android_device_signature(
            &public_key_base64,
            device_id,
            counter,
            old_timestamp,
            claim,
            &signature_base64,
        );
        assert!(result.is_err());
        let err = format!("{}", result.unwrap_err());
        assert!(err.contains("timestamp"));
    }

    #[test]
    fn test_verify_android_device_signature_invalid_key() {
        use base64::{engine::general_purpose::STANDARD, Engine};

        let now_ms = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_millis() as i64;

        let result = verify_android_device_signature(
            &STANDARD.encode(b"not-a-real-public-key"),
            "device",
            1,
            now_ms,
            "claim",
            &STANDARD.encode(b"not-a-real-signature"),
        );
        assert!(result.is_err());
    }
}