ad-time 1.0.1

/// SMB2 NEGOTIATE time source — fallback on TCP/445.
///
/// Protocol Specifications:
/// - **MS-SMB2 §2.2.3**: SMB2 NEGOTIATE Request
/// - **MS-SMB2 §2.2.4**: SMB2 NEGOTIATE Response
///
/// Sends SMB2 NEGOTIATE request and reads SystemTime from the response.
use std::io::{Read, Write};
use std::net::{SocketAddr, TcpStream};
use std::time::{Duration, Instant, SystemTime};

use super::common::{filetime_to_system_time, system_time_to_us};
use crate::time_src::{OffsetMicros, TimeSource, TimeSourceError};

pub struct SmbSource;

// SMB2 capabilities: DFS | LEASING | LARGE_MTU | MULTI_CHANNEL | PERSISTENT_HANDLES | DIR_LEASING | ENCRYPTION
const SMB2_CAPABILITIES: u32 = 0x7F;

/// Sequential field reader for little-endian binary structs.
struct FieldReader<'a> {
    buf: &'a [u8],
    pos: usize,
}

impl<'a> FieldReader<'a> {
    fn new(buf: &'a [u8]) -> Self {
        Self { buf, pos: 0 }
    }

    fn read_u16_le(&mut self) -> Result<u16, TimeSourceError> {
        let b = self.next_bytes(2)?;
        Ok(u16::from_le_bytes([b[0], b[1]]))
    }

    fn read_u32_le(&mut self) -> Result<u32, TimeSourceError> {
        let b = self.next_bytes(4)?;
        Ok(u32::from_le_bytes([b[0], b[1], b[2], b[3]]))
    }

    fn read_u64_le(&mut self) -> Result<u64, TimeSourceError> {
        let b = self.next_bytes(8)?;
        Ok(u64::from_le_bytes([
            b[0], b[1], b[2], b[3], b[4], b[5], b[6], b[7],
        ]))
    }

    fn skip(&mut self, n: usize) -> Result<(), TimeSourceError> {
        self.next_bytes(n)?;
        Ok(())
    }

    fn next_bytes(&mut self, n: usize) -> Result<&'a [u8], TimeSourceError> {
        let end = self
            .pos
            .checked_add(n)
            .ok_or_else(|| TimeSourceError::Parse("FieldReader overflow".into()))?;
        if end > self.buf.len() {
            return Err(TimeSourceError::Parse("SMB body overruns buffer".into()));
        }
        let b = &self.buf[self.pos..end];
        self.pos = end;
        Ok(b)
    }
}

impl TimeSource for SmbSource {
    fn name(&self) -> &'static str {
        "smb"
    }

    fn fetch(
        &self,
        target: SocketAddr,
        timeout: Duration,
    ) -> Result<OffsetMicros, TimeSourceError> {
        let smb_addr: SocketAddr = (target.ip(), 445).into();
        fetch_smb(smb_addr, timeout)
    }
}

fn fetch_smb(addr: SocketAddr, timeout: Duration) -> Result<OffsetMicros, TimeSourceError> {
    let mut stream = TcpStream::connect_timeout(&addr, timeout).map_err(map_io_err)?;
    stream
        .set_read_timeout(Some(timeout))
        .map_err(|e| TimeSourceError::Protocol(e.to_string()))?;
    stream
        .set_write_timeout(Some(timeout))
        .map_err(|e| TimeSourceError::Protocol(e.to_string()))?;

    let t_send = Instant::now();
    let t_send_sys = SystemTime::now();

    let request = build_negotiate_request();
    stream
        .write_all(&request)
        .map_err(|e| TimeSourceError::Protocol(e.to_string()))?;

    // Read NetBIOS header (4 bytes) to know response length.
    let mut nb_header = [0u8; 4];
    stream.read_exact(&mut nb_header).map_err(map_io_err)?;
    // NetBIOS session message: byte 0 = 0x00, bytes 1..4 = 24-bit big-endian length.
    let msg_len = u32::from_be_bytes(nb_header) & 0x00FF_FFFF;
    if msg_len > 65536 {
        return Err(TimeSourceError::Protocol(format!(
            "implausibly large SMB2 response: {} bytes",
            msg_len
        )));
    }
    if msg_len < 64 + 65 {
        return Err(TimeSourceError::Parse(format!(
            "SMB2 response too short: {} bytes",
            msg_len
        )));
    }

    let mut body = vec![0u8; msg_len as usize];
    stream.read_exact(&mut body).map_err(map_io_err)?;

    let rtt = t_send.elapsed();

    // body[0..64] is SMB2 header; body[64..] is NEGOTIATE_RESPONSE.
    let negotiate = &body[64..];
    let server_time = parse_negotiate_response(negotiate)?;

    // Single-point approximation: server timestamp ≈ midpoint of our send/recv window.
    // Precision: ±RTT/2 — sufficient for Kerberos 5-minute skew window.
    let t_mid_us = system_time_to_us(t_send_sys)? + (rtt.as_micros() as i64) / 2;
    let server_us = system_time_to_us(server_time)?;

    Ok(server_us - t_mid_us)
}

