//! Milenage authentication algorithm
//! as proposed by ETSI SAGE for 3G authentication.
//!
//! See 3GPP TS [35.205 (General)](https://www.3gpp.org/ftp/Specs/archive/35_series/35.205/),
//! [3GPP TS 35.206 (Algorithm specification)](https://www.3gpp.org/ftp/Specs/archive/35_series/35.206/)
//! and [3GPP TS 35.208 (Design conformance test data)](https://www.3gpp.org/ftp/Specs/archive/35_series/35.208/).
//!
//!
//! # Usage example
//! ```
//!use hex_literal::hex;
//!use milenage::Milenage;
//!
//!fn main() {
//!        // Use Test set 2 from 3GPP 35.208
//!        let k = hex!("465b5ce8b199b49faa5f0a2ee238a6bc");
//!        let op = hex!("cdc202d5123e20f62b6d676ac72cb318");
//!        let rand = hex!("23553cbe9637a89d218ae64dae47bf35");
//!
//!        let mut m = Milenage::new_with_op(k, op);
//!        let (res, ck, ik, ak) = m.f2345(&rand);
//!
//!        assert_eq!(m.res(), Some(&hex!("a54211d5e3ba50bf")));
//!        // or
//!        assert_eq!(res, hex!("a54211d5e3ba50bf"));
//!        assert_eq!(ck, hex!("b40ba9a3c58b2a05bbf0d987b21bf8cb"));
//!        assert_eq!(ik, hex!("f769bcd751044604127672711c6d3441"));
//!        assert_eq!(ak, hex!("aa689c648370"));
//!}
//!
//! ```

#[cfg(all(feature = "aes", feature = "openssl"))]
compile_error!("feature \"aes\" and feature \"openssl\" cannot be enabled at the same time");

#[cfg(feature = "aes")]
use aes::Aes128;

#[cfg(feature = "openssl")]
use {
    openssl::aes::{aes_ige, AesKey},
    openssl::symm::Mode,
};

use hmac::{Hmac, KeyInit as _, Mac};
use sha2::Sha256;
use std::fmt;

/// Errors that can occur during Milenage operations.
#[derive(Debug, PartialEq)]
pub enum MilenageError {
    /// CK is missing, f2345() must be called before compute_res_star().
    MissingCk,
    /// IK is missing, f2345() must be called before compute_res_star().
    MissingIk,
    /// The provided MNC is invalid (must be 2 or 3 digits).
    InvalidMnc(String),
    /// The provided MCC is invalid (must be 3 digits).
    InvalidMcc(String),
}

impl fmt::Display for MilenageError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            MilenageError::MissingCk => {
                write!(f, "missing CK, run f2345() before compute_res_star()")
            }
            MilenageError::MissingIk => {
                write!(f, "missing IK, run f2345() before compute_res_star()")
            }
            MilenageError::InvalidMnc(mnc) => write!(f, "invalid MNC: {}", mnc),
            MilenageError::InvalidMcc(mcc) => write!(f, "invalid MCC: {}", mcc),
        }
    }
}

impl std::error::Error for MilenageError {}

/// xor two 16 bytes array
fn xor(a1: &[u8; 16], a2: &[u8; 16]) -> [u8; 16] {
    let mut output = [0u8; 16];
    for i in 0..16 {
        output[i] = a1[i] ^ a2[i];
    }
    output
}

/// Milenage instance
#[derive(Debug, Default)]
pub struct Milenage {
    /// AK is a 48-bit anonymity key that is the output of either of the functions f5.
    ak: Option<[u8; 6]>,
    /// CK is a 128-bit confidentiality key that is the output of the function f3.
    ck: Option<[u8; 16]>,
    /// IK is a 128-bit integrity key that is the output of the function f4.
    ik: Option<[u8; 16]>,
    /// K is a 128-bit subscriber key that is an input to the functions f1, f1*, f2, f3, f4, f5 and f5*.
    k: [u8; 16],
    /// MACA is a 64-bit network authentication code that is the output of the function f1.
    maca: Option<[u8; 8]>,
    /// MACS is a 64-bit resynchronisation authentication code that is the output of the function f1*.
    macs: Option<[u8; 8]>,
    /// OPc is a 128-bit value derived from OP and K and used within the computation of the functions.
    opc: [u8; 16],
    /// RES is a 64-bit signed response that is the output of the function f2.
    res: Option<[u8; 8]>,
    /// RES* is a 128-bit response that is used in 5G.
    res_star: Option<[u8; 16]>,
}

