/* Copyright (C) 2023-2024 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.

// Author: Sascha Steinbiss <sascha@steinbiss.name>

*/
#[cfg(feature = "ja4")]
use crate::jsonbuilder::HEX;
#[cfg(feature = "ja4")]
use digest::Digest;
#[cfg(feature = "ja4")]
use sha2::Sha256;
#[cfg(feature = "ja4")]
use std::cmp::min;
#[cfg(feature = "ja4")]
use tls_parser::{TlsExtensionType, TlsVersion};

use crate::handshake::HandshakeParams;

pub const JA4_HEX_LEN: usize = 36;

pub(crate) trait JA4Impl {
    fn try_new(hs: &HandshakeParams) -> Option<JA4>;
}

#[derive(Debug, PartialEq)]
pub struct JA4 {
    hash: String,
}

impl AsRef<str> for JA4 {
    fn as_ref(&self) -> &str {
        &self.hash
    }
}

#[cfg(feature = "ja4")]
impl JA4 {
    #[inline]
    fn version_to_ja4code(val: Option<TlsVersion>) -> &'static str {
        match val {
            Some(TlsVersion::Tls13) => "13",
            Some(TlsVersion::Tls12) => "12",
            Some(TlsVersion::Tls11) => "11",
            Some(TlsVersion::Tls10) => "10",
            Some(TlsVersion::Ssl30) => "s3",
            // the TLS parser does not support SSL 1.0 and 2.0 hence no
            // support for "s1"/"s2"
            _ => "00",
        }
    }

    fn format_alpn(alpn: Option<&Vec<u8>>) -> [char; 2] {
        let mut ret = ['0', '0'];

        if let Some(alpn) = alpn {
            if !alpn.is_empty() {
                // If the first ALPN value is only a single character, then that character is treated as both the first and last character.
                if alpn.len() == 2 {
                    // GREASE values are 2 bytes, so this could be one -- check
                    let v: u16 = ((alpn[0] as u16) << 8) | alpn[alpn.len() - 1] as u16;
                    if HandshakeParams::is_grease(v) {
                        return ret;
                    }
                }
                if !alpn[0].is_ascii_alphanumeric() || !alpn[alpn.len() - 1].is_ascii_alphanumeric()
                {
                    // If the first or last byte of the first ALPN is non-alphanumeric (meaning not 0x30-0x39, 0x41-0x5A, or 0x61-0x7A), then we print the first and last characters of the hex representation of the first ALPN instead.
                    ret[0] = char::from(HEX[(alpn[0] >> 4) as usize]);
                    ret[1] = char::from(HEX[(alpn[alpn.len() - 1] & 0xF) as usize]);
                } else {
                    ret[0] = char::from(alpn[0]);
                    ret[1] = char::from(alpn[alpn.len() - 1]);
                }
            }
        }
        ret
    }
}

#[cfg(feature = "ja4")]
impl JA4Impl for JA4 {
    fn try_new(hs: &HandshakeParams) -> Option<Self> {
        // All non-GREASE extensions are stored to produce a more verbose, complete output
        // of extensions but we need to omit ALPN & SNI extensions from the JA4_a hash.
        let mut exts = hs
            .extensions
            .iter()
            .filter(|&ext| {
                *ext != TlsExtensionType::ApplicationLayerProtocolNegotiation
                    && *ext != TlsExtensionType::ServerName
            })
            .collect::<Vec<&TlsExtensionType>>();

        let alpn = Self::format_alpn(hs.alpns.first());

        // Calculate JA4_a
        let ja4_a = format!(
            "{proto}{version}{sni}{nof_c:02}{nof_e:02}{al1}{al2}",
            proto = if hs.quic { "q" } else { "t" },
            version = Self::version_to_ja4code(hs.tls_version),
            sni = if hs.domain { "d" } else { "i" },
            nof_c = min(99, hs.ciphersuites.len()),
            nof_e = min(99, hs.extensions.len()),
            al1 = alpn[0],
            al2 = alpn[1]
        );

        // Calculate JA4_b
        let mut sorted_ciphers = hs.ciphersuites.to_vec();
        sorted_ciphers.sort_by(|a, b| u16::from(*a).cmp(&u16::from(*b)));
        let sorted_cipherstrings: Vec<String> = sorted_ciphers
            .iter()
            .map(|v| format!("{:04x}", u16::from(*v)))
            .collect();
        let mut sha = Sha256::new();
        let ja4_b_raw = sorted_cipherstrings.join(",");
        sha.update(&ja4_b_raw);
        let mut ja4_b = format!("{:x}", sha.finalize_reset());
        ja4_b.truncate(12);

        // Calculate JA4_c
        exts.sort_by(|&a, &b| u16::from(*a).cmp(&u16::from(*b)));
        let sorted_extstrings: Vec<String> = exts
            .into_iter()
            .map(|&v| format!("{:04x}", u16::from(v)))
            .collect();
        let ja4_c1_raw = sorted_extstrings.join(",");
        let unsorted_sigalgostrings: Vec<String> = hs
            .signature_algorithms
            .iter()
            .map(|v| format!("{:04x}", (*v)))
            .collect();
        let ja4_c2_raw = unsorted_sigalgostrings.join(",");
        let ja4_c_raw = format!("{}_{}", ja4_c1_raw, ja4_c2_raw);
        sha.update(&ja4_c_raw);
        let mut ja4_c = format!("{:x}", sha.finalize());
        ja4_c.truncate(12);

        Some(Self {
            hash: format!("{}_{}_{}", ja4_a, ja4_b, ja4_c),
        })
    }
}

