//! Minimal Modbus/TCP client for dynamic ICS probing.
//!
//! Modbus/TCP (port 502) has no authentication or encryption in the protocol, so
//! an endpoint that answers requests is, by design, open to any host that can
//! reach it. The probe only *reads* — a Read Holding Registers request and a Read
//! Device Identification request — and never writes to the live process.

use std::time::Duration;

use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpStream;
use tokio::time::timeout;

/// Outcome of probing a Modbus/TCP endpoint.
#[derive(Debug, Default, PartialEq, Eq)]
pub struct ModbusProbe {
    /// A TCP connection to the port was established.
    pub reachable: bool,
    /// The endpoint answered a Modbus request (a normal reply or a Modbus
    /// exception) — i.e. it speaks Modbus, unauthenticated.
    pub speaks_modbus: bool,
    /// Device identity, if disclosed via Read Device Identification (FC 43 / 14).
    pub device: Option<DeviceId>,
    /// Coils returned by a Read Coils of the first block, if that address range
    /// exists. Coils are read/write process bits, so an exposed block is an
    /// unauthenticated write surface on the live process.
    pub coils_readable: Option<u16>,
    /// Holding registers returned by a Read Holding Registers of the first block,
    /// if that range exists. Holding registers are read/write process words.
    pub holding_registers_readable: Option<u16>,
}

/// Vendor / product / revision from Read Device Identification.
#[derive(Debug, Default, PartialEq, Eq)]
pub struct DeviceId {
    pub vendor: Option<String>,
    pub product: Option<String>,
    pub revision: Option<String>,
}

/// How many coils / holding registers to request when enumerating the exposed
/// process surface. Read-only: a normal reply means the block exists and is,
/// over unauthenticated Modbus/TCP, also writable.
const ENUM_QTY: u16 = 16;

/// Probe a Modbus/TCP endpoint. Read-only: issues Read Holding Registers, Read
/// Coils, and Read Device Identification requests; never writes to the device.
pub async fn probe(host: &str, port: u16, budget: Duration) -> ModbusProbe {
    let mut out = ModbusProbe::default();
    let Ok(Ok(mut stream)) = timeout(budget, TcpStream::connect((host, port))).await else {
        return out; // unreachable
    };
    out.reachable = true;

    // Read Holding Registers (FC 0x03), unit 1, addr 0 — a benign read that also
    // enumerates the exposed register block.
    let rhr = [0x03u8, 0x00, 0x00, (ENUM_QTY >> 8) as u8, ENUM_QTY as u8];
    if let Some(resp) = txn(&mut stream, 1, &rhr, budget).await {
        // A normal reply (0x03) or an exception (0x83) both prove it speaks Modbus.
        if matches!(resp.first(), Some(0x03) | Some(0x83)) {
            out.speaks_modbus = true;
        }
        if resp.first() == Some(&0x03) {
            out.holding_registers_readable = Some(register_count_from_reply(&resp));
        }
    }

    // Read Coils (FC 0x01), addr 0 — enumerates the exposed coil (bit) block.
    let rc = [0x01u8, 0x00, 0x00, (ENUM_QTY >> 8) as u8, ENUM_QTY as u8];
    if let Some(resp) = txn(&mut stream, 1, &rc, budget).await {
        if matches!(resp.first(), Some(0x01) | Some(0x81)) {
            out.speaks_modbus = true;
        }
        if resp.first() == Some(&0x01) {
            out.coils_readable = Some(coil_count_from_reply(&resp));
        }
    }

    // Read Device Identification (FC 0x2B / MEI 0x0E), basic (0x01), object 0.
    let rdi = [0x2Bu8, 0x0E, 0x01, 0x00];
    if let Some(resp) = txn(&mut stream, 1, &rdi, budget).await {
        if resp.first() == Some(&0x2B) {
            out.speaks_modbus = true;
            out.device = parse_device_id(&resp);
        }
    }
    out
}

/// Coils reported by a Read Coils reply `[0x01, byte_count, data…]` (8 per byte).
fn coil_count_from_reply(pdu: &[u8]) -> u16 {
    pdu.get(1).map(|&b| u16::from(b) * 8).unwrap_or(0)
}

/// Registers reported by a Read Holding Registers reply `[0x03, byte_count,
/// data…]` (2 bytes per register).
fn register_count_from_reply(pdu: &[u8]) -> u16 {
    pdu.get(1).map(|&b| u16::from(b) / 2).unwrap_or(0)
}

/// Send one Modbus PDU and return the response PDU (function code + data), or
/// `None` on timeout / malformed reply.
async fn txn(stream: &mut TcpStream, unit: u8, pdu: &[u8], budget: Duration) -> Option<Vec<u8>> {
    // MBAP header: transaction id (2) + protocol id (2) = 0 + length (2) + unit (1),
    // then the PDU. `length` counts the unit byte plus the PDU.
    let len = (pdu.len() + 1) as u16;
    let mut frame = Vec::with_capacity(7 + pdu.len());
    frame.extend_from_slice(&[0x00, 0x01]); // transaction id
    frame.extend_from_slice(&[0x00, 0x00]); // protocol id
    frame.extend_from_slice(&len.to_be_bytes());
    frame.push(unit);
    frame.extend_from_slice(pdu);
    timeout(budget, stream.write_all(&frame)).await.ok()?.ok()?;

    let mut hdr = [0u8; 7];
    timeout(budget, stream.read_exact(&mut hdr))
        .await
        .ok()?
        .ok()?;
    // Reject non-Modbus replies (protocol id must be 0).
    if hdr[2] != 0 || hdr[3] != 0 {
        return None;
    }
    let plen = u16::from_be_bytes([hdr[4], hdr[5]]) as usize;
    if !(2..=260).contains(&plen) {
        return None;
    }
    let mut body = vec![0u8; plen - 1]; // minus the unit id already in hdr[6]
    timeout(budget, stream.read_exact(&mut body))
        .await
        .ok()?
        .ok()?;
    Some(body)
}