impl Milenage {
    ///  Returns a new initialized Milenage from K and OP
    pub fn new_with_op(k: [u8; 16], op: [u8; 16]) -> Milenage {
        let mut m = Milenage {
            k,
            ..Default::default()
        };
        m.compute_opc(op);
        m
    }

    ///  Returns a new initialized Milenage from K and OPc
    pub fn new_with_opc(k: [u8; 16], opc: [u8; 16]) -> Milenage {
        Milenage {
            k,
            opc,
            ..Default::default()
        }
    }

    /// Returns the 48-bit anonymity key (output of f5 or f5*), if computed.
    pub fn ak(&self) -> Option<&[u8; 6]> {
        self.ak.as_ref()
    }

    /// Returns the 128-bit confidentiality key (output of f3), if computed.
    pub fn ck(&self) -> Option<&[u8; 16]> {
        self.ck.as_ref()
    }

    /// Returns the 128-bit integrity key (output of f4), if computed.
    pub fn ik(&self) -> Option<&[u8; 16]> {
        self.ik.as_ref()
    }

    /// Returns the 64-bit network authentication code (output of f1), if computed.
    pub fn maca(&self) -> Option<&[u8; 8]> {
        self.maca.as_ref()
    }

    /// Returns the 64-bit resynchronisation authentication code (output of f1*), if computed.
    pub fn macs(&self) -> Option<&[u8; 8]> {
        self.macs.as_ref()
    }

    /// Returns the 128-bit OPc value derived from OP and K.
    pub fn opc(&self) -> &[u8; 16] {
        &self.opc
    }

    /// Returns the 64-bit signed response (output of f2), if computed.
    pub fn res(&self) -> Option<&[u8; 8]> {
        self.res.as_ref()
    }

    /// Returns the 128-bit 5G response (RES*), if computed.
    pub fn res_star(&self) -> Option<&[u8; 16]> {
        self.res_star.as_ref()
    }

    /// F1 is the network authentication function.
    /// F1 computes network authentication code MAC-A from key K, random challenge RAND,
    /// sequence number SQN and authentication management field AMF.
    pub fn f1(&mut self, rand: &[u8; 16], sqn: &[u8; 6], amf: &[u8; 2]) -> [u8; 8] {
        let mac = self.f1base(rand, sqn, amf);
        let mut maca = [0u8; 8];
        maca.copy_from_slice(&mac[..8]);

        self.maca = Some(maca);
        maca
    }

    /// F1Star is the re-synchronisation message authentication function.
    /// F1Star computes resynch authentication code MAC-S from key K, random challenge RAND,
    /// sequence number SQN and authentication management field AMF.
    pub fn f1star(&mut self, rand: &[u8; 16], sqn: &[u8; 6], amf: &[u8; 2]) -> [u8; 8] {
        let mac = self.f1base(rand, sqn, amf);
        let mut macs = [0u8; 8];
        macs.copy_from_slice(&mac[8..]);

        self.macs = Some(macs);
        macs
    }

    /// Used by f1 and f1star
    fn f1base(&self, rand: &[u8; 16], sqn: &[u8; 6], amf: &[u8; 2]) -> [u8; 16] {
        let rijndael_input: [u8; 16] = xor(&self.opc, rand);
        let temp = self.rijndael_encrypt(&rijndael_input);

        let mut in1 = [0u8; 16];
        in1[..6].copy_from_slice(sqn);
        in1[6..8].copy_from_slice(amf);
        in1[8..14].copy_from_slice(sqn);
        in1[14..16].copy_from_slice(amf);

        let mut rijndael_input = [0u8; 16];

        /* XOR op_c and in1, rotate by r1=64, and XOR *
         * on the constant c1 (which is all zeroes)   */

        for i in 0..16 {
            rijndael_input[(i + 8) % 16] = in1[i] ^ self.opc[i];
        }

        /* XOR on the value temp computed before */
        for (i, elem) in rijndael_input.iter_mut().enumerate() {
            *elem ^= temp[i];
        }

        let mut out1 = self.rijndael_encrypt(&rijndael_input);

        for (i, elem) in out1.iter_mut().enumerate() {
            *elem ^= &self.opc[i];
        }

        out1
    }

