telnet_wrapper 0.1.0

//! # `telnet_wrapper`
//!
//! A telnet server that wraps terminal applications, exposing them to telnet clients
//! with proper protocol support including window size negotiation (NAWS).
//!
//! ## Features
//!
//! - Wraps any terminal application and exposes it via telnet
//! - Full telnet protocol support using libtelnet-rs
//! - Window size negotiation (NAWS) - automatically resizes the PTY when client window changes
//! - Character-by-character input (no line buffering)
//! - Supports multiple concurrent clients
//!
//! ## Example
//!
//! ```bash
//! # Expose htop on port 2323
//! telnet_wrapper --port 2323 /usr/bin/htop
//!
//! # Connect from another terminal
//! telnet localhost 2323
//! ```

#![warn(
    clippy::all,
    clippy::pedantic,
    clippy::nursery,
    clippy::unwrap_used,
    clippy::panic,
    clippy::dbg_macro,
    clippy::missing_const_for_fn,
    clippy::needless_pass_by_value,
    clippy::redundant_pub_crate
)]
#![allow(
    clippy::missing_errors_doc,
    clippy::must_use_candidate,
    clippy::multiple_crate_versions,
    clippy::missing_panics_doc,
    clippy::option_if_let_else
)]

use bytes::Bytes;
use clap::Parser as ClapParser;
use libtelnet_rs::compatibility::CompatibilityTable;
use libtelnet_rs::events::TelnetEvents;
use libtelnet_rs::telnet::op_option as opt;
use libtelnet_rs::Parser as TelnetParser;
use portable_pty::{CommandBuilder, MasterPty, NativePtySystem, PtyPair, PtySize, PtySystem};
use std::io::{Read, Write};
use std::net::{SocketAddr, TcpStream};
use std::sync::{Arc, Mutex};
use std::thread::JoinHandle;
use tokio::net::TcpListener;

#[derive(ClapParser, Debug)]
#[command(author, version, about, long_about = None)]
struct Cli {
    /// The application command to run within the PTY
    app_command: String,

    /// Arguments to pass to the application command
    #[arg(allow_hyphen_values = true)]
    app_args: Vec<String>,

    /// Initial number of rows for the PTY
    #[arg(long, default_value = "24")]
    rows: u16,

    /// Initial number of columns for the PTY
    #[arg(long, default_value = "80")]
    cols: u16,

    /// Port to listen on for telnet connections
    #[arg(long, default_value = "23")]
    port: u16,
}

/// Send telnet events to the TCP socket
fn send_events(writer: &mut TcpStream, events: Vec<TelnetEvents>) {
    for event in events {
        if let TelnetEvents::DataSend(data) = event {
            let _ = writer.write_all(&data);
        }
    }
    let _ = writer.flush();
}

/// Parse NAWS subnegotiation data and return (cols, rows)
fn parse_naws_data(data: &Bytes) -> Option<(u16, u16)> {
    if data.len() >= 4 {
        let cols = u16::from_be_bytes([data[0], data[1]]);
        let rows = u16::from_be_bytes([data[2], data[3]]);
        Some((cols, rows))
    } else {
        None
    }
}

/// Scan raw bytes for NAWS subnegotiation sequences and return (cols, rows) if found
/// Format: IAC(255) SB(250) NAWS(31) <col-high> <col-low> <row-high> <row-low> IAC(255) SE(240)
fn scan_for_naws(data: &[u8]) -> Option<(u16, u16)> {
    // Need at least 9 bytes for a complete NAWS sequence
    if data.len() < 9 {
        return None;
    }
    // Look for IAC SB NAWS pattern
    for i in 0..=data.len() - 9 {
        if data[i] == 255
            && data[i + 1] == 250
            && data[i + 2] == 31
            && data[i + 7] == 255
            && data[i + 8] == 240
        {
            let cols = u16::from_be_bytes([data[i + 3], data[i + 4]]);
            let rows = u16::from_be_bytes([data[i + 5], data[i + 6]]);
            return Some((cols, rows));
        }
    }
    None
}

/// Resize the PTY to the given dimensions
fn resize_pty(master: &Arc<Mutex<Box<dyn MasterPty + Send>>>, cols: u16, rows: u16) {
    if let Ok(master) = master.lock() {
        let _ = master.resize(PtySize {
            rows,
            cols,
            pixel_width: 0,
            pixel_height: 0,
        });
    }
}

