bark-apns 0.2.0

use aes::{Aes128, Aes192, Aes256};
use aes_gcm::{
    Aes128Gcm, Aes256Gcm, AesGcm, Nonce,
    aead::{Aead, KeyInit, consts::U12},
};
use base64::{Engine, engine::general_purpose::STANDARD};
use cbc::Encryptor as CbcEncryptor;
use cipher::{BlockEncryptMut, KeyIvInit, block_padding::Pkcs7};
use ecb::Encryptor as EcbEncryptor;

use crate::error::{Error, Result};

/// AES algorithm selected in Bark's "Push Encryption" settings.
///
/// Bark currently exposes AES-128, AES-192, and AES-256. The selected variant
/// determines the required UTF-8 byte length of the shared key.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum EncryptionAlgorithm {
    /// AES-128, requiring a 16-byte key.
    AES128,
    /// AES-192, requiring a 24-byte key.
    AES192,
    /// AES-256, requiring a 32-byte key.
    AES256,
}

impl EncryptionAlgorithm {
    /// Returns the key length, in bytes, required by this algorithm.
    pub const fn key_len(self) -> usize {
        match self {
            Self::AES128 => 16,
            Self::AES192 => 24,
            Self::AES256 => 32,
        }
    }

    /// Returns the algorithm name used by Bark's client-side settings.
    pub const fn as_bark_str(self) -> &'static str {
        match self {
            Self::AES128 => "AES128",
            Self::AES192 => "AES192",
            Self::AES256 => "AES256",
        }
    }
}

/// AES block mode selected in Bark's "Push Encryption" settings.
///
/// Bark supports CBC, ECB, and GCM. CBC uses PKCS#7 padding, ECB uses PKCS#7
/// padding and no IV, and GCM uses CryptoSwift's combined mode, where the
/// authentication tag is appended to the ciphertext before Base64 encoding.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum EncryptionMode {
    /// AES-CBC with a 16-byte IV and PKCS#7 padding.
    CBC,
    /// AES-ECB with no IV and PKCS#7 padding.
    ECB,
    /// AES-GCM with a 12-byte nonce/IV and a combined authentication tag.
    GCM,
}

impl EncryptionMode {
    /// Returns the generated IV length required by this mode.
    ///
    /// ECB does not use an IV and returns `None`.
    pub const fn iv_len(self) -> Option<usize> {
        match self {
            Self::CBC => Some(16),
            Self::ECB => None,
            Self::GCM => Some(12),
        }
    }

    /// Returns the mode name used by Bark's client-side settings.
    pub const fn as_bark_str(self) -> &'static str {
        match self {
            Self::CBC => "CBC",
            Self::ECB => "ECB",
            Self::GCM => "GCM",
        }
    }
}

/// Bark iOS device target and its device-level settings.
///
/// Device tokens are normalized when a `Device` is created: surrounding angle
/// brackets and ASCII whitespace are removed. Encryption settings, when
/// present, should match the Push Encryption configuration in the Bark app for
/// this device.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Device {
    token: String,
    encryption_algorithm: Option<EncryptionAlgorithm>,
    encryption_mode: Option<EncryptionMode>,
    encryption_key: Option<String>,
}

impl Device {
    /// Creates a device target from an iOS device token.
    pub fn new<T>(token: T) -> Self
    where
        T: Into<String>,
    {
        Self {
            token: normalize_device_token(token.into()),
            encryption_algorithm: None,
            encryption_mode: None,
            encryption_key: None,
        }
    }