    /// F2345 takes key K and random challenge RAND, and returns response RES,
    /// confidentiality key CK, integrity key IK and anonymity key AK.
    pub fn f2345(&mut self, rand: &[u8; 16]) -> ([u8; 8], [u8; 16], [u8; 16], [u8; 6]) {
        let rijndael_input = xor(&self.opc, rand);
        let temp = self.rijndael_encrypt(&rijndael_input);

        // To obtain output block OUT2: XOR OPc and TEMP, rotate by r2=0, and XOR on the
        // constant c2 (which is all zeroes except that the last bit is 1).

        let mut rijndael_input = xor(&temp, &self.opc);
        rijndael_input[15] ^= 1;

        let out = self.rijndael_encrypt(&rijndael_input);
        let tmp = xor(&out, &self.opc);

        let mut res = [0u8; 8];
        let mut ak = [0u8; 6];

        res.copy_from_slice(&tmp[8..]);
        ak.copy_from_slice(&tmp[..6]);

        // To obtain output block OUT3: XOR OPc and TEMP, rotate by r3=32, and XOR on the
        // constant c3 (which is all zeroes except that the next to last bit is 1).

        let mut rijndael_input = [0u8; 16];
        for i in 0..16 {
            rijndael_input[(i + 12) % 16] = temp[i] ^ self.opc[i]
        }
        rijndael_input[15] ^= 2;

        let out = self.rijndael_encrypt(&rijndael_input);
        let ck = xor(&out, &self.opc);

        // To obtain output block OUT4: XOR OPc and TEMP, rotate by r4=64, and XOR on the
        // constant c4 (which is all zeroes except that the 2nd from last bit is 1).

        let mut rijndael_input = [0u8; 16];
        for i in 0..16 {
            rijndael_input[(i + 8) % 16] = temp[i] ^ self.opc[i]
        }
        rijndael_input[15] ^= 4;

        let out = self.rijndael_encrypt(&rijndael_input);
        let ik = xor(&out, &self.opc);

        self.res = Some(res);
        self.ck = Some(ck);
        self.ik = Some(ik);
        self.ak = Some(ak);
        (res, ck, ik, ak)
    }

    /// F5Star is the anonymity key derivation function for the re-synchronisation message.
    /// F5Star takes key K and random challenge RAND, and returns resynch anonymity key AK.
    pub fn f5star(&mut self, rand: &[u8; 16]) -> [u8; 6] {
        let mut rijndael_input = xor(&self.opc, rand);
        let temp = self.rijndael_encrypt(&rijndael_input);

        // To obtain output block OUT5: XOR OPc and TEMP, rotate by r5=96, and XOR on the
        // constant c5 (which is all zeroes except that the 3rd from last bit is 1).
        for i in 0..16 {
            rijndael_input[(i + 4) % 16] = temp[i] ^ self.opc[i];
        }
        rijndael_input[15] ^= 8;

        let mut out = self.rijndael_encrypt(&rijndael_input);
        for (i, elem) in out.iter_mut().enumerate() {
            *elem ^= &self.opc[i];
        }

        let mut ak = [0u8; 6];
        ak.copy_from_slice(&out[..6]);

        self.ak = Some(ak);
        ak
    }

    /// Computes RESStar from serving network name, RAND and RES as described in
    /// A.4 RES* and XRES* derivation function, TS 33.501.
    pub fn compute_res_star(
        &mut self,
        mcc: &str,
        mnc: &str,
        rand: &[u8; 16],
        res: &[u8; 8],
    ) -> Result<[u8; 16], MilenageError> {
        let mut n: String = mnc.to_string();
        if mnc.len() == 2 {
            n = format!("0{}", mnc);
        } else if mnc.len() != 3 {
            return Err(MilenageError::InvalidMnc(mnc.to_string()));
        };

        if mcc.len() != 3 {
            return Err(MilenageError::InvalidMcc(mcc.to_string()));
        };

        let snn = format!("5G:mnc{}.mcc{}.3gppnetwork.org", n, mcc);

        let mut data = [0u8; 63];
        data[0] = 0x6bu8;

        data[1..33].copy_from_slice(snn.as_bytes());
        data[33..35].copy_from_slice(&[0x00u8, 0x20u8]);

        data[35..51].copy_from_slice(rand);
        data[51..53].copy_from_slice(&[0x00u8, 0x10u8]);

        data[53..61].copy_from_slice(res);
        data[61..63].copy_from_slice(&[0x00u8, 0x08u8]);

        let ck = self.ck.ok_or(MilenageError::MissingCk)?;
        let ik = self.ik.ok_or(MilenageError::MissingIk)?;

        let mut k = [0u8; 32];
        k[0..16].copy_from_slice(&ck);
        k[16..32].copy_from_slice(&ik);

        type HmacSha256 = Hmac<Sha256>;
        let mut mac = HmacSha256::new_from_slice(&k).expect("HMAC can take key of any size");
        mac.update(&data);
        let result = mac.finalize().into_bytes();

        let mut res_star = [0u8; 16];
        res_star.copy_from_slice(&result[16..32]);
        self.res_star = Some(res_star);
        Ok(res_star)
    }