/// Open a PTY with the given dimensions
fn open_pty(addr: SocketAddr, cols: u16, rows: u16) -> Option<PtyPair> {
    let pty_system = NativePtySystem::default();
    eprintln!("[{addr}] Opening PTY with size {cols}x{rows}.");

    match pty_system.openpty(PtySize {
        rows,
        cols,
        pixel_width: 0,
        pixel_height: 0,
    }) {
        Ok(pair) => {
            eprintln!("[{addr}] PTY opened successfully.");
            Some(pair)
        }
        Err(e) => {
            eprintln!("[{addr}] Failed to open PTY: {e}");
            None
        }
    }
}

/// Spawn the child process in the PTY
fn spawn_child(
    addr: SocketAddr,
    pair: &PtyPair,
    app_cmd: &str,
    app_args: &[String],
) -> Option<Box<dyn portable_pty::Child + Send + Sync>> {
    let mut cmd = CommandBuilder::new(app_cmd);
    cmd.args(app_args.iter().cloned());
    cmd.env("TERM", "xterm-256color");
    eprintln!("[{addr}] Preparing to spawn command: {app_cmd:?} {app_args:?}");

    match pair.slave.spawn_command(cmd) {
        Ok(child) => {
            eprintln!("[{addr}] App command spawned successfully.");
            Some(child)
        }
        Err(e) => {
            eprintln!("[{addr}] Failed to spawn app command: {e}");
            None
        }
    }
}

/// Send initial telnet negotiations
fn send_telnet_negotiations(addr: SocketAddr, writer: &mut TcpStream) {
    let mut compat_table = CompatibilityTable::default();
    compat_table.support_local(opt::ECHO);
    compat_table.support_local(opt::SGA);
    compat_table.support_remote(opt::NAWS);

    let mut telnet_parser = TelnetParser::with_support(compat_table);

    eprintln!("[{addr}] Sending Telnet negotiations.");

    if let Some(event) = telnet_parser._will(opt::ECHO) {
        send_events(writer, vec![event]);
    }

    if let Some(event) = telnet_parser._will(opt::SGA) {
        send_events(writer, vec![event]);
    }

    if let Some(event) = telnet_parser._do(opt::NAWS) {
        send_events(writer, vec![event]);
    }

    if let Some(event) = telnet_parser._dont(opt::LINEMODE) {
        send_events(writer, vec![event]);
    }

    eprintln!("[{addr}] Telnet negotiations sent.");
}

/// Start the output thread (PTY -> TCP)
fn start_output_thread(
    addr: SocketAddr,
    mut pty_reader: Box<dyn Read + Send>,
    mut writer: TcpStream,
) -> JoinHandle<()> {
    std::thread::spawn(move || {
        eprintln!("[{addr}] Output thread started (PTY -> TCP).");
        let mut buf = [0u8; 8192];
        loop {
            match pty_reader.read(&mut buf) {
                Ok(0) => {
                    eprintln!("[{addr}] PTY reader returned 0 bytes (EOF).");
                    break;
                }
                Ok(n) => {
                    let data: Vec<u8> = buf[..n].to_vec();
                    let escaped = TelnetParser::escape_iac(data);
                    if let Err(e) = writer.write_all(&escaped) {
                        eprintln!("[{addr}] Error writing to socket in output thread: {e}");
                        break;
                    }
                    if let Err(e) = writer.flush() {
                        eprintln!("[{addr}] Error flushing socket in output thread: {e}");
                        break;
                    }
                }
                Err(e) => {
                    eprintln!("[{addr}] Error reading from PTY in output thread: {e}");
                    break;
                }
            }
        }
        eprintln!("[{addr}] Output thread finished.");
    })
}

