red_smb 0.0.2

use nom::bytes::complete::{tag, take};
use nom::number::complete::{le_u16, le_u32, le_u8};
use crate::Result;
use std::convert::TryInto;
use super::flags::SMB_FLAGS_REPLY;

pub const SMB1_ID: [u8; 4] = [0xff, 0x53, 0x4d, 0x42];

#[derive(Clone, Debug, PartialEq)]
pub struct Smb1Header {
    pub command: u8,
    pub status: u32,
    pub flags: u8,
    pub flags2: u16,
    pub pid_high: u16,
    pub security_features: [u8; 8],
    pub reserved: u16,
    pub tid: u16,
    pub pid_low: u16,
    pub uid: u16,
    pub mid: u16,
}

impl Smb1Header {

    pub fn new(command: u8) -> Self {
        return Self {
            command,
            status: 0,
            flags: 0,
            flags2: 0,
            pid_high: 0,
            security_features: [0; 8],
            reserved: 0,
            tid: 0,
            pid_low: 0,
            uid: 0,
            mid: 0,
        }
    }

    pub fn is_response(&self) -> bool {
        return (self.flags & SMB_FLAGS_REPLY) != 0;
    }

    pub fn is_request(&self) -> bool {
        return !self.is_response();
    }

    pub fn build(&self) -> Vec<u8> {
        let mut bytes = SMB1_ID.to_vec();

        bytes.push(self.command);

        bytes.append(&mut self.status.to_le_bytes().to_vec());
        bytes.push(self.flags);

        bytes.append(&mut self.flags2.to_le_bytes().to_vec());

        bytes.append(&mut self.pid_high.to_le_bytes().to_vec());
        bytes.append(&mut self.security_features.to_vec());
        bytes.append(&mut self.reserved.to_le_bytes().to_vec());
        bytes.append(&mut self.tid.to_le_bytes().to_vec());
        bytes.append(&mut self.pid_low.to_le_bytes().to_vec());
        bytes.append(&mut self.uid.to_le_bytes().to_vec());
        bytes.append(&mut self.mid.to_le_bytes().to_vec());

        return bytes;
    }

    pub fn parse(raw: &[u8]) -> Result<(&[u8], Self)> {
        let (raw, _) = tag(&SMB1_ID[..])(raw)?;
        let (raw, command) = le_u8(raw)?;
        let (raw, status) = le_u32(raw)?;
        let (raw, flags) = le_u8(raw)?;
        let (raw, flags2) = le_u16(raw)?;
        let (raw, pid_high) = le_u16(raw)?;
        let (raw, security_features) = take(8usize)(raw)?;
        let (raw, reserved) = le_u16(raw)?;
        let (raw, tid) = le_u16(raw)?;
        let (raw, pid_low) = le_u16(raw)?;
        let (raw, uid) = le_u16(raw)?;
        let (raw, mid) = le_u16(raw)?;

        return Ok((
            raw,
            Self {
                command,
                status,
                flags,
                flags2,
                pid_high,
                security_features: security_features.try_into().unwrap(),
                reserved,
                tid,
                pid_low,
                uid,
                mid,
            },
        ));
    }
}

impl Default for Smb1Header {
    fn default() -> Self {
        return Self {
            command: 0,
            status: 0,
            flags: 0,
            flags2: 0,
            pid_high: 0,
            security_features: [0; 8],
            reserved: 0,
            tid: 0xffff,
            pid_low: 0,
            uid: 0,
            mid: 0,
        };
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::smb1::header::flags::{
        SMB_FLAGS2_EXTENDED_SECURITY, SMB_FLAGS2_LONG_NAMES,
        SMB_FLAGS2_NT_STATUS, SMB_FLAGS_CANONICALIZED_PATHS,
        SMB_FLAGS_CASE_INSENSITIVE,
    };
    use crate::smb1::header::commands::SMB_COM_NEGOTIATE;

    const RAW_HEADER: &'static [u8] = &[
        0xff, 0x53, 0x4d, 0x42, 0x72, 0x00, 0x00, 0x00, 0x00, 0x18, 0x01, 0x48,
        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
        0xff, 0xff, 0x0c, 0x04, 0x00, 0x00, 0x00, 0x00,
    ];

    #[test]
    fn test_build_smb_header() {
        let mut header = Smb1Header::default();
        header.command = SMB_COM_NEGOTIATE;

        header.flags =
            SMB_FLAGS_CANONICALIZED_PATHS | SMB_FLAGS_CASE_INSENSITIVE;

        header.flags2 = SMB_FLAGS2_NT_STATUS
            | SMB_FLAGS2_LONG_NAMES
            | SMB_FLAGS2_EXTENDED_SECURITY;

        header.pid_low = 0x040c;

        assert_eq!(RAW_HEADER.to_vec(), header.build());
    }

    #[test]
    fn test_parse_smb_header() {
        let mut header = Smb1Header::default();
        header.command = SMB_COM_NEGOTIATE;

        header.flags =
            SMB_FLAGS_CANONICALIZED_PATHS | SMB_FLAGS_CASE_INSENSITIVE;

        header.flags2 = SMB_FLAGS2_NT_STATUS
            | SMB_FLAGS2_LONG_NAMES
            | SMB_FLAGS2_EXTENDED_SECURITY;

        header.pid_low = 0x040c;

        assert_eq!(header, Smb1Header::parse(&RAW_HEADER).unwrap().1);
    }
}