#[cfg(not(feature = "ja4"))]
impl JA4Impl for JA4 {
    fn try_new(_hs: &HandshakeParams) -> Option<Self> {
        None
    }
}

// C ABI
#[cfg(feature = "ja4")]
#[no_mangle]
pub unsafe extern "C" fn SCJA4GetHash(hs: &HandshakeParams, out: &mut [u8; JA4_HEX_LEN]) {
    if let Some(ja4) = JA4::try_new(hs) {
        out[0..JA4_HEX_LEN].copy_from_slice(ja4.as_ref().as_bytes());
    }
}

#[cfg(test)]
#[cfg(feature = "ja4")]
mod tests {
    use super::*;
    use tls_parser::{TlsCipherSuiteID, TlsExtensionType, TlsVersion};

    #[test]
    fn test_format_alpn_ascii() {
        let res = JA4::format_alpn(Some(&"http/1.1".as_bytes().to_vec()));
        assert_eq!(res, ['h', '1']);
    }

    #[test]
    fn test_add_alpn_non_ascii_first_or_last() {
        let res = JA4::format_alpn(Some(&vec![0x01, b'T', b'E', 0x7f])); // non-alphanumeric start and end
        assert_eq!(res, [HEX[0x0], HEX[0xF]].map(|b| b as char)); // 0x01 -> 0, 0x7f -> f
    }

    #[test]
    fn test_add_alpn_grease_pair_filtered() {
        let res = JA4::format_alpn(Some(&vec![0x2a, 0x2a])); // 0x2a2a GREASE
        assert_eq!(res, ['0', '0']);
    }

    #[test]
    fn test_hash_limit_numbers() {
        // Test whether the limitation of the extension and ciphersuite
        // count to 99 is reflected correctly.
        let mut hs = HandshakeParams::default();

        for i in 1..200 {
            hs.add_cipher_suite(TlsCipherSuiteID(i));
        }
        for i in 1..200 {
            hs.add_extension(TlsExtensionType(i));
        }

        let ja4 = JA4::try_new(&hs).expect("JA4 create failure");

        // Only testing the ja4_a portion of the hash, we we truncate to
        // ensure we're only testing this
        let mut ja4_hash = ja4.as_ref().to_string();
        ja4_hash.truncate(10);

        assert_eq!(ja4_hash, "t00i999900");
    }

