rust_pwntools 0.1.0

A Rust crate inspired by Pwntools, providing powerful tools for binary exploitation, reverse engineering, and CTF challenges.
Documentation
use tokio::net::{TcpStream, UdpSocket};
use std::io::{self, Result};
use tokio::io::{AsyncReadExt, AsyncWriteExt};

pub struct Remote {
    tcp_stream: Option<TcpStream>,
    udp_socket: Option<UdpSocket>,
}

impl Remote {
    pub async fn connect_tcp(address: &str) -> Result<Self> {
        let tcp_stream = TcpStream::connect(address).await?;
        Ok(Remote {
            tcp_stream: Some(tcp_stream),
            udp_socket: None,
        })
    }

    pub async fn connect_udp(address: &str) -> Result<Self> {
        let udp_socket = UdpSocket::bind("0.0.0.0:0").await?;
        udp_socket.connect(address).await?;
        Ok(Remote {
            tcp_stream: None,
            udp_socket: Some(udp_socket),
        })
    }

    pub async fn send(&mut self, data: &[u8]) -> Result<()> {
        if let Some(ref mut stream) = self.tcp_stream {
            stream.write_all(data).await?;
        } else if let Some(ref mut socket) = self.udp_socket {
            socket.send(data).await?;
        }
        Ok(())
    }

    pub async fn receive(&mut self, buffer: &mut [u8]) -> Result<usize> {
        if let Some(ref mut stream) = self.tcp_stream {
            let n = stream.read(buffer).await?;
            Ok(n)
        } else if let Some(ref mut socket) = self.udp_socket {
            let (n, _) = socket.recv_from(buffer).await?;
            Ok(n)
        } else {
            Err(io::Error::new(io::ErrorKind::Other, "No connection established"))
        }
    }
}