/* Copyright (C) 2021 Open Information Security Foundation

 *

 * You can copy, redistribute or modify this Program under the terms of

 * the GNU General Public License version 2 as published by the Free

 * Software Foundation.

 *

 * This program is distributed in the hope that it will be useful,

 * but WITHOUT ANY WARRANTY; without even the implied warranty of

 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the

 * GNU General Public License for more details.

 *

 * You should have received a copy of the GNU General Public License

 * version 2 along with this program; if not, write to the Free Software

 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA

 * 02110-1301, USA.

 */

use aes::cipher::generic_array::GenericArray;

use aes::Aes128;

use aes::BlockEncrypt;

use aes::NewBlockCipher;

use aes_gcm::AeadInPlace;

use aes_gcm::Aes128Gcm;

use aes_gcm::NewAead;

use hkdf::Hkdf;

use sha2::Sha256;

pub const AES128_KEY_LEN: usize = 16;

pub const AES128_TAG_LEN: usize = 16;

pub const AES128_IV_LEN: usize = 12;

pub struct HeaderProtectionKey(Aes128);

impl HeaderProtectionKey {

    fn new(secret: &[u8], version: u32) -> Self {

        let hk = Hkdf::<Sha256>::from_prk(secret).unwrap();

        let mut secret = [0u8; AES128_KEY_LEN];

        let quichp = if version == 0x6b3343cf {

            b"quicv2 hp" as &[u8]

        } else {

            b"quic hp" as &[u8]

        };

        hkdf_expand_label(&hk, quichp, &mut secret, AES128_KEY_LEN as u16);

        return Self(Aes128::new(GenericArray::from_slice(&secret)));

    }

    pub fn decrypt_in_place(

        &self, sample: &[u8], first: &mut u8, packet_number: &mut [u8],

    ) -> Result<(), ()> {

        let mut mask = GenericArray::clone_from_slice(sample);

        self.0.encrypt_block(&mut mask);

        let (first_mask, pn_mask) = mask.split_first().unwrap();

        let bits = if (*first & 0x80) != 0 {

            0x0f // Long header: 4 bits masked

        } else {

            0x1f // Short header: 5 bits masked

        };

        *first ^= first_mask & bits;

        let pn_len = (*first & 0x03) as usize + 1;

        for (dst, m) in packet_number.iter_mut().zip(pn_mask).take(pn_len) {

            *dst ^= m;

        }

        Ok(())

    }

}

pub struct PacketKey {

    key: Aes128Gcm,

    iv: [u8; AES128_IV_LEN],

}

impl PacketKey {

    fn new(secret: &[u8], version: u32) -> Self {

        let hk = Hkdf::<Sha256>::from_prk(secret).unwrap();

        let mut secret = [0u8; AES128_KEY_LEN];

        let quickey = if version == 0x6b3343cf {

            b"quicv2 key" as &[u8]

        } else {

            b"quic key" as &[u8]

        };

        hkdf_expand_label(&hk, quickey, &mut secret, AES128_KEY_LEN as u16);

        let key = Aes128Gcm::new(GenericArray::from_slice(&secret));

        let mut r = PacketKey {

            key,

            iv: [0u8; AES128_IV_LEN],

        };

        let quiciv = if version == 0x6b3343cf {

            b"quicv2 iv" as &[u8]

        } else {

            b"quic iv" as &[u8]

        };

        hkdf_expand_label(&hk, quiciv, &mut r.iv, AES128_IV_LEN as u16);

        return r;

    }

    pub fn decrypt_in_place<'a>(

        &self, packet_number: u64, header: &[u8], payload: &'a mut [u8],

    ) -> Result<&'a [u8], ()> {

        if payload.len() < AES128_TAG_LEN {

            return Err(());

        }

        let mut nonce = [0; AES128_IV_LEN];

        nonce[4..].copy_from_slice(&packet_number.to_be_bytes());

        for (nonce, inp) in nonce.iter_mut().zip(self.iv.iter()) {

            *nonce ^= inp;

        }

        let tag_pos = payload.len() - AES128_TAG_LEN;

        let (buffer, tag) = payload.split_at_mut(tag_pos);

        let taga = GenericArray::from_slice(tag);

        self.key

            .decrypt_in_place_detached(GenericArray::from_slice(&nonce), header, buffer, taga)

            .map_err(|_| ())?;

        Ok(&payload[..tag_pos])

    }

}

