Add timeout to pivot.ncat

docs/_docs/user-guide/eldritch.md

@@ -209,7 +209,7 @@ Key must be 16 Bytes (Characters)
 
 ### crypto.encode_b64
 
-`crypto.encode_b64(content: str, encode_type: Optional<str>) -> str`
+`crypto.encode_b64(content: str, encode_type: Option<str>) -> str`
 
 The <b>crypto.encode_b64</b> method encodes the given text using the given base64 encoding method. Valid methods include:
 
@@ -220,7 +220,7 @@ The <b>crypto.encode_b64</b> method encodes the given text using the given base6
 
 ### crypto.decode_b64
 
-`crypto.decode_b64(content: str, decode_type: Optional<str>) -> str`
+`crypto.decode_b64(content: str, decode_type: Option<str>) -> str`
 
 The <b>crypto.decode_b64</b> method encodes the given text using the given base64 decoding method. Valid methods include:
 
@@ -563,10 +563,12 @@ The <b>pivot.bind_proxy</b> method is being proposed to provide users another op
 
 ### pivot.ncat
 
-`pivot.ncat(address: str, port: int, data: str, protocol: str ) -> str`
+`pivot.ncat(address: str, port: int, data: str, protocol: str, timeout: Option<int>) -> str`
 
 The <b>pivot.ncat</b> method allows a user to send arbitrary data over TCP/UDP to a host. If the server responds that response will be returned.
 
+Timeout is the number of seconds to wait for the ncat connection to succeed before closing it. Defaults to 2 seconds if not specified.
+
 `protocol` must be `tcp`, or `udp` anything else will return an error `Protocol not supported please use: udp or tcp.`.
 
 ### pivot.port_forward
@@ -614,7 +616,7 @@ NOTE: Windows scans against `localhost`/`127.0.0.1` can behave unexpectedly or e
 
 ### pivot.reverse_shell_pty
 
-`pivot.reverse_shell_pty(cmd: Optional<str>) -> None`
+`pivot.reverse_shell_pty(cmd: Option<str>) -> None`
 
 The **pivot.reverse_shell_pty** method spawns the provided command in a cross-platform PTY and opens a reverse shell over the agent's current transport (e.g. gRPC). If no command is provided, Windows will use `cmd.exe`. On other platforms, `/bin/bash` is used as a default, but if it does not exist then `/bin/sh` is used.
 
@@ -636,7 +638,7 @@ The file directory the `dst` file exists in must exist in order for ssh_copy to
 
 ### pivot.ssh_exec
 
-`pivot.ssh_exec(target: str, port: int, command: str, username: str, password: Optional<str>, key: Optional<str>, key_password: Optional<str>, timeout: Optional<int>) -> List<Dict>`
+`pivot.ssh_exec(target: str, port: int, command: str, username: str, password: Option<str>, key: Option<str>, key_password: Option<str>, timeout: Option<int>) -> List<Dict>`
 
 The <b>pivot.ssh_exec</b> method executes a command string on the remote host using the default shell.
 Stdout returns the string result from the command output.
@@ -657,7 +659,7 @@ Status will be equal to the code returned by the command being run and -1 in the
 
 ### process.info
 
-`process.info(pid: Optional<int>) -> Dict`
+`process.info(pid: Option<int>) -> Dict`
 
 The <b>process.info</b> method returns all information on a given process ID. Default is the current process.
 

implants/lib/eldritch/src/pivot/mod.rs

@@ -78,8 +78,8 @@ fn methods(builder: &mut MethodsBuilder) {
     }
 
     #[allow(unused_variables)]
