// Web Push (RFC 8030 / 8291 / 8292): the phone (an installed home-screen page) gets a notification
// when a bot needs the person or finishes while they're away. The payload is encrypted for the
// phone (aes128gcm); the push service only sees ciphertext.

use base64::Engine;
use ring::rand::SecureRandom;
use ring::{aead, agreement, hkdf, signature};
use serde_json::{json, Value};
use std::path::Path;

fn b64u(b: &[u8]) -> String {
    base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(b)
}

fn unb64u(s: &str) -> Option<Vec<u8>> {
    base64::engine::general_purpose::URL_SAFE_NO_PAD.decode(s.trim_end_matches('=')).ok()
}

struct Len(usize);
impl hkdf::KeyType for Len {
    fn len(&self) -> usize {
        self.0
    }
}

fn hkdf(salt: &[u8], ikm: &[u8], info: &[u8], len: usize) -> Vec<u8> {
    let prk = hkdf::Salt::new(hkdf::HKDF_SHA256, salt).extract(ikm);
    let mut out = vec![0u8; len];
    prk.expand(&[info], Len(len)).expect("hkdf").fill(&mut out).expect("hkdf");
    out
}

/// The server's signing key (made once, kept in the data folder).
pub struct Vapid {
    pair: signature::EcdsaKeyPair,
    pub public_b64: String,
}

impl Vapid {
    pub fn load(dir: &Path) -> Vapid {
        let rng = ring::rand::SystemRandom::new();
        let path = dir.join("vapid.pk8");
        let pk8 = std::fs::read(&path).ok().unwrap_or_else(|| {
            let doc = signature::EcdsaKeyPair::generate_pkcs8(&signature::ECDSA_P256_SHA256_FIXED_SIGNING, &rng).expect("key");
            let _ = std::fs::write(&path, doc.as_ref());
            doc.as_ref().to_vec()
        });
        let pair = signature::EcdsaKeyPair::from_pkcs8(&signature::ECDSA_P256_SHA256_FIXED_SIGNING, &pk8, &rng).expect("vapid key");
        let public_b64 = b64u(signature::KeyPair::public_key(&pair).as_ref());
        Vapid { pair, public_b64 }
    }

    /// `Authorization: vapid t=…, k=…` for one push service.
    fn header(&self, endpoint: &str, contact: &str) -> String {
        let aud = endpoint.splitn(4, '/').take(3).collect::<Vec<_>>().join("/");
        let exp = chrono::Utc::now().timestamp() + 12 * 3600;
        let head = b64u(br#"{"typ":"JWT","alg":"ES256"}"#);
        let claims = b64u(json!({"aud": aud, "exp": exp, "sub": contact}).to_string().as_bytes());
        let msg = format!("{head}.{claims}");
        let sig = self.pair.sign(&ring::rand::SystemRandom::new(), msg.as_bytes()).expect("sign");
        format!("vapid t={msg}.{}, k={}", b64u(sig.as_ref()), self.public_b64)
    }
}

/// Encrypt `payload` for one subscription (RFC 8291, one record).
pub fn encrypt(p256dh: &str, auth: &str, payload: &[u8]) -> Option<Vec<u8>> {
    let rng = ring::rand::SystemRandom::new();
    let ua_public = unb64u(p256dh)?;
    let auth_secret = unb64u(auth)?;
    let eph = agreement::EphemeralPrivateKey::generate(&agreement::ECDH_P256, &rng).ok()?;
    let as_public = eph.compute_public_key().ok()?.as_ref().to_vec();
    let ecdh = agreement::agree_ephemeral(eph, &agreement::UnparsedPublicKey::new(&agreement::ECDH_P256, &ua_public), |s| s.to_vec()).ok()?;
    let mut key_info = b"WebPush: info\0".to_vec();
    key_info.extend(&ua_public);
    key_info.extend(&as_public);
    let ikm = hkdf(&auth_secret, &ecdh, &key_info, 32);
    let mut salt = [0u8; 16];
    rng.fill(&mut salt).ok()?;
    let cek = hkdf(&salt, &ikm, b"Content-Encoding: aes128gcm\0", 16);
    let nonce = hkdf(&salt, &ikm, b"Content-Encoding: nonce\0", 12);
    let key = aead::LessSafeKey::new(aead::UnboundKey::new(&aead::AES_128_GCM, &cek).ok()?);
    let mut data = payload.to_vec();
    data.push(2); // last record
    key.seal_in_place_append_tag(aead::Nonce::try_assume_unique_for_key(&nonce).ok()?, aead::Aad::empty(), &mut data).ok()?;
    let mut body = salt.to_vec();
    body.extend(4096u32.to_be_bytes());
    body.push(as_public.len() as u8);
    body.extend(&as_public);
    body.extend(data);
    Some(body)
}

/// Send one notification. Err(true) = the subscription is gone (forget it).
pub fn send(v: &Vapid, sub: &Value, payload: &Value, contact: &str) -> Result<(), bool> {
    let endpoint = sub["endpoint"].as_str().ok_or(true)?;
    if !endpoint.starts_with("https://") {
        return Err(true);
    }
    let body = encrypt(sub["keys"]["p256dh"].as_str().unwrap_or(""), sub["keys"]["auth"].as_str().unwrap_or(""), payload.to_string().as_bytes()).ok_or(true)?;
    let agent: ureq::Agent = ureq::Agent::config_builder().timeout_global(Some(std::time::Duration::from_secs(15))).http_status_as_error(false).build().into();
    let r = agent
        .post(endpoint)
        .header("Authorization", &v.header(endpoint, contact))
        .header("Content-Encoding", "aes128gcm")
        .header("Content-Type", "application/octet-stream")
        .header("TTL", "86400")
        .header("Urgency", "high")
        .send(&body[..]);
    match r {
        Ok(resp) => match resp.status().as_u16() {
            200..=299 => Ok(()),
            404 | 410 => Err(true),
            _ => Err(false),
        },
        Err(_) => Err(false),
    }
}

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