pub struct DirectionalKeys {

    pub header: HeaderProtectionKey,

    pub packet: PacketKey,

}

impl DirectionalKeys {

    fn new(secret: &[u8], version: u32) -> Self {

        Self {

            header: HeaderProtectionKey::new(secret, version),

            packet: PacketKey::new(secret, version),

        }

    }

}

pub struct QuicKeys {

    pub local: DirectionalKeys,

    pub remote: DirectionalKeys,

}

fn hkdf_expand_label(hk: &Hkdf<Sha256>, label: &[u8], okm: &mut [u8], olen: u16) {

    const LABEL_PREFIX: &[u8] = b"tls13 ";

    let output_len = u16::to_be_bytes(olen);

    let label_len = u8::to_be_bytes((LABEL_PREFIX.len() + label.len()) as u8);

    let context_len = u8::to_be_bytes(0);

    let info = &[

        &output_len[..],

        &label_len[..],

        LABEL_PREFIX,

        label,

        &context_len[..],

    ];

    hk.expand_multi_info(info, okm).unwrap();

}

pub fn quic_keys_initial(version: u32, client_dst_connection_id: &[u8]) -> Option<QuicKeys> {

    let salt = match version {

        0x51303530 => &[

            0x50, 0x45, 0x74, 0xEF, 0xD0, 0x66, 0xFE, 0x2F, 0x9D, 0x94, 0x5C, 0xFC, 0xDB, 0xD3,

            0xA7, 0xF0, 0xD3, 0xB5, 0x6B, 0x45,

        ],

        0xff00_001d..=0xff00_0020 => &[

            // https://datatracker.ietf.org/doc/html/draft-ietf-quic-tls-32#section-5.2

            0xaf, 0xbf, 0xec, 0x28, 0x99, 0x93, 0xd2, 0x4c, 0x9e, 0x97, 0x86, 0xf1, 0x9c, 0x61,

            0x11, 0xe0, 0x43, 0x90, 0xa8, 0x99,

        ],

        0xfaceb002 | 0xff00_0017..=0xff00_001c => &[

            // https://datatracker.ietf.org/doc/html/draft-ietf-quic-tls-23#section-5.2

            0xc3, 0xee, 0xf7, 0x12, 0xc7, 0x2e, 0xbb, 0x5a, 0x11, 0xa7, 0xd2, 0x43, 0x2b, 0xb4,

            0x63, 0x65, 0xbe, 0xf9, 0xf5, 0x02,

        ],

        0x0000_0001 | 0xff00_0021..=0xff00_0022 => &[

            // https://www.rfc-editor.org/rfc/rfc9001.html#name-initial-secrets

            0x38, 0x76, 0x2c, 0xf7, 0xf5, 0x59, 0x34, 0xb3, 0x4d, 0x17, 0x9a, 0xe6, 0xa4, 0xc8,

            0x0c, 0xad, 0xcc, 0xbb, 0x7f, 0x0a,

        ],

        0x6b3343cf => &[

            // https://www.rfc-editor.org/rfc/rfc9369.html#section-3.3.1

            0x0d, 0xed, 0xe3, 0xde, 0xf7, 0x00, 0xa6, 0xdb, 0x81, 0x93, 0x81, 0xbe, 0x6e, 0x26,

            0x9d, 0xcb, 0xf9, 0xbd, 0x2e, 0xd9,

        ],

        _ => {

            return None;

        }

    };

    let hk = Hkdf::<Sha256>::new(Some(salt), client_dst_connection_id);

    let mut client_secret = [0u8; 32];

    hkdf_expand_label(&hk, b"client in", &mut client_secret, 32);

    let mut server_secret = [0u8; 32];

    hkdf_expand_label(&hk, b"server in", &mut server_secret, 32);

    return Some(QuicKeys {

        local: DirectionalKeys::new(&server_secret, version),

        remote: DirectionalKeys::new(&client_secret, version),

    });

}