/// Start the input thread (TCP -> PTY)
fn start_input_thread(
    addr: SocketAddr,
    mut reader: TcpStream,
    mut pty_writer: Box<dyn Write + Send>,
    mut response_writer: TcpStream,
    master: Arc<Mutex<Box<dyn MasterPty + Send>>>,
) -> JoinHandle<()> {
    std::thread::spawn(move || {
        eprintln!("[{addr}] Input thread started (TCP -> PTY).");
        let mut buf = [0u8; 1024];

        let mut input_compat = CompatibilityTable::default();
        input_compat.support_remote(opt::NAWS);
        input_compat.support_local(opt::ECHO);
        input_compat.support_local(opt::SGA);
        let mut input_parser = TelnetParser::with_support(input_compat);

        loop {
            match reader.read(&mut buf) {
                Ok(0) => {
                    eprintln!("[{addr}] Socket reader returned 0 bytes (EOF).");
                    break;
                }
                Ok(n) => {
                    if let Some((cols, rows)) = scan_for_naws(&buf[..n]) {
                        eprintln!("[{addr}] NAWS detected in raw bytes: {cols}x{rows}");
                        resize_pty(&master, cols, rows);
                    }

                    let events = input_parser.receive(&buf[..n]);

                    for event in events {
                        match event {
                            TelnetEvents::DataReceive(data) => {
                                if let Err(e) = pty_writer.write_all(&data) {
                                    eprintln!("[{addr}] Error writing to PTY: {e}");
                                    return;
                                }
                            }
                            TelnetEvents::Subnegotiation(subneg) => {
                                if subneg.option == opt::NAWS {
                                    if let Some((cols, rows)) = parse_naws_data(&subneg.buffer) {
                                        eprintln!("[{addr}] NAWS received: {cols}x{rows}");
                                        resize_pty(&master, cols, rows);
                                    }
                                }
                            }
                            TelnetEvents::DataSend(data) => {
                                if let Err(e) = response_writer.write_all(&data) {
                                    eprintln!("[{addr}] Error sending response: {e}");
                                }
                                let _ = response_writer.flush();
                            }
                            TelnetEvents::Negotiation(_)
                            | TelnetEvents::IAC(_)
                            | TelnetEvents::DecompressImmediate(_) => {}
                        }
                    }
                }
                Err(e) => {
                    eprintln!("[{addr}] Error reading from socket: {e}");
                    break;
                }
            }
        }
        eprintln!("[{addr}] Input thread finished.");
    })
}

/// Handle a single client connection
async fn handle_client(
    socket: tokio::net::TcpStream,
    addr: SocketAddr,
    app_cmd: String,
    app_args: Vec<String>,
    initial_cols: u16,
    initial_rows: u16,
) {
    eprintln!("[{addr}] Spawning task for client.");

    let Some(pair) = open_pty(addr, initial_cols, initial_rows) else {
        return;
    };

    let Some(mut child) = spawn_child(addr, &pair, &app_cmd, &app_args) else {
        return;
    };

    let master: Arc<Mutex<Box<dyn MasterPty + Send>>> = Arc::new(Mutex::new(pair.master));

    let Ok(std_socket) = socket.into_std() else {
        eprintln!("[{addr}] Failed to convert socket to std");
        return;
    };

    if std_socket.set_nonblocking(false).is_err() {
        eprintln!("[{addr}] Failed to set socket to blocking");
        return;
    }

    let Ok(mut writer) = std_socket.try_clone() else {
        eprintln!("[{addr}] Failed to clone writer");
        return;
    };
    let reader = std_socket;

    send_telnet_negotiations(addr, &mut writer);

    let Ok(pty_reader) = master
        .lock()
        .map_err(|_| ())
        .and_then(|m| m.try_clone_reader().map_err(|_| ()))
    else {
        eprintln!("[{addr}] Failed to get PTY reader");
        return;
    };

    let Ok(pty_writer) = master
        .lock()
        .map_err(|_| ())
        .and_then(|m| m.take_writer().map_err(|_| ()))
    else {
        eprintln!("[{addr}] Failed to get PTY writer");
        return;
    };

    let Ok(writer_for_output) = writer.try_clone() else {
        eprintln!("[{addr}] Failed to clone writer for output");
        return;
    };

    let t1 = start_output_thread(addr, pty_reader, writer_for_output);
    let t2 = start_input_thread(addr, reader, pty_writer, writer, Arc::clone(&master));

    eprintln!("[{addr}] Starting persistence check.");
    loop {
        if let Ok(Some(status)) = child.try_wait() {
            eprintln!("[{addr}] Child process exited with status: {status:?}");
            break;
        }
        if t1.is_finished() && t2.is_finished() {
            eprintln!("[{addr}] Both I/O threads finished.");
            break;
        }
        tokio::time::sleep(std::time::Duration::from_millis(100)).await;
    }
    eprintln!("[{addr}] Persistence check finished. Killing child process if still alive.");
    let _ = child.kill();
    eprintln!("[{addr}] Client handling task finished.");
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    eprintln!("Telnet Gateway starting...");
    let cli = Cli::parse();
    let port = cli.port;
    let listener = TcpListener::bind(format!("0.0.0.0:{port}")).await?;
    eprintln!("Telnet Gateway active on port {port}. Waiting for connections...");

    let initial_rows = cli.rows;
    let initial_cols = cli.cols;

    loop {
        let (socket, addr) = listener.accept().await?;
        eprintln!("New client connected from: {addr}");
        let app_cmd = cli.app_command.clone();
        let app_args = cli.app_args.clone();

        tokio::spawn(handle_client(
            socket,
            addr,
            app_cmd,
            app_args,
            initial_cols,
            initial_rows,
        ));
    }
}