pwn_helper/

io.rs

use std::{
    io::{stdin, stdout, ErrorKind, Read, Write},
    net::TcpStream,
    sync::mpsc::TryRecvError,
    thread,
    time::Duration,
};
use thiserror::Error;

#[derive(Debug, Error)]
/// Errors that can be encountered during IO operations
pub enum PwnIoError {
    #[error("IO: {0}")]
    Io(std::io::Error),
    #[error("End of file")]
    Eof,
}

impl From<std::io::Error> for PwnIoError {
    fn from(x: std::io::Error) -> Self {
        match x.kind() {
            std::io::ErrorKind::UnexpectedEof => Self::Eof,
            _ => Self::Io(x),
        }
    }
}

/// Trait for read IO operations
pub trait PwnIoRead {
    /// Attempts to read from the stream until the pattern specified in `bytes` is reached, optionally trimming the pattern from the returned bytes
    fn receive_until(&mut self, pattern: &[u8], trim_pattern: bool) -> Result<Vec<u8>, PwnIoError>;

    /// Attempts to read from the stream until it reaches a newline character, optionally trimming the newline from the returned bytes
    fn receive_line(&mut self, trim_newline: bool) -> Result<Vec<u8>, PwnIoError>;

    /// Attempts to read all bytes currently available
    fn receive(&mut self) -> Result<Vec<u8>, PwnIoError>;

    /// Attempts to read all bytes until the stream ends
    fn receive_all(&mut self) -> Result<Vec<u8>, PwnIoError>;

    /// Attempts to read `count` bytes until the stream ends
    fn receive_count(&mut self, count: usize) -> Result<Vec<u8>, PwnIoError>;
}

impl<T> PwnIoRead for T
where
    T: Read,
{
    fn receive_until(&mut self, pattern: &[u8], trim_pattern: bool) -> Result<Vec<u8>, PwnIoError> {
        log::debug!("Receiving until: {:?}", pattern);
        let mut out = Vec::new();
        let mut buffer = [0u8];
        while !out.ends_with(pattern) {
            self.read_exact(&mut buffer)?;
            out.push(buffer[0]);
        }

        if trim_pattern {
            for _ in 0..pattern.len() {
                out.pop();
            }
        }
        Ok(out)
    }

    fn receive_line(&mut self, trim_newline: bool) -> Result<Vec<u8>, PwnIoError> {
        log::debug!("Receiving line");
        self.receive_until(b"\n", trim_newline)
    }

    fn receive(&mut self) -> Result<Vec<u8>, PwnIoError> {
        log::debug!("Receiving");
        let mut out = Vec::new();
        let mut buffer = [0u8];
        loop {
            if self.read(&mut buffer)? == 0 {
                return Ok(out);
            } else {
                out.push(buffer[0]);
            }
        }
    }

    fn receive_all(&mut self) -> Result<Vec<u8>, PwnIoError> {
        log::debug!("Receiving all");
        let mut out = Vec::new();
        let mut buffer = [0u8];
        loop {
            match self.read_exact(&mut buffer) {
                Ok(_) => out.push(buffer[0]),
                Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => return Ok(out),
                Err(e) => return Err(PwnIoError::Io(e)),
            }
        }
    }

    fn receive_count(&mut self, count: usize) -> Result<Vec<u8>, PwnIoError> {
        log::debug!("Receiving {:x} bytes", count);
        let mut out = Vec::with_capacity(count);
        self.read_exact(&mut out)?;
        Ok(out)
    }
}

/// Trait for write IO operations
pub trait PwnIoWrite {
    /// Attempts to write data to the stream
    fn send(&mut self, data: &[u8]) -> Result<(), std::io::Error>;
    /// Attempts to write data to the stream followed by a newline character
    fn send_line(&mut self, data: &[u8]) -> Result<(), std::io::Error>;
}

impl<T> PwnIoWrite for T
where
    T: Write,
{
    fn send_line(&mut self, data: &[u8]) -> Result<(), std::io::Error> {
        self.write_all(data)?;
        self.write_all(b"\n")?;
        Ok(())
    }

    fn send(&mut self, data: &[u8]) -> Result<(), std::io::Error> {
        self.write_all(data)?;
        Ok(())
    }
}

pub trait PwnIoInteractive {
    /// Attempts to enter into interactive mode via stdin/stdout
    fn interactive(&mut self) -> Result<(), std::io::Error>;
}

impl PwnIoInteractive for TcpStream {
    fn interactive(&mut self) -> Result<(), std::io::Error> {
        log::debug!("Entering interactive");
        let mut buffer = [0u8; 2048];
        let mut send_buffer: Option<u8> = None;

        let (s, r) = std::sync::mpsc::channel();

        thread::spawn(move || {
            let mut buffer = [0; 2048];
            let mut stdin = stdin();
            'outer_loop: loop {
                let read_amount = stdin.read(&mut buffer);

                match read_amount {
                    Ok(read_amount) if read_amount > 0 => {
                        for x in buffer.iter().copied().take(read_amount) {
                            s.send(Ok(x)).unwrap();
                        }
                        buffer = [0; 2048];
                    }
                    Ok(_) => {}
                    Err(x) => {
                        s.send(Err(x)).unwrap();
                        break 'outer_loop;
                    }
                };
            }
        });

        'outer_loop: loop {
            self.set_read_timeout(Some(Duration::from_millis(200)))?;
            self.set_write_timeout(Some(Duration::from_millis(200)))?;
            let read_amount = match self.read(&mut buffer) {
                Ok(x) => x,
                Err(x) if x.kind() == ErrorKind::WouldBlock => 0,
                Err(x) => return Err(x),
            };

            if read_amount > 0 {
                print!("{}", String::from_utf8_lossy(&buffer[0..read_amount]));
                stdout().flush()?;
            }

            match send_buffer {
                Some(x) => match self.write(&[x]) {
                    Ok(_) => {
                        send_buffer = None;
                    }
                    Err(err) if err.kind() == ErrorKind::WouldBlock => {}
                    Err(err) => return Err(err),
                },
                None => 'inner_loop: loop {
                    match r.try_recv() {
                        Ok(Ok(x)) => match self.write(&[x]) {
                            Ok(_) => {}
                            Err(err) if err.kind() == ErrorKind::WouldBlock => {
                                send_buffer = Some(x);
                                break 'inner_loop;
                            }
                            Err(err) => return Err(err),
                        },
                        Ok(Err(_)) | Err(TryRecvError::Disconnected) => {
                            break 'outer_loop;
                        }
                        Err(TryRecvError::Empty) => {
                            break 'inner_loop;
                        }
                    }
                },
            };
        }

        Ok(())
    }
}