/// Build SMB2 NEGOTIATE request wrapped in a NetBIOS session message.
fn build_negotiate_request() -> Vec<u8> {
    // Dialects: SMB 3.0, 2.1, 2.0.2. Dropped 3.1.1 because it requires Negotiate Contexts to be OPSEC safe.
    let dialects: &[u16] = &[0x0300, 0x0210, 0x0202];
    let dialect_count = dialects.len() as u16;

    // SMB2 NEGOTIATE request body (MS-SMB2 §2.2.3):
    // StructureSize (2) + DialectCount (2) + SecurityMode (2) + Reserved (2) +
    // Capabilities (4) + ClientGuid (16) + ClientStartTime/NegotiateContextOffset/Count (8) +
    // Dialects (2*n)
    let body_size = 2 + 2 + 2 + 2 + 4 + 16 + 8 + (2 * dialect_count as usize);
    let smb2_header_size = 64usize;
    let total = smb2_header_size + body_size;

    let mut pkt = vec![0u8; 4 + total]; // 4-byte NetBIOS prefix

    // NetBIOS session message header (type=0x00, 24-bit big-endian length)
    pkt[1] = ((total >> 16) & 0xFF) as u8;
    pkt[2] = ((total >> 8) & 0xFF) as u8;
    pkt[3] = (total & 0xFF) as u8;

    let h = &mut pkt[4..4 + smb2_header_size];
    // ProtocolId: 0xFE 'S' 'M' 'B'
    h[0..4].copy_from_slice(b"\xfeSMB");
    // StructureSize = 64
    h[4..6].copy_from_slice(&64u16.to_le_bytes());
    // Command = NEGOTIATE (0x0000)
    h[12..14].copy_from_slice(&0u16.to_le_bytes());
    // Flags = 0 (client)
    // CreditRequest = 1
    h[18..20].copy_from_slice(&1u16.to_le_bytes());
    // MessageId = 1
    h[28..36].copy_from_slice(&1u64.to_le_bytes());

    let b = &mut pkt[4 + smb2_header_size..];
    // StructureSize = 36
    b[0..2].copy_from_slice(&36u16.to_le_bytes());
    // DialectCount
    b[2..4].copy_from_slice(&dialect_count.to_le_bytes());
    // SecurityMode = 1 (signing enabled, but not required - matches Windows default)
    b[4..6].copy_from_slice(&1u16.to_le_bytes());
    b[8..12].copy_from_slice(&SMB2_CAPABILITIES.to_le_bytes());
    // OPSEC: Random ClientGuid (UUIDv4)
    let mut guid = [0u8; 16];
    for b_out in guid.iter_mut() {
        *b_out = rand::random();
    }
    guid[6] = (guid[6] & 0x0F) | 0x40; // Version 4
    guid[8] = (guid[8] & 0x3F) | 0x80; // Variant 10xx
    b[12..28].copy_from_slice(&guid);
    // Dialects start at offset 36 from body start
    for (i, &d) in dialects.iter().enumerate() {
        let off = 36 + i * 2;
        b[off..off + 2].copy_from_slice(&d.to_le_bytes());
    }

    pkt
}