/// Parse vendor / product / revision from a Read Device Identification PDU:
/// `[0x2B, 0x0E, readDevIdCode, conformity, moreFollows, nextObjId, numObjects,
/// (objId, len, bytes…)…]`.
fn parse_device_id(pdu: &[u8]) -> Option<DeviceId> {
    if pdu.len() < 7 {
        return None;
    }
    let num = pdu[6] as usize;
    let mut i = 7;
    let mut dev = DeviceId::default();
    for _ in 0..num {
        if i + 2 > pdu.len() {
            break;
        }
        let id = pdu[i];
        let l = pdu[i + 1] as usize;
        i += 2;
        if i + l > pdu.len() {
            break;
        }
        let val = String::from_utf8_lossy(&pdu[i..i + l]).trim().to_string();
        i += l;
        match id {
            0x00 => dev.vendor = Some(val),
            0x01 => dev.product = Some(val),
            0x02 => dev.revision = Some(val),
            _ => {}
        }
    }
    if dev == DeviceId::default() {
        None
    } else {
        Some(dev)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tokio::net::TcpListener;

    /// A one-shot mock Modbus/TCP server that answers a Read Holding Registers
    /// request and a Read Device Identification request on one connection.
    async fn mock_server(with_device: bool) -> std::net::SocketAddr {
        let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind");
        let addr = listener.local_addr().expect("addr");
        tokio::spawn(async move {
            let (mut sock, _) = listener.accept().await.expect("accept");
            loop {
                let mut hdr = [0u8; 7];
                if sock.read_exact(&mut hdr).await.is_err() {
                    break;
                }
                let plen = u16::from_be_bytes([hdr[4], hdr[5]]) as usize;
                let mut pdu = vec![0u8; plen - 1];
                if sock.read_exact(&mut pdu).await.is_err() {
                    break;
                }
                let reply_pdu: Vec<u8> = match pdu.first() {
                    Some(0x03) => vec![0x03, 0x02, 0x00, 0x00], // 1 register (byte_count 2)
                    Some(0x01) => vec![0x01, 0x02, 0xFF, 0xFF], // 16 coils (byte_count 2)
                    Some(0x2B) if with_device => vec![
                        0x2B, 0x0E, 0x01, 0x81, 0x00, 0x00, 0x02, // 2 objects
                        0x00, 0x04, b'A', b'C', b'M', b'E', // vendor
                        0x01, 0x03, b'P', b'L', b'C', // product
                    ],
                    _ => vec![pdu[0] | 0x80, 0x01], // exception
                };
                let len = (reply_pdu.len() + 1) as u16;
                let mut frame = vec![hdr[0], hdr[1], 0x00, 0x00];
                frame.extend_from_slice(&len.to_be_bytes());
                frame.push(hdr[6]);
                frame.extend_from_slice(&reply_pdu);
                if sock.write_all(&frame).await.is_err() {
                    break;
                }
            }
        });
        addr
    }

    #[tokio::test]
    async fn probe_detects_a_modbus_endpoint_and_reads_device_id() {
        let addr = mock_server(true).await;
        let p = probe(&addr.ip().to_string(), addr.port(), Duration::from_secs(2)).await;
        assert!(p.reachable && p.speaks_modbus);
        let dev = p.device.expect("device id");
        assert_eq!(dev.vendor.as_deref(), Some("ACME"));
        assert_eq!(dev.product.as_deref(), Some("PLC"));
    }

    #[tokio::test]
    async fn probe_enumerates_exposed_process_points() {
        let addr = mock_server(false).await;
        let p = probe(&addr.ip().to_string(), addr.port(), Duration::from_secs(2)).await;
        assert!(p.speaks_modbus);
        // The mock returns a 2-byte holding-register block (1 register) and a
        // 2-byte coil block (16 coils).
        assert_eq!(p.holding_registers_readable, Some(1));
        assert_eq!(p.coils_readable, Some(16));
    }

    #[test]
    fn reply_counts_decode_byte_counts() {
        assert_eq!(register_count_from_reply(&[0x03, 0x08]), 4); // 8 bytes → 4 regs
        assert_eq!(coil_count_from_reply(&[0x01, 0x03]), 24); // 3 bytes → 24 coils
        assert_eq!(register_count_from_reply(&[0x03]), 0); // malformed → 0
    }

    #[tokio::test]
    async fn probe_reports_unreachable_for_a_closed_port() {
        // 127.0.0.1:1 is (almost certainly) closed.
        let p = probe("127.0.0.1", 1, Duration::from_millis(500)).await;
        assert!(!p.reachable && !p.speaks_modbus);
    }

    #[test]
    fn parses_device_identification_objects() {
        let pdu = [
            0x2B, 0x0E, 0x01, 0x81, 0x00, 0x00, 0x01, // 1 object
            0x02, 0x05, b'v', b'1', b'.', b'2', b'3', // revision
        ];
        let dev = parse_device_id(&pdu).expect("device");
        assert_eq!(dev.revision.as_deref(), Some("v1.23"));
        assert!(dev.vendor.is_none());
    }
}