    /// Configures Bark Push Encryption for this device.
    ///
    /// The algorithm, mode, and key should match this device's Push Encryption
    /// settings in the Bark app. The key is validated immediately against the
    /// selected AES algorithm, so invalid device configuration fails before any
    /// APNs request is attempted.
    ///
    /// This method only stores the device-level encryption configuration. It
    /// does not force every message sent to this device to be encrypted; mark
    /// individual messages with [`Message::encrypt`] when encrypted delivery is
    /// desired.
    ///
    /// CBC and GCM IVs are generated for each encrypted APNs payload, not stored
    /// on the device. CBC uses a 16-byte IV, GCM uses a 12-byte nonce/IV, and
    /// ECB sends no IV.
    ///
    /// # Errors
    ///
    /// Returns [`Error::InvalidKeyLength`] if the key length does not match the
    /// selected algorithm.
    ///
    /// [`Error::InvalidKeyLength`]: crate::Error::InvalidKeyLength
    /// [`Message::encrypt`]: crate::Message::encrypt
    pub fn encrypt<K>(
        mut self,
        algorithm: EncryptionAlgorithm,
        mode: EncryptionMode,
        key: K,
    ) -> Result<Self>
    where
        K: Into<String>,
    {
        let key = key.into();
        let expected = algorithm.key_len();
        let actual = key.len();
        if actual != expected {
            return Err(Error::InvalidKeyLength {
                algorithm: algorithm.as_bark_str(),
                expected,
                actual,
            });
        }

        self.encryption_algorithm = Some(algorithm);
        self.encryption_mode = Some(mode);
        self.encryption_key = Some(key);
        Ok(self)
    }

    /// Returns the normalized iOS device token.
    pub fn token(&self) -> &str {
        &self.token
    }

    /// Returns whether this device has Bark Push Encryption configured.
    pub(crate) fn has_encryption(&self) -> bool {
        self.encryption().is_some()
    }

    pub(crate) fn encrypt_bark_json(&self, plaintext: &[u8]) -> Result<EncryptedPayload> {
        let (algorithm, mode, key) =
            self.encryption()
                .ok_or_else(|| Error::MissingDeviceEncryption {
                    device: self.token.clone(),
                })?;
        let iv = match mode.iv_len() {
            Some(len) => random_ascii_iv(len)?,
            None => String::new(),
        };

        let encrypted = match (algorithm, mode) {
            (EncryptionAlgorithm::AES128, EncryptionMode::CBC) => {
                CbcEncryptor::<Aes128>::new_from_slices(key.as_bytes(), iv.as_bytes())
                    .map_err(|_| Error::Encryption)?
                    .encrypt_padded_vec_mut::<Pkcs7>(plaintext)
            }
            (EncryptionAlgorithm::AES192, EncryptionMode::CBC) => {
                CbcEncryptor::<Aes192>::new_from_slices(key.as_bytes(), iv.as_bytes())
                    .map_err(|_| Error::Encryption)?
                    .encrypt_padded_vec_mut::<Pkcs7>(plaintext)
            }
            (EncryptionAlgorithm::AES256, EncryptionMode::CBC) => {
                CbcEncryptor::<Aes256>::new_from_slices(key.as_bytes(), iv.as_bytes())
                    .map_err(|_| Error::Encryption)?
                    .encrypt_padded_vec_mut::<Pkcs7>(plaintext)
            }
            (EncryptionAlgorithm::AES128, EncryptionMode::ECB) => {
                EcbEncryptor::<Aes128>::new_from_slice(key.as_bytes())
                    .map_err(|_| Error::Encryption)?
                    .encrypt_padded_vec_mut::<Pkcs7>(plaintext)
            }
            (EncryptionAlgorithm::AES192, EncryptionMode::ECB) => {
                EcbEncryptor::<Aes192>::new_from_slice(key.as_bytes())
                    .map_err(|_| Error::Encryption)?
                    .encrypt_padded_vec_mut::<Pkcs7>(plaintext)
            }
            (EncryptionAlgorithm::AES256, EncryptionMode::ECB) => {
                EcbEncryptor::<Aes256>::new_from_slice(key.as_bytes())
                    .map_err(|_| Error::Encryption)?
                    .encrypt_padded_vec_mut::<Pkcs7>(plaintext)
            }
            (EncryptionAlgorithm::AES128, EncryptionMode::GCM) => {
                let cipher =
                    Aes128Gcm::new_from_slice(key.as_bytes()).map_err(|_| Error::Encryption)?;
                cipher
                    .encrypt(Nonce::from_slice(iv.as_bytes()), plaintext)
                    .map_err(|_| Error::Encryption)?
            }
            (EncryptionAlgorithm::AES192, EncryptionMode::GCM) => {
                let cipher = AesGcm::<Aes192, U12>::new_from_slice(key.as_bytes())
                    .map_err(|_| Error::Encryption)?;
                cipher
                    .encrypt(Nonce::from_slice(iv.as_bytes()), plaintext)
                    .map_err(|_| Error::Encryption)?
            }
            (EncryptionAlgorithm::AES256, EncryptionMode::GCM) => {
                let cipher =
                    Aes256Gcm::new_from_slice(key.as_bytes()).map_err(|_| Error::Encryption)?;
                cipher
                    .encrypt(Nonce::from_slice(iv.as_bytes()), plaintext)
                    .map_err(|_| Error::Encryption)?
            }
        };

        Ok(EncryptedPayload {
            ciphertext: STANDARD.encode(encrypted),
            iv: mode.iv_len().map(|_| iv),
        })
    }