/// Parse SMB2 NEGOTIATE_RESPONSE (MS-SMB2 §2.2.4) and extract SystemTime.
fn parse_negotiate_response(b: &[u8]) -> Result<SystemTime, TimeSourceError> {
    let mut r = FieldReader::new(b);
    // Fields are little-endian; read sequentially per MS-SMB2 §2.2.4.
    let structure_size = r.read_u16_le()?; //  0: StructureSize (must be 65)
    let _security_mode = r.read_u16_le()?; //  2: SecurityMode
    let _dialect_revision = r.read_u16_le()?; //  4: DialectRevision
    let _negotiate_ctx_cnt = r.read_u16_le()?; //  6: NegotiateContextCount/Reserved
    r.skip(16)?; //  8: ServerGuid ([u8; 16])
    let _capabilities = r.read_u32_le()?; // 24: Capabilities
    let _max_transact = r.read_u32_le()?; // 28: MaxTransactSize
    let _max_read = r.read_u32_le()?; // 32: MaxReadSize
    let _max_write = r.read_u32_le()?; // 36: MaxWriteSize
    let system_time = r.read_u64_le()?; // 40: SystemTime (FILETIME)

    if structure_size != 65 {
        return Err(TimeSourceError::Protocol(format!(
            "unexpected SMB2 NEGOTIATE_RESPONSE StructureSize: {}",
            structure_size
        )));
    }

    filetime_to_system_time(system_time)
}

fn map_io_err(e: std::io::Error) -> TimeSourceError {
    use std::io::ErrorKind::*;
    match e.kind() {
        TimedOut | WouldBlock => TimeSourceError::Timeout,
        ConnectionRefused => TimeSourceError::Refused,
        _ => TimeSourceError::Protocol(e.to_string()),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::time::UNIX_EPOCH;

    #[test]
    fn filetime_unix_epoch() {
        // FILETIME of Unix epoch = 116444736000000000
        let ft: u64 = 116_444_736_000_000_000;
        let st = filetime_to_system_time(ft).unwrap();
        assert_eq!(st, UNIX_EPOCH);
    }

    #[test]
    fn filetime_2024_01_01() {
        // 2024-01-01 00:00:00 UTC as FILETIME
        // Unix timestamp = 1704067200
        // FILETIME = (1704067200 + 11644473600) * 10_000_000 = 133485408000000000
        let ft: u64 = 133_485_408_000_000_000;
        let st = filetime_to_system_time(ft).unwrap();
        let unix_secs = st.duration_since(UNIX_EPOCH).unwrap().as_secs();
        assert_eq!(unix_secs, 1_704_067_200);
    }

    #[test]
    fn filetime_before_unix_epoch_errors() {
        assert!(filetime_to_system_time(0).is_err());
        assert!(filetime_to_system_time(100).is_err());
    }

    #[test]
    fn negotiate_response_too_short() {
        assert!(parse_negotiate_response(&[0u8; 10]).is_err());
    }

    #[test]
    fn negotiate_response_bad_structure_size() {
        let mut b = vec![0u8; 50];
        // StructureSize = 99 (wrong)
        b[0..2].copy_from_slice(&99u16.to_le_bytes());
        assert!(parse_negotiate_response(&b).is_err());
    }

    #[test]
    fn build_negotiate_request_has_random_guid() {
        // ClientGuid is at offset 4 (NetBIOS) + 64 (SMB2 header) + 12 (body offset) = 80..96
        let r1 = build_negotiate_request();
        let r2 = build_negotiate_request();
        assert_ne!(
            &r1[80..96],
            &r2[80..96],
            "ClientGuid must differ between calls"
        );
        // Sanity: neither is all-zero (overwhelmingly probable)
        assert_ne!(&r1[80..96], &[0u8; 16]);
    }

    #[test]
    fn fetch_smb_rejects_large_msg_len() {
        // Simulate a NetBIOS header claiming a 128 KB response body (> 65536 limit).
        // We cannot call fetch_smb (needs a real socket), but we can verify the
        // guard arithmetic: msg_len field is 24-bit from bytes [1..4].
        let large: u32 = 0x0002_0000; // 131072 bytes
        assert!(large > 65536);
        // Confirm the mask used in production: u32::from_be_bytes([0, 2, 0, 0]) & 0x00FF_FFFF = 131072
        let nb = [0x00u8, 0x02, 0x00, 0x00];
        let msg_len = u32::from_be_bytes(nb) & 0x00FF_FFFF;
        assert_eq!(msg_len, 131072);
        assert!(msg_len > 65536);
    }
}