    #[test]
    fn test_short_alpn() {
        let mut hs = HandshakeParams::default();
        hs.add_alpn("b".as_bytes());
        let mut s = JA4::try_new(&hs)
            .expect("JA4 create failure")
            .as_ref()
            .to_string();
        s.truncate(10);
        assert_eq!(s, "t00i0000bb");

        let mut hs = HandshakeParams::default();
        hs.add_alpn("h2".as_bytes());
        let mut s = JA4::try_new(&hs)
            .expect("JA4 create failure")
            .as_ref()
            .to_string();
        s.truncate(10);
        assert_eq!(s, "t00i0000h2");

        // from https://github.com/FoxIO-LLC/ja4/blob/main/technical_details/JA4.md#alpn-extension-value
        let mut hs = HandshakeParams::default();
        hs.add_alpn(&[0xab]);
        let mut s = JA4::try_new(&hs)
            .expect("JA4 create failure")
            .as_ref()
            .to_string();
        s.truncate(10);
        assert_eq!(s, "t00i0000ab");

        let mut hs = HandshakeParams::default();
        hs.add_alpn(&[0xab, 0xcd]);
        let mut s = JA4::try_new(&hs)
            .expect("JA4 create failure")
            .as_ref()
            .to_string();
        s.truncate(10);
        assert_eq!(s, "t00i0000ad");

        let mut hs = HandshakeParams::default();
        hs.add_alpn(&[0x30, 0xab]);
        let mut s = JA4::try_new(&hs)
            .expect("JA4 create failure")
            .as_ref()
            .to_string();
        s.truncate(10);
        assert_eq!(s, "t00i00003b");

        let mut hs = HandshakeParams::default();
        hs.add_alpn(&[0x30, 0x31, 0xab, 0xcd]);
        let mut s = JA4::try_new(&hs)
            .expect("JA4 create failure")
            .as_ref()
            .to_string();
        s.truncate(10);
        assert_eq!(s, "t00i00003d");

        let mut hs = HandshakeParams::default();
        hs.add_alpn(&[0x30, 0xab, 0xcd, 0x31]);
        let mut s = JA4::try_new(&hs)
            .expect("JA4 create failure")
            .as_ref()
            .to_string();
        s.truncate(10);
        assert_eq!(s, "t00i000001");
    }

    #[test]
    fn test_get_hash() {
        let mut hs = HandshakeParams::default();

        // the empty JA4 hash
        let s = JA4::try_new(&hs).expect("JA4 create failure");
        assert_eq!(s.as_ref(), "t00i000000_e3b0c44298fc_d2e2adf7177b");

        // set TLS version
        hs.set_tls_version(TlsVersion::Tls12);
        let s = JA4::try_new(&hs).expect("JA4 create failure");
        assert_eq!(s.as_ref(), "t12i000000_e3b0c44298fc_d2e2adf7177b");

        // set QUIC
        hs.quic = true;
        let s = JA4::try_new(&hs).expect("JA4 create failure");
        assert_eq!(s.as_ref(), "q12i000000_e3b0c44298fc_d2e2adf7177b");

        // set GREASE extension, should be ignored
        hs.add_extension(TlsExtensionType(0x0a0a));
        let s = JA4::try_new(&hs).expect("JA4 create failure");
        assert_eq!(s.as_ref(), "q12i000000_e3b0c44298fc_d2e2adf7177b");

        // set SNI extension, should only increase count and change i->d
        hs.add_extension(TlsExtensionType(0x0000));
        let s = JA4::try_new(&hs).expect("JA4 create failure");
        assert_eq!(s.as_ref(), "q12d000100_e3b0c44298fc_d2e2adf7177b");

        // set ALPN extension, should only increase count and set end of JA4_a
        hs.add_alpn(b"h3-16");
        hs.add_extension(TlsExtensionType::ApplicationLayerProtocolNegotiation);
        let s = JA4::try_new(&hs).expect("JA4 create failure");
        assert_eq!(s.as_ref(), "q12d0002h6_e3b0c44298fc_d2e2adf7177b");

        // set some ciphers
        hs.add_cipher_suite(TlsCipherSuiteID(0x1111));
        hs.add_cipher_suite(TlsCipherSuiteID(0x0a20));
        hs.add_cipher_suite(TlsCipherSuiteID(0xbada));
        let s = JA4::try_new(&hs).expect("JA4 create failure");
        assert_eq!(s.as_ref(), "q12d0302h6_f500716053f9_d2e2adf7177b");

        // set some extensions and signature algorithms
        hs.add_extension(TlsExtensionType(0xface));
        hs.add_extension(TlsExtensionType(0x0121));
        hs.add_extension(TlsExtensionType(0x1234));
        hs.add_signature_algorithm(0x6666);
        let s = JA4::try_new(&hs).expect("JA4 create failure");
        assert_eq!(s.as_ref(), "q12d0305h6_f500716053f9_2debc8880bae");
    }
}