    /// Decrypt as a browser would (RFC 8291), with the browser's private key.
    #[test]
    fn payload_decrypts_on_the_phone_side() {
        let rng = ring::rand::SystemRandom::new();
        let ua = agreement::EphemeralPrivateKey::generate(&agreement::ECDH_P256, &rng).unwrap();
        let ua_pub = ua.compute_public_key().unwrap().as_ref().to_vec();
        let mut auth = [0u8; 16];
        rng.fill(&mut auth).unwrap();
        let body = encrypt(&b64u(&ua_pub), &b64u(&auth), "需要你允許".as_bytes()).unwrap();
        let (salt, rest) = body.split_at(16);
        let idlen = rest[4] as usize;
        let as_pub = &rest[5..5 + idlen];
        let ct = &rest[5 + idlen..];
        let ecdh = agreement::agree_ephemeral(ua, &agreement::UnparsedPublicKey::new(&agreement::ECDH_P256, as_pub), |s| s.to_vec()).unwrap();
        let mut info = b"WebPush: info\0".to_vec();
        info.extend(&ua_pub);
        info.extend(as_pub);
        let ikm = hkdf(&auth, &ecdh, &info, 32);
        let cek = hkdf(salt, &ikm, b"Content-Encoding: aes128gcm\0", 16);
        let nonce = hkdf(salt, &ikm, b"Content-Encoding: nonce\0", 12);
        let key = aead::LessSafeKey::new(aead::UnboundKey::new(&aead::AES_128_GCM, &cek).unwrap());
        let mut data = ct.to_vec();
        let plain = key.open_in_place(aead::Nonce::try_assume_unique_for_key(&nonce).unwrap(), aead::Aad::empty(), &mut data).unwrap();
        assert_eq!(&plain[..plain.len() - 1], "需要你允許".as_bytes());
        assert_eq!(*plain.last().unwrap(), 2);
    }

    #[test]
    fn vapid_token_verifies() {
        let dir = std::env::temp_dir().join(format!("anybot-vapid-{}", std::process::id()));
        std::fs::create_dir_all(&dir).unwrap();
        let v = Vapid::load(&dir);
        let h = v.header("https://web.push.apple.com/abc/def", "mailto:x@example.com");
        let t = h.strip_prefix("vapid t=").unwrap().split(", k=").next().unwrap();
        let (msg, sig) = t.rsplit_once('.').unwrap();
        let pubkey = unb64u(&v.public_b64).unwrap();
        signature::UnparsedPublicKey::new(&signature::ECDSA_P256_SHA256_FIXED, &pubkey).verify(msg.as_bytes(), &unb64u(sig).unwrap()).unwrap();
        let claims: Value = serde_json::from_slice(&unb64u(msg.split('.').nth(1).unwrap()).unwrap()).unwrap();
        assert_eq!(claims["aud"], "https://web.push.apple.com");
        // Same key after a restart.
        assert_eq!(Vapid::load(&dir).public_b64, v.public_b64);
        let _ = std::fs::remove_dir_all(&dir);
    }
}