use crate::flags::{NTLM_NEG_UNICODE, NTLM_NEG_VERSION, NTLM_REQUEST_TARGET};
use crate::messages::NTLM_SIGN;
use crate::time::NtlmTime;
use crate::utils;
use crate::utils::from_unicode;
use crate::AvPairs;
use crate::Result;
use crate::Version;
use nom::bytes::complete::{tag, take};
use nom::number::complete::{le_u16, le_u32};
use std::convert::TryInto;

const NTLM_MSG_CHALLENGE: u32 = 0x00000002;

const CHALLENGE_MSG_FIXED_SIZE: u32 = 56;

#[derive(Clone, Debug, Default, PartialEq)]
pub struct ChallengeMsg {
    pub flags: u32,

    /// A nonce generated by the authentication server.
    pub server_challenge: [u8; 8],

    /// The Operating System version of the host that sends the message. It
    /// must be sent when NTLM_NEG_VERSION is set.
    pub version: Option<Version>,

    /// Indicates the server authentication realm. If server is in a domain,
    /// then is the domain name, else is the server name. It must be sent when
    /// NTLMSSP_REQUEST_TARGET is set.
    pub target_name: Option<String>,

    /// An array with several information about the server.
    pub target_info: AvPairs,
}

impl ChallengeMsg {

    /// Indicates if the given flag is set in the message.
    pub fn has_flag(&self, flag: u32) -> bool {
        return (self.flags & flag) != 0;
    }

    /// Returns the NetBIOS computer name if it is in target_info.
    pub fn nb_computer_name(&self) -> Option<&String> {
        return self.target_info.nb_computer_name();
    }

    /// Returns the NetBIOS domain name if it is in target_info.
    pub fn nb_domain_name(&self) -> Option<&String> {
        return self.target_info.nb_domain_name();
    }

    /// Returns the DNS computer name if it is in target_info.
    pub fn dns_computer_name(&self) -> Option<&String> {
        return self.target_info.dns_computer_name();
    }

    /// Returns the DNS domain name if it is in target_info.
    pub fn dns_domain_name(&self) -> Option<&String> {
        return self.target_info.dns_domain_name();
    }

    /// Returns the DNS tree name (forest name) if it is in target_info.
    pub fn dns_tree_name(&self) -> Option<&String> {
        return self.target_info.dns_tree_name();
    }

    /// Returns the timestamp if it is in target_info.
    pub fn timestamp(&self) -> Option<&NtlmTime> {
        return self.target_info.timestamp();
    }

    /// Parse network bytes to build a Challenge message.
    pub fn parse(raw: &[u8]) -> Result<Self> {
        let (rest, _) = tag(NTLM_SIGN)(raw)?;
        let (rest, _) = tag(NTLM_MSG_CHALLENGE.to_le_bytes())(rest)?;

        let (rest, tn_len, tn_offset) = dec_len_offset!(rest);

        let (rest, flags) = le_u32(rest)?;
        let (rest, server_challenge) = take(8usize)(rest)?;

        // reserved, should be zeros
        let (rest, _) = take(8usize)(rest)?;

        let (rest, ti_len, ti_offset) = dec_len_offset!(rest);

        let version = if (flags & NTLM_NEG_VERSION) != 0 {
            Some(Version::parse(rest)?.1)
        } else {
            None
        };

        let target_name = if (flags & NTLM_REQUEST_TARGET) != 0 {
            let (raw_tn, _) = take(tn_offset as usize)(raw)?;
            let (_, tn_bytes) = take(tn_len as usize)(raw_tn)?;

            Some(if (flags & NTLM_NEG_UNICODE) != 0 {
                from_unicode(tn_bytes)?
            } else {
                String::from_utf8(tn_bytes.to_vec())?
            })
        } else {
            None
        };

        let target_info = if ti_len != 0 {
            let (raw_ti, _) = take(ti_offset as usize)(raw)?;
            AvPairs::parse(raw_ti)?.1
        } else {
            AvPairs::default()
        };

        return Ok(Self {
            flags,
            server_challenge: server_challenge.try_into().unwrap(),
            version,
            target_name,
            target_info,
        });
    }

