red_smb 0.0.2

use crate::smb2::header::commands::SMB2_NEGOTIATE_PROTOCOL;
use crate::smb2::Smb2Header;
use crate::Result;
use nom::bytes::complete::{tag, take};
use nom::number::complete::{le_u16, le_u32, le_u64};
use std::convert::TryInto;

const SMB2_NEG_REQ_SIZE: u16 = 36;

#[derive(Clone, Debug, Default, PartialEq)]
pub struct Smb2NegReq {
    pub header: Smb2Header,
    pub body: Smb2NegReqBody,
}

impl Smb2NegReq {
    pub fn new() -> Self {
        return Self {
            header: Smb2Header::new(SMB2_NEGOTIATE_PROTOCOL),
            body: Smb2NegReqBody::default(),
        };
    }

    pub fn parse(raw: &[u8]) -> Result<Self> {
        let (raw, header) = Smb2Header::parse(raw)?;
        let (_, body) = Smb2NegReqBody::parse(raw)?;

        return Ok(Self { header, body });
    }

    pub fn build(&self) -> Vec<u8> {
        let mut raw = self.header.build();
        raw.extend(self.body.build());
        return raw;
    }
}

#[derive(Clone, Debug, Default, PartialEq)]
pub struct Smb2NegReqBody {
    pub security_mode: u16,
    pub capabilities: u32,
    pub client_start_time: u64,
    pub client_guid: [u8; 16],
    pub dialects: Vec<u16>,
}

impl Smb2NegReqBody {
    pub fn build(&self) -> Vec<u8> {
        let mut raw = SMB2_NEG_REQ_SIZE.to_le_bytes().to_vec();

        raw.extend(&(self.dialects.len() as u16).to_le_bytes());
        raw.extend(&self.security_mode.to_le_bytes());

        // reserved
        raw.extend(&[0; 2]);

        raw.extend(&self.capabilities.to_le_bytes());
        raw.extend(&self.client_guid);

        // ClientStartTime MUST be set to 0
        raw.extend(&self.client_start_time.to_le_bytes());

        for dialect in &self.dialects {
            raw.extend(&dialect.to_le_bytes());
        }

        return raw;
    }

    pub fn parse(raw: &[u8]) -> Result<(&[u8], Self)> {
        let (raw, _) = tag(&SMB2_NEG_REQ_SIZE.to_le_bytes()[..])(raw)?;

        let (raw, dialect_count) = le_u16(raw)?;
        let (raw, security_mode) = le_u16(raw)?;
        let (raw, _) = le_u16(raw)?;
        let (raw, capabilities) = le_u32(raw)?;

        let (raw, client_guid) = take(16usize)(raw)?;
        let (mut raw, client_start_time) = le_u64(raw)?;

        let mut dialects = Vec::new();
        for _ in 0..dialect_count {
            let (r, dialect) = le_u16(raw)?;
            raw = r;
            dialects.push(dialect);
        }

        return Ok((
            raw,
            Self {
                security_mode,
                capabilities,
                client_start_time,
                client_guid: client_guid.try_into().unwrap(),
                dialects,
            },
        ));
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::smb2::negotiate::SMB2_DIA_202;
    use crate::smb2::negotiate::SMB2_DIA_210;
    use crate::smb2::negotiate::SMB2_DIA_300;
    use crate::smb2::negotiate::SMB2_GLOBAL_CAP_ENCRYPTION;
    use crate::smb2::negotiate::SMB2_NEGOTIATE_SIGNING_ENABLED;

    const RAW_SMB2_NEG_REQ_BODY: &'static [u8] = &[
        0x24, 0x00, 0x03, 0x00, 0x01, 0x00, 0x00, 0x00, 0x40, 0x00, 0x00, 0x00,
        0x4f, 0x49, 0x7a, 0x4d, 0x6c, 0x4d, 0x4f, 0x43, 0x59, 0x77, 0x69, 0x42,
        0x69, 0x54, 0x67, 0x76, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
        0x02, 0x02, 0x10, 0x02, 0x00, 0x03,
    ];

    #[test]
    fn test_build_smb2_neg_req() {
        let mut neg_req = Smb2NegReqBody::default();
        neg_req.security_mode = SMB2_NEGOTIATE_SIGNING_ENABLED;

        neg_req.capabilities = SMB2_GLOBAL_CAP_ENCRYPTION;

        neg_req.client_guid = [
            0x4f, 0x49, 0x7a, 0x4d, 0x6c, 0x4d, 0x4f, 0x43, 0x59, 0x77, 0x69,
            0x42, 0x69, 0x54, 0x67, 0x76,
        ];

        neg_req.dialects = vec![SMB2_DIA_202, SMB2_DIA_210, SMB2_DIA_300];

        assert_eq!(RAW_SMB2_NEG_REQ_BODY.to_vec(), neg_req.build());
    }

    #[test]
    fn test_parse_smb2_neg_req() {
        let mut neg_req = Smb2NegReqBody::default();
        neg_req.security_mode = SMB2_NEGOTIATE_SIGNING_ENABLED;

        neg_req.capabilities = SMB2_GLOBAL_CAP_ENCRYPTION;

        neg_req.client_guid = [
            0x4f, 0x49, 0x7a, 0x4d, 0x6c, 0x4d, 0x4f, 0x43, 0x59, 0x77, 0x69,
            0x42, 0x69, 0x54, 0x67, 0x76,
        ];

        neg_req.dialects = vec![SMB2_DIA_202, SMB2_DIA_210, SMB2_DIA_300];

        assert_eq!(
            neg_req,
            Smb2NegReqBody::parse(RAW_SMB2_NEG_REQ_BODY).unwrap().1
        );
    }
}