-    fn ncat(this: &PivotLibrary, address: String, port: i32, data: String, protocol: String) ->  anyhow::Result<String> {
-        ncat_impl::ncat(address, port, data, protocol)
+    fn ncat(this: &PivotLibrary, address: String, port: i32, data: String, protocol: String, timeout: Option<u32>) ->  anyhow::Result<String> {
+        ncat_impl::ncat(address, port, data, protocol, timeout)
     }
 
     #[allow(unused_variables)]

implants/lib/eldritch/src/pivot/ncat_impl.rs

@@ -1,27 +1,34 @@
-use std::net::Ipv4Addr;
+use std::{
+    io::{BufReader, Read, Write},
+    net::{Ipv4Addr, SocketAddr, TcpStream, UdpSocket},
+    time::Duration,
+};
 
 use anyhow::Result;
-use tokio::io::{AsyncReadExt, AsyncWriteExt, BufReader};
-use tokio::net::{TcpStream, UdpSocket};
 
-// Since we cannot go from async (test) -> sync (ncat) `block_on` -> async (handle_ncat) without getting an error "cannot create runtime in current runtime since current thread is calling async code."
-async fn handle_ncat(address: String, port: i32, data: String, protocol: String) -> Result<String> {
+async fn handle_ncat_timeout(
+    address: String,
+    port: i32,
+    data: String,
+    protocol: String,
+    duration: Duration,
+) -> Result<String> {
     // If the response is longer than 4096 bytes it will be  truncated.
     let mut response_buffer: Vec<u8> = Vec::new();
     let result_string: String;
 
-    let address_and_port = format!("{}:{}", address, port);
+    let address_and_port = format!("{}:{}", address, port).parse::<SocketAddr>()?;
 
     if protocol == "tcp" {
         // Connect to remote host
-        let mut connection = TcpStream::connect(&address_and_port).await?;
+        let mut connection = TcpStream::connect_timeout(&address_and_port, duration)?;
 
         // Write our meessage
-        connection.write_all(data.as_bytes()).await?;
+        connection.write_all(data.as_bytes())?;
 
         // Read server response
         let mut read_stream = BufReader::new(connection);
-        read_stream.read_buf(&mut response_buffer).await?;
+        read_stream.read_to_end(&mut response_buffer)?;
 
         // We  need to take a buffer of bytes, turn it into a String but that string has null bytes.
         // To remove the null bytes we're using trim_matches.
@@ -36,16 +43,16 @@ async fn handle_ncat(address: String, port: i32, data: String, protocol: String)
 
         // Setting the bind address to unspecified should leave it up to the OS to decide.
         // https://stackoverflow.com/a/67084977
-        let sock = UdpSocket::bind((Ipv4Addr::UNSPECIFIED, 0)).await?;
+        let sock = UdpSocket::bind((Ipv4Addr::UNSPECIFIED, 0))?;
+        sock.set_read_timeout(Some(duration))?;
+        sock.set_write_timeout(Some(duration))?;
 
         // Send bytes to remote host
-        let _bytes_sent = sock
-            .send_to(data.as_bytes(), address_and_port.clone())
-            .await?;
+        let _bytes_sent = sock.send_to(data.as_bytes(), address_and_port)?;
 
         // Recieve any response from remote host
         let mut response_buffer = [0; 1024];
-        let (_bytes_copied, _addr) = sock.recv_from(&mut response_buffer).await?;
+        let (_bytes_copied, _addr) = sock.recv_from(&mut response_buffer)?;
 
         // We  need to take a buffer of bytes, turn it into a String but that string has null bytes.
         // To remove the null bytes we're using trim_matches.
@@ -59,14 +66,53 @@ async fn handle_ncat(address: String, port: i32, data: String, protocol: String)
     }
 }
 
+// Since we cannot go from async (test) -> sync (ncat) `block_on` -> async (handle_ncat) without getting an error "cannot create runtime in current runtime since current thread is calling async code."
+async fn handle_ncat(
+    address: String,
+    port: i32,
+    data: String,
+    protocol: String,
+    timeout: u32,
+) -> Result<String> {
+    let duration = std::time::Duration::from_secs(timeout as u64);
+    let res = match tokio::time::timeout(
+        duration,
+        handle_ncat_timeout(address, port, data, protocol, duration),
+    )
+    .await?
+    {
+        Ok(local_res) => local_res,
+        Err(local_err) => {
+            return Err(anyhow::anyhow!(
+                "Failed to run handle_ncat_timeout: {}",
+                local_err.to_string()
+            ))
+        }
+    };
+
+    Ok(res)
+}
+
 // We do not want to make this async since it would require we make all of the starlark bindings async.
 // Instead we have a handle_ncat function that we call with block_on
-pub fn ncat(address: String, port: i32, data: String, protocol: String) -> Result<String> {
+pub fn ncat(
+    address: String,
+    port: i32,
+    data: String,
+    protocol: String,
+    timeout: Option<u32>,
+) -> Result<String> {
+    let default_timeout = 2;
+    let timeout_u32 = match timeout {
+        Some(res) => res,
+        None => default_timeout,
+    };
+
     let runtime = tokio::runtime::Builder::new_current_thread()
         .enable_all()
         .build()?;
 
-    let response = runtime.block_on(handle_ncat(address, port, data, protocol));
+    let response = runtime.block_on(handle_ncat(address, port, data, protocol, timeout_u32));
 
     match response {
         Ok(_) => Ok(response.unwrap()),
@@ -76,6 +122,7 @@ pub fn ncat(address: String, port: i32, data: String, protocol: String) -> Resul
 
 #[cfg(test)]
 mod tests {
+
     use super::*;
     use anyhow::Context;
     use tokio::io::copy;
@@ -166,6 +213,7 @@ mod tests {
             test_port,
             expected_response.clone(),
             String::from("tcp"),
+            5,
         ));
 
         // Will this create a race condition where the sender sends before the listener starts?
@@ -194,6 +242,7 @@ mod tests {
             test_port,
             expected_response.clone(),
             String::from("udp"),
+            5,
         ));
 
         // Will this create a race condition where the sender sends before the listener starts?
@@ -204,6 +253,33 @@ mod tests {
         assert_eq!(expected_response, actual_response.unwrap().unwrap());
         Ok(())
     }
+    #[tokio::test]
+    async fn test_ncat_timeout_exceeded() -> anyhow::Result<()> {
+        let test_port = allocate_localhost_unused_ports(1, "udp".to_string()).await?[0];
+
+        // Setup a test echo server
+        let expected_response = String::from("Hello world!");
+
+        // Setup a sender
+        let send_task = task::spawn(handle_ncat(
+            String::from("127.0.0.1"),
+            test_port,
+            expected_response.clone(),
+            String::from("udp"),
+            2,
+        ))
+        .await?;
+
+        assert!(send_task.is_err());
+        let err_string = send_task.unwrap_err().to_string();
+        println!("{}", err_string);
+        assert!(
+            err_string.contains("deadline has elapsed")
+                || err_string.contains("not properly respond after a period of time")
+        );
+
+        Ok(())
+    }
     // #[test]
     // fn test_ncat_not_handle() -> anyhow::Result<()> {
     //     let runtime = tokio::runtime::Builder::new_current_thread()