    /// Create the raw representation in bytes to be transmitted over the
    /// network.
    pub fn build(&self) -> Vec<u8> {
        let mut raw = NTLM_SIGN.to_vec();
        raw.extend(&NTLM_MSG_CHALLENGE.to_le_bytes());

        let mut offset = CHALLENGE_MSG_FIXED_SIZE;

        let raw_tn = match &self.target_name {
            Some(tn) => utils::to_neg_enc(tn, self.flags),
            None => Vec::new(),
        };
        let tn_offset = offset;
        offset += raw_tn.len() as u32;

        enc_len_offset!(raw, raw_tn, tn_offset);
        raw.extend(&self.flags.to_le_bytes());
        raw.extend(&self.server_challenge);

        // reserved, should be zeros
        raw.extend(&[0; 8]);

        let raw_ti = self.target_info.build();
        let ti_offset = offset;
        enc_len_offset!(raw, raw_ti, ti_offset);

        if let Some(v) = &self.version {
            raw.extend(v.build());
        }

        raw.extend(raw_tn);
        raw.extend(raw_ti);

        return raw;
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::flags::*;
    use crate::time::new_time;
    use crate::AvPair;

    const RAW_NTLM_CHALLENGE: &'static [u8] = &[
        0x4e, 0x54, 0x4c, 0x4d, 0x53, 0x53, 0x50, 0x00, 0x02, 0x00, 0x00, 0x00,
        0x0e, 0x00, 0x0e, 0x00, 0x38, 0x00, 0x00, 0x00, 0x05, 0x02, 0x89, 0xa2,
        0x99, 0x15, 0x63, 0xc0, 0x40, 0xe1, 0xf7, 0x08, 0x00, 0x00, 0x00, 0x00,
        0x00, 0x00, 0x00, 0x00, 0x9e, 0x00, 0x9e, 0x00, 0x46, 0x00, 0x00, 0x00,
        0x0a, 0x00, 0x61, 0x4a, 0x00, 0x00, 0x00, 0x0f, 0x43, 0x00, 0x4f, 0x00,
        0x4e, 0x00, 0x54, 0x00, 0x4f, 0x00, 0x53, 0x00, 0x4f, 0x00, 0x02, 0x00,
        0x0e, 0x00, 0x43, 0x00, 0x4f, 0x00, 0x4e, 0x00, 0x54, 0x00, 0x4f, 0x00,
        0x53, 0x00, 0x4f, 0x00, 0x01, 0x00, 0x0e, 0x00, 0x57, 0x00, 0x53, 0x00,
        0x30, 0x00, 0x31, 0x00, 0x2d, 0x00, 0x31, 0x00, 0x30, 0x00, 0x04, 0x00,
        0x1a, 0x00, 0x63, 0x00, 0x6f, 0x00, 0x6e, 0x00, 0x74, 0x00, 0x6f, 0x00,
        0x73, 0x00, 0x6f, 0x00, 0x2e, 0x00, 0x6c, 0x00, 0x6f, 0x00, 0x63, 0x00,
        0x61, 0x00, 0x6c, 0x00, 0x03, 0x00, 0x2a, 0x00, 0x77, 0x00, 0x73, 0x00,
        0x30, 0x00, 0x31, 0x00, 0x2d, 0x00, 0x31, 0x00, 0x30, 0x00, 0x2e, 0x00,
        0x63, 0x00, 0x6f, 0x00, 0x6e, 0x00, 0x74, 0x00, 0x6f, 0x00, 0x73, 0x00,
        0x6f, 0x00, 0x2e, 0x00, 0x6c, 0x00, 0x6f, 0x00, 0x63, 0x00, 0x61, 0x00,
        0x6c, 0x00, 0x05, 0x00, 0x1a, 0x00, 0x63, 0x00, 0x6f, 0x00, 0x6e, 0x00,
        0x74, 0x00, 0x6f, 0x00, 0x73, 0x00, 0x6f, 0x00, 0x2e, 0x00, 0x6c, 0x00,
        0x6f, 0x00, 0x63, 0x00, 0x61, 0x00, 0x6c, 0x00, 0x07, 0x00, 0x08, 0x00,
        0x8d, 0x0d, 0xad, 0x0f, 0xb4, 0xf0, 0xd6, 0x01, 0x00, 0x00, 0x00, 0x00,
    ];

    #[test]
    fn test_parse_challenge_msg() {
        let mut challenge = ChallengeMsg::default();
        challenge.target_name = Some("CONTOSO".to_string());
        challenge.flags = NTLM_NEG_56
            | NTLM_NEG_128
            | NTLM_NEG_VERSION
            | NTLM_NEG_TARGET_INFO
            | NTLM_NEG_EXTENDED_SECURITY
            | NTLM_NEG_TARGET_TYPE_DOMAIN
            | NTLM_NEG_NTLM
            | NTLM_REQUEST_TARGET
            | NTLM_NEG_UNICODE;

        challenge.server_challenge =
            [0x99, 0x15, 0x63, 0xc0, 0x40, 0xe1, 0xf7, 0x08];
        challenge.version = Some(Version {
            major: 10,
            minor: 0,
            build: 19041,
            revision: 15,
        });

        challenge.target_info = vec![
            AvPair::NbDomainName("CONTOSO".to_string()),
            AvPair::NbComputerName("WS01-10".to_string()),
            AvPair::DnsDomainName("contoso.local".to_string()),
            AvPair::DnsComputerName("ws01-10.contoso.local".to_string()),
            AvPair::DnsTreeName("contoso.local".to_string()),
            AvPair::Timestamp(new_time(2021, 01, 22, 11, 45, 20, 178727700)),
        ]
        .into();

        assert_eq!(
            challenge,
            ChallengeMsg::parse(RAW_NTLM_CHALLENGE).unwrap()
        );
    }

    #[test]
    fn test_build_challenge_msg() {
        let mut challenge = ChallengeMsg::default();
        challenge.target_name = Some("CONTOSO".to_string());
        challenge.flags = NTLM_NEG_56
            | NTLM_NEG_128
            | NTLM_NEG_VERSION
            | NTLM_NEG_TARGET_INFO
            | NTLM_NEG_EXTENDED_SECURITY
            | NTLM_NEG_TARGET_TYPE_DOMAIN
            | NTLM_NEG_NTLM
            | NTLM_REQUEST_TARGET
            | NTLM_NEG_UNICODE;

        challenge.server_challenge =
            [0x99, 0x15, 0x63, 0xc0, 0x40, 0xe1, 0xf7, 0x08];
        challenge.version = Some(Version {
            major: 10,
            minor: 0,
            build: 19041,
            revision: 15,
        });

        challenge.target_info = vec![
            AvPair::NbDomainName("CONTOSO".to_string()),
            AvPair::NbComputerName("WS01-10".to_string()),
            AvPair::DnsDomainName("contoso.local".to_string()),
            AvPair::DnsComputerName("ws01-10.contoso.local".to_string()),
            AvPair::DnsTreeName("contoso.local".to_string()),
            AvPair::Timestamp(new_time(2021, 01, 22, 11, 45, 20, 178727700)),
        ]
        .into();

        assert_eq!(RAW_NTLM_CHALLENGE.to_vec(), challenge.build(),);
    }
}