    fn encryption(&self) -> Option<(EncryptionAlgorithm, EncryptionMode, &str)> {
        Some((
            self.encryption_algorithm?,
            self.encryption_mode?,
            self.encryption_key.as_deref()?,
        ))
    }
}

/// Bark encrypted APNs payload fields generated for one device.
pub(crate) struct EncryptedPayload {
    /// Base64-encoded encrypted Bark request JSON.
    pub(crate) ciphertext: String,
    /// Per-payload IV sent through APNs for CBC and GCM.
    pub(crate) iv: Option<String>,
}

fn random_ascii_iv(len: usize) -> Result<String> {
    const ALPHABET: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_";

    let mut bytes = vec![0; len];
    getrandom::getrandom(&mut bytes)?;

    Ok(bytes
        .into_iter()
        .map(|byte| ALPHABET[usize::from(byte) % ALPHABET.len()] as char)
        .collect())
}

fn normalize_device_token(device: String) -> String {
    device
        .trim()
        .trim_start_matches('<')
        .trim_end_matches('>')
        .chars()
        .filter(|ch| !ch.is_ascii_whitespace())
        .collect()
}

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

    #[test]
    fn normalizes_device_tokens() {
        let device = Device::new("<aa bb>");

        assert_eq!(device.token(), "aabb");
    }

    #[test]
    fn keeps_device_encryption() {
        let device = Device::new("aabb")
            .encrypt(
                EncryptionAlgorithm::AES128,
                EncryptionMode::CBC,
                "1234567890123456",
            )
            .unwrap();

        assert_eq!(
            device.encryption_algorithm,
            Some(EncryptionAlgorithm::AES128)
        );
        assert_eq!(device.encryption_mode, Some(EncryptionMode::CBC));
        assert_eq!(device.encryption_key.as_deref(), Some("1234567890123456"));
        assert!(device.has_encryption());
    }

    #[test]
    fn validates_encryption_key_length() {
        let err = Device::new("aabb")
            .encrypt(EncryptionAlgorithm::AES128, EncryptionMode::CBC, "short")
            .unwrap_err();

        assert!(matches!(
            err,
            crate::Error::InvalidKeyLength {
                algorithm: "AES128",
                expected: 16,
                actual: 5
            }
        ));
    }

    #[test]
    fn generates_mode_specific_payload_iv() {
        let cbc = Device::new("aabb")
            .encrypt(
                EncryptionAlgorithm::AES128,
                EncryptionMode::CBC,
                "1234567890123456",
            )
            .unwrap();
        let gcm = Device::new("aabb")
            .encrypt(
                EncryptionAlgorithm::AES128,
                EncryptionMode::GCM,
                "1234567890123456",
            )
            .unwrap();
        let ecb = Device::new("aabb")
            .encrypt(
                EncryptionAlgorithm::AES128,
                EncryptionMode::ECB,
                "1234567890123456",
            )
            .unwrap();

        assert_eq!(
            cbc.encrypt_bark_json(b"test").unwrap().iv.unwrap().len(),
            16
        );
        assert_eq!(
            gcm.encrypt_bark_json(b"test").unwrap().iv.unwrap().len(),
            12
        );
        assert_eq!(ecb.encrypt_bark_json(b"test").unwrap().iv, None);
    }

    #[test]
    fn generates_fresh_iv_for_each_encrypted_payload() {
        let device = Device::new("aabb")
            .encrypt(
                EncryptionAlgorithm::AES128,
                EncryptionMode::CBC,
                "1234567890123456",
            )
            .unwrap();

        let first = device.encrypt_bark_json(b"test").unwrap();
        let second = device.encrypt_bark_json(b"test").unwrap();

        assert_ne!(first.iv, second.iv);
        assert_ne!(first.ciphertext, second.ciphertext);
    }
}