    /// Derive OP with K to produce OPc
    fn compute_opc(&mut self, op: [u8; 16]) {
        let ciphered_opc = self.rijndael_encrypt(&op);
        self.opc = xor(&ciphered_opc, &op);
    }

    #[cfg(feature = "aes")]
    fn rijndael_encrypt(&self, input: &[u8; 16]) -> [u8; 16] {
        use aes::cipher::array::Array;
        use aes::cipher::{BlockCipherEncrypt as _, KeyInit as _};

        let cipher = Aes128::new((&self.k).into());
        let mut block: Array<_, _> = (*input).into();
        cipher.encrypt_block(&mut block);
        let mut output = [0u8; 16];
        output.copy_from_slice(&block);

        output
    }

    #[cfg(feature = "openssl")]
    fn rijndael_encrypt(&self, input: &[u8; 16]) -> [u8; 16] {
        let key = AesKey::new_encrypt(&self.k).unwrap();
        let mut iv = [0; 32];
        let mut output = [0u8; 16];
        aes_ige(input, &mut output, &key, &mut iv, Mode::Encrypt);

        output
    }
}

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

    // =========================================================================
    // Test vectors from 3GPP TS 35.208 - Test Sets 1 through 6
    // =========================================================================

    // --- Test Set 1 (3GPP TS 35.208 Section 4) ---

    #[test]
    fn test_set1_opc() {
        let k = hex!("465b5ce8b199b49faa5f0a2ee238a6bc");
        let op = hex!("cdc202d5123e20f62b6d676ac72cb318");
        let m = Milenage::new_with_op(k, op);
        assert_eq!(m.opc(), &hex!("cd63cb71954a9f4e48a5994e37a02baf"));
    }

    #[test]
    fn test_set1_f1() {
        let k = hex!("465b5ce8b199b49faa5f0a2ee238a6bc");
        let op = hex!("cdc202d5123e20f62b6d676ac72cb318");
        let rand = hex!("23553cbe9637a89d218ae64dae47bf35");
        let sqn = hex!("ff9bb4d0b607");
        let amf = hex!("b9b9");
        let mut m = Milenage::new_with_op(k, op);
        let maca = m.f1(&rand, &sqn, &amf);
        assert_eq!(maca, hex!("4a9ffac354dfafb3"));
    }

    #[test]
    fn test_set1_f1star() {
        let k = hex!("465b5ce8b199b49faa5f0a2ee238a6bc");
        let op = hex!("cdc202d5123e20f62b6d676ac72cb318");
        let rand = hex!("23553cbe9637a89d218ae64dae47bf35");
        let sqn = hex!("ff9bb4d0b607");
        let amf = hex!("b9b9");
        let mut m = Milenage::new_with_op(k, op);
        let macs = m.f1star(&rand, &sqn, &amf);
        assert_eq!(macs, hex!("01cfaf9ec4e871e9"));
    }

    #[test]
    fn test_set1_f2345() {
        let k = hex!("465b5ce8b199b49faa5f0a2ee238a6bc");
        let op = hex!("cdc202d5123e20f62b6d676ac72cb318");
        let rand = hex!("23553cbe9637a89d218ae64dae47bf35");
        let mut m = Milenage::new_with_op(k, op);
        let (res, ck, ik, ak) = m.f2345(&rand);
        assert_eq!(res, hex!("a54211d5e3ba50bf"));
        assert_eq!(ck, hex!("b40ba9a3c58b2a05bbf0d987b21bf8cb"));
        assert_eq!(ik, hex!("f769bcd751044604127672711c6d3441"));
        assert_eq!(ak, hex!("aa689c648370"));
    }

    #[test]
    fn test_set1_f5star() {
        let k = hex!("465b5ce8b199b49faa5f0a2ee238a6bc");
        let op = hex!("cdc202d5123e20f62b6d676ac72cb318");
        let rand = hex!("23553cbe9637a89d218ae64dae47bf35");
        let mut m = Milenage::new_with_op(k, op);
        let ak = m.f5star(&rand);
        assert_eq!(ak, hex!("451e8beca43b"));
    }

    // --- Test Set 2 (3GPP TS 35.208 Section 5) ---

    #[test]
    fn test_set2_opc() {
        let k = hex!("0396eb317b6d1c36f19c1c84cd6ffd16");
        let op = hex!("ff53bade17df5d4e793073ce9d7579fa");
        let m = Milenage::new_with_op(k, op);
        assert_eq!(m.opc(), &hex!("53c15671c60a4b731c55b4a441c0bde2"));
    }

    #[test]
    fn test_set2_f1() {
        let k = hex!("0396eb317b6d1c36f19c1c84cd6ffd16");
        let op = hex!("ff53bade17df5d4e793073ce9d7579fa");
        let rand = hex!("c00d603103dcee52c4478119494202e8");
        let sqn = hex!("fd8eef40df7d");
        let amf = hex!("af17");
        let mut m = Milenage::new_with_op(k, op);
        let maca = m.f1(&rand, &sqn, &amf);
        assert_eq!(maca, hex!("5df5b31807e258b0"));
    }

    #[test]
    fn test_set2_f1star() {
        let k = hex!("0396eb317b6d1c36f19c1c84cd6ffd16");
        let op = hex!("ff53bade17df5d4e793073ce9d7579fa");
        let rand = hex!("c00d603103dcee52c4478119494202e8");
        let sqn = hex!("fd8eef40df7d");
        let amf = hex!("af17");
        let mut m = Milenage::new_with_op(k, op);
        let macs = m.f1star(&rand, &sqn, &amf);
        assert_eq!(macs, hex!("a8c016e51ef4a343"));
    }

    #[test]
    fn test_set2_f2345() {
        let k = hex!("0396eb317b6d1c36f19c1c84cd6ffd16");
        let op = hex!("ff53bade17df5d4e793073ce9d7579fa");
        let rand = hex!("c00d603103dcee52c4478119494202e8");
        let mut m = Milenage::new_with_op(k, op);
        let (res, ck, ik, ak) = m.f2345(&rand);
        assert_eq!(res, hex!("d3a628ed988620f0"));
        assert_eq!(ck, hex!("58c433ff7a7082acd424220f2b67c556"));
        assert_eq!(ik, hex!("21a8c1f929702adb3e738488b9f5c5da"));
        assert_eq!(ak, hex!("c47783995f72"));
    }

    #[test]
    fn test_set2_f5star() {
        let k = hex!("0396eb317b6d1c36f19c1c84cd6ffd16");
        let op = hex!("ff53bade17df5d4e793073ce9d7579fa");
        let rand = hex!("c00d603103dcee52c4478119494202e8");
        let mut m = Milenage::new_with_op(k, op);
        let ak = m.f5star(&rand);
        assert_eq!(ak, hex!("30f1197061c1"));
    }

    // --- Test Set 3 (3GPP TS 35.208 Section 6) ---

    #[test]
    fn test_set3_opc() {
        let k = hex!("fec86ba6eb707ed08905757b1bb44b8f");
        let op = hex!("dbc59adcb6f9a0ef735477b7fadf8374");
        let m = Milenage::new_with_op(k, op);
        assert_eq!(m.opc(), &hex!("1006020f0a478bf6b699f15c062e42b3"));
    }

    #[test]
    fn test_set3_f1() {
        let k = hex!("fec86ba6eb707ed08905757b1bb44b8f");
        let op = hex!("dbc59adcb6f9a0ef735477b7fadf8374");
        let rand = hex!("9f7c8d021accf4db213ccff0c7f71a6a");
        let sqn = hex!("9d0277595ffc");
        let amf = hex!("725c");
        let mut m = Milenage::new_with_op(k, op);
        let maca = m.f1(&rand, &sqn, &amf);
        assert_eq!(maca, hex!("9cabc3e99baf7281"));
    }

    #[test]
    fn test_set3_f1star() {
        let k = hex!("fec86ba6eb707ed08905757b1bb44b8f");
        let op = hex!("dbc59adcb6f9a0ef735477b7fadf8374");
        let rand = hex!("9f7c8d021accf4db213ccff0c7f71a6a");
        let sqn = hex!("9d0277595ffc");
        let amf = hex!("725c");
        let mut m = Milenage::new_with_op(k, op);
        let macs = m.f1star(&rand, &sqn, &amf);
        assert_eq!(macs, hex!("95814ba2b3044324"));
    }

    #[test]
    fn test_set3_f2345() {
        let k = hex!("fec86ba6eb707ed08905757b1bb44b8f");
        let op = hex!("dbc59adcb6f9a0ef735477b7fadf8374");
        let rand = hex!("9f7c8d021accf4db213ccff0c7f71a6a");
        let mut m = Milenage::new_with_op(k, op);
        let (res, ck, ik, ak) = m.f2345(&rand);
        assert_eq!(res, hex!("8011c48c0c214ed2"));
        assert_eq!(ck, hex!("5dbdbb2954e8f3cde665b046179a5098"));
        assert_eq!(ik, hex!("59a92d3b476a0443487055cf88b2307b"));
        assert_eq!(ak, hex!("33484dc2136b"));
    }

    #[test]
    fn test_set3_f5star() {
        let k = hex!("fec86ba6eb707ed08905757b1bb44b8f");
        let op = hex!("dbc59adcb6f9a0ef735477b7fadf8374");
        let rand = hex!("9f7c8d021accf4db213ccff0c7f71a6a");
        let mut m = Milenage::new_with_op(k, op);
        let ak = m.f5star(&rand);
        assert_eq!(ak, hex!("deacdd848cc6"));
    }

    // --- Test Set 4 (3GPP TS 35.208 Section 7) ---

    #[test]
    fn test_set4_opc() {
        let k = hex!("9e5944aea94b81165c82fbf9f32db751");
        let op = hex!("223014c5806694c007ca1eeef57f004f");
        let m = Milenage::new_with_op(k, op);
        assert_eq!(m.opc(), &hex!("a64a507ae1a2a98bb88eb4210135dc87"));
    }

    #[test]
    fn test_set4_f1() {
        let k = hex!("9e5944aea94b81165c82fbf9f32db751");
        let op = hex!("223014c5806694c007ca1eeef57f004f");
        let rand = hex!("ce83dbc54ac0274a157c17f80d017bd6");
        let sqn = hex!("0b604a81eca8");
        let amf = hex!("9e09");
        let mut m = Milenage::new_with_op(k, op);
        let maca = m.f1(&rand, &sqn, &amf);
        assert_eq!(maca, hex!("74a58220cba84c49"));
    }

    #[test]
    fn test_set4_f1star() {
        let k = hex!("9e5944aea94b81165c82fbf9f32db751");
        let op = hex!("223014c5806694c007ca1eeef57f004f");
        let rand = hex!("ce83dbc54ac0274a157c17f80d017bd6");
        let sqn = hex!("0b604a81eca8");
        let amf = hex!("9e09");
        let mut m = Milenage::new_with_op(k, op);
        let macs = m.f1star(&rand, &sqn, &amf);
        assert_eq!(macs, hex!("ac2cc74a96871837"));
    }

    #[test]
    fn test_set4_f2345() {
        let k = hex!("9e5944aea94b81165c82fbf9f32db751");
        let op = hex!("223014c5806694c007ca1eeef57f004f");
        let rand = hex!("ce83dbc54ac0274a157c17f80d017bd6");
        let mut m = Milenage::new_with_op(k, op);
        let (res, ck, ik, ak) = m.f2345(&rand);
        assert_eq!(res, hex!("f365cd683cd92e96"));
        assert_eq!(ck, hex!("e203edb3971574f5a94b0d61b816345d"));
        assert_eq!(ik, hex!("0c4524adeac041c4dd830d20854fc46b"));
        assert_eq!(ak, hex!("f0b9c08ad02e"));
    }

    #[test]
    fn test_set4_f5star() {
        let k = hex!("9e5944aea94b81165c82fbf9f32db751");
        let op = hex!("223014c5806694c007ca1eeef57f004f");
        let rand = hex!("ce83dbc54ac0274a157c17f80d017bd6");
        let mut m = Milenage::new_with_op(k, op);
        let ak = m.f5star(&rand);
        assert_eq!(ak, hex!("6085a86c6f63"));
    }

    // --- Test Set 5 (3GPP TS 35.208 Section 8) ---

    #[test]
    fn test_set5_opc() {
        let k = hex!("4ab1deb05ca6ceb051fc98e77d026a84");
        let op = hex!("2d16c5cd1fdf6b22383584e3bef2a8d8");
        let m = Milenage::new_with_op(k, op);
        assert_eq!(m.opc(), &hex!("dcf07cbd51855290b92a07a9891e523e"));
    }

    #[test]
    fn test_set5_f1() {
        let k = hex!("4ab1deb05ca6ceb051fc98e77d026a84");
        let op = hex!("2d16c5cd1fdf6b22383584e3bef2a8d8");
        let rand = hex!("74b0cd6031a1c8339b2b6ce2b8c4a186");
        let sqn = hex!("e880a1b580b6");
        let amf = hex!("9f07");
        let mut m = Milenage::new_with_op(k, op);
        let maca = m.f1(&rand, &sqn, &amf);
        assert_eq!(maca, hex!("49e785dd12626ef2"));
    }

    #[test]
    fn test_set5_f1star() {
        let k = hex!("4ab1deb05ca6ceb051fc98e77d026a84");
        let op = hex!("2d16c5cd1fdf6b22383584e3bef2a8d8");
        let rand = hex!("74b0cd6031a1c8339b2b6ce2b8c4a186");
        let sqn = hex!("e880a1b580b6");
        let amf = hex!("9f07");
        let mut m = Milenage::new_with_op(k, op);
        let macs = m.f1star(&rand, &sqn, &amf);
        assert_eq!(macs, hex!("9e85790336bb3fa2"));
    }

    #[test]
    fn test_set5_f2345() {
        let k = hex!("4ab1deb05ca6ceb051fc98e77d026a84");
        let op = hex!("2d16c5cd1fdf6b22383584e3bef2a8d8");
        let rand = hex!("74b0cd6031a1c8339b2b6ce2b8c4a186");
        let mut m = Milenage::new_with_op(k, op);
        let (res, ck, ik, ak) = m.f2345(&rand);
        assert_eq!(res, hex!("5860fc1bce351e7e"));
        assert_eq!(ck, hex!("7657766b373d1c2138f307e3de9242f9"));
        assert_eq!(ik, hex!("1c42e960d89b8fa99f2744e0708ccb53"));
        assert_eq!(ak, hex!("31e11a609118"));
    }

    #[test]
    fn test_set5_f5star() {
        let k = hex!("4ab1deb05ca6ceb051fc98e77d026a84");
        let op = hex!("2d16c5cd1fdf6b22383584e3bef2a8d8");
        let rand = hex!("74b0cd6031a1c8339b2b6ce2b8c4a186");
        let mut m = Milenage::new_with_op(k, op);
        let ak = m.f5star(&rand);
        assert_eq!(ak, hex!("fe2555e54aa9"));
    }

    // --- Test Set 6 (3GPP TS 35.208 Section 9) ---

    #[test]
    fn test_set6_opc() {
        let k = hex!("6c38a116ac280c454f59332ee35c8c4f");
        let op = hex!("1ba00a1a7c6700ac8c3ff3e96ad08725");
        let m = Milenage::new_with_op(k, op);
        assert_eq!(m.opc(), &hex!("3803ef5363b947c6aaa225e58fae3934"));
    }

    #[test]
    fn test_set6_f1() {
        let k = hex!("6c38a116ac280c454f59332ee35c8c4f");
        let op = hex!("1ba00a1a7c6700ac8c3ff3e96ad08725");
        let rand = hex!("ee6466bc96202c5a557abbeff8babf63");
        let sqn = hex!("414b98222181");
        let amf = hex!("4464");
        let mut m = Milenage::new_with_op(k, op);
        let maca = m.f1(&rand, &sqn, &amf);
        assert_eq!(maca, hex!("078adfb488241a57"));
    }

    #[test]
    fn test_set6_f1star() {
        let k = hex!("6c38a116ac280c454f59332ee35c8c4f");
        let op = hex!("1ba00a1a7c6700ac8c3ff3e96ad08725");
        let rand = hex!("ee6466bc96202c5a557abbeff8babf63");
        let sqn = hex!("414b98222181");
        let amf = hex!("4464");
        let mut m = Milenage::new_with_op(k, op);
        let macs = m.f1star(&rand, &sqn, &amf);
        assert_eq!(macs, hex!("80246b8d0186bcf1"));
    }

    #[test]
    fn test_set6_f2345() {
        let k = hex!("6c38a116ac280c454f59332ee35c8c4f");
        let op = hex!("1ba00a1a7c6700ac8c3ff3e96ad08725");
        let rand = hex!("ee6466bc96202c5a557abbeff8babf63");
        let mut m = Milenage::new_with_op(k, op);
        let (res, ck, ik, ak) = m.f2345(&rand);
        assert_eq!(res, hex!("16c8233f05a0ac28"));
        assert_eq!(ck, hex!("3f8c7587fe8e4b233af676aede30ba3b"));
        assert_eq!(ik, hex!("a7466cc1e6b2a1337d49d3b66e95d7b4"));
        assert_eq!(ak, hex!("45b0f69ab06c"));
    }

    #[test]
    fn test_set6_f5star() {
        let k = hex!("6c38a116ac280c454f59332ee35c8c4f");
        let op = hex!("1ba00a1a7c6700ac8c3ff3e96ad08725");
        let rand = hex!("ee6466bc96202c5a557abbeff8babf63");
        let mut m = Milenage::new_with_op(k, op);
        let ak = m.f5star(&rand);
        assert_eq!(ak, hex!("1f53cd2b1113"));
    }

    // =========================================================================
    // Additional tests
    // =========================================================================

    #[test]
    fn test_new_with_opc() {
        let k = hex!("465b5ce8b199b49faa5f0a2ee238a6bc");
        let opc = hex!("cd63cb71954a9f4e48a5994e37a02baf");
        let rand = hex!("23553cbe9637a89d218ae64dae47bf35");
        let mut m = Milenage::new_with_opc(k, opc);
        let (res, ck, ik, ak) = m.f2345(&rand);
        assert_eq!(res, hex!("a54211d5e3ba50bf"));
        assert_eq!(ck, hex!("b40ba9a3c58b2a05bbf0d987b21bf8cb"));
        assert_eq!(ik, hex!("f769bcd751044604127672711c6d3441"));
        assert_eq!(ak, hex!("aa689c648370"));
    }

    #[test]
    fn test_compute_res_star() {
        let k = hex!("465b5ce8b199b49faa5f0a2ee238a6bc");
        let op = hex!("cdc202d5123e20f62b6d676ac72cb318");
        let rand = hex!("23553cbe9637a89d218ae64dae47bf35");
        let mut m = Milenage::new_with_op(k, op);
        let (res, _, _, _) = m.f2345(&rand);
        let res_star = m.compute_res_star("001", "01", &rand, &res).unwrap();
        assert_eq!(res_star, hex!("f236a7417272bfb2d66d4d670733b527"));
    }

    #[test]
    fn test_compute_res_star_missing_ck() {
        let k = hex!("465b5ce8b199b49faa5f0a2ee238a6bc");
        let op = hex!("cdc202d5123e20f62b6d676ac72cb318");
        let rand = hex!("23553cbe9637a89d218ae64dae47bf35");
        let mut m = Milenage::new_with_op(k, op);
        let res = hex!("a54211d5e3ba50bf");
        let err = m.compute_res_star("001", "01", &rand, &res).unwrap_err();
        assert_eq!(err, MilenageError::MissingCk);
    }

    #[test]
    fn test_compute_res_star_invalid_mnc() {
        let k = hex!("465b5ce8b199b49faa5f0a2ee238a6bc");
        let op = hex!("cdc202d5123e20f62b6d676ac72cb318");
        let rand = hex!("23553cbe9637a89d218ae64dae47bf35");
        let mut m = Milenage::new_with_op(k, op);
        let (res, _, _, _) = m.f2345(&rand);
        let err = m.compute_res_star("001", "1", &rand, &res).unwrap_err();
        assert_eq!(err, MilenageError::InvalidMnc("1".to_string()));
    }

    #[test]
    fn test_compute_res_star_invalid_mcc() {
        let k = hex!("465b5ce8b199b49faa5f0a2ee238a6bc");
        let op = hex!("cdc202d5123e20f62b6d676ac72cb318");
        let rand = hex!("23553cbe9637a89d218ae64dae47bf35");
        let mut m = Milenage::new_with_op(k, op);
        let (res, _, _, _) = m.f2345(&rand);
        let err = m.compute_res_star("01", "01", &rand, &res).unwrap_err();
        assert_eq!(err, MilenageError::InvalidMcc("01".to_string()));
    }
}
