bark-apns 0.1.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 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",
        }
    }
}

/// Encryption settings for a Bark encrypted push.
///
/// The algorithm, mode, key, and IV mirror Bark's `CryptoSettingFields`.
/// Create this value with [`Encryption::new`] to generate a fresh IV for CBC or
/// GCM automatically, or with [`Encryption::with_iv`] when an exact IV is needed
/// for tests or interoperability checks.
///
/// The IV is not secret. For CBC and GCM this crate serializes it as the
/// top-level APNs `iv` field so Bark's notification service extension can use it
/// instead of the IV stored in the app settings.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Encryption {
    algorithm: EncryptionAlgorithm,
    mode: EncryptionMode,
    key: String,
    iv: String,
}

impl Encryption {
    /// Creates encryption settings and generates a mode-appropriate IV.
    ///
    /// CBC generates a 16-byte ASCII IV, GCM generates a 12-byte ASCII IV, and
    /// ECB stores no IV. The key is validated immediately against the selected
    /// algorithm's required byte length.
    ///
    /// # Errors
    ///
    /// Returns [`Error::InvalidKeyLength`] if the key length does not match the
    /// algorithm, or [`Error::Random`] if the operating system cannot provide
    /// randomness for an IV.
    ///
    /// [`Error::InvalidKeyLength`]: crate::Error::InvalidKeyLength
    /// [`Error::Random`]: crate::Error::Random
    pub fn new<K>(algorithm: EncryptionAlgorithm, mode: EncryptionMode, key: K) -> Result<Self>
    where
        K: Into<String>,
    {
        let iv = match mode.iv_len() {
            Some(len) => random_ascii_iv(len)?,
            None => String::new(),
        };

        Self::with_iv(algorithm, mode, key, iv)
    }

    /// Creates encryption settings with an explicit IV.
    ///
    /// Use this when reproducing Bark documentation examples or when the caller
    /// deliberately wants to control the IV. Normal sends should prefer
    /// [`Encryption::new`] so CBC and GCM do not reuse IVs.
    ///
    /// For ECB, pass an empty string because Bark's ECB mode does not use an IV.
    ///
    /// # Errors
    ///
    /// Returns [`Error::InvalidKeyLength`] when the key has the wrong length, or
    /// [`Error::InvalidIvLength`] when the IV length does not match the selected
    /// mode.
    ///
    /// [`Error::InvalidKeyLength`]: crate::Error::InvalidKeyLength
    /// [`Error::InvalidIvLength`]: crate::Error::InvalidIvLength
    pub fn with_iv<K, I>(
        algorithm: EncryptionAlgorithm,
        mode: EncryptionMode,
        key: K,
        iv: I,
    ) -> Result<Self>
    where
        K: Into<String>,
        I: Into<String>,
    {
        let key = key.into();
        let iv = iv.into();
        let expected_key_len = algorithm.key_len();
        let actual_key_len = key.len();
        if actual_key_len != expected_key_len {
            return Err(Error::InvalidKeyLength {
                algorithm: algorithm.as_bark_str(),
                expected: expected_key_len,
                actual: actual_key_len,
            });
        }

        let expected_iv_len = mode.iv_len().unwrap_or(0);
        let actual_iv_len = iv.len();
        if actual_iv_len != expected_iv_len {
            return Err(Error::InvalidIvLength {
                mode: mode.as_bark_str(),
                expected: expected_iv_len,
                actual: actual_iv_len,
            });
        }

        Ok(Self {
            algorithm,
            mode,
            key,
            iv,
        })
    }

    /// Returns the selected Bark encryption algorithm.
    pub const fn algorithm(&self) -> EncryptionAlgorithm {
        self.algorithm
    }

    /// Returns the selected Bark encryption mode.
    pub const fn mode(&self) -> EncryptionMode {
        self.mode
    }

    /// Returns the shared encryption key.
    ///
    /// This is the exact UTF-8 key string that must also be configured on the
    /// iOS device in Bark's push encryption settings.
    pub fn key(&self) -> &str {
        &self.key
    }

    /// Returns the IV that will be sent with the APNs payload.
    ///
    /// CBC and GCM return `Some`, while ECB returns `None` because it has no IV.
    pub fn iv(&self) -> Option<&str> {
        self.mode.iv_len().map(|_| self.iv.as_str())
    }

    pub(crate) fn apns_iv(&self) -> Option<&str> {
        self.iv()
    }

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

        Ok(STANDARD.encode(encrypted))
    }
}

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())
}

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

    #[test]
    fn validates_key_length() {
        let err =
            Encryption::new(EncryptionAlgorithm::AES128, EncryptionMode::CBC, "short").unwrap_err();

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

    #[test]
    fn validates_cbc_iv_length() {
        let err = Encryption::with_iv(
            EncryptionAlgorithm::AES128,
            EncryptionMode::CBC,
            "1234567890123456",
            "short",
        )
        .unwrap_err();

        assert!(matches!(
            err,
            Error::InvalidIvLength {
                mode: "CBC",
                expected: 16,
                actual: 5
            }
        ));
    }

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

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

    #[test]
    fn encrypts_like_bark_docs_cbc_example() {
        let encryption = Encryption::with_iv(
            EncryptionAlgorithm::AES128,
            EncryptionMode::CBC,
            "1234567890123456",
            "1111111111111111",
        )
        .unwrap();

        let ciphertext = encryption
            .encrypt_bark_json(br#"{"body": "test", "sound": "birdsong"}"#)
            .unwrap();

        assert_eq!(
            ciphertext,
            "d3QhjQjP5majvNt5CjsvFWwqqj2gKl96RFj5OO+u6ynTt7lkyigDYNA3abnnCLpr"
        );
    }
}