Add bindings to agent environment

.gitignore

@@ -64,6 +64,7 @@ docs/_build/
 
 # PyCharm
 .idea/
+*.iml
 
 # VSCode
 .vscode/

README.md

@@ -2,3 +2,15 @@ BinHarness
 ===
 
 BinHarness is a framework to facilitate analyzing binary programs in various environments.
+
+## Project layout
+```
+binharness
+| crates: Main directory for rust code
+  | bh_agent_client: Client code for communicating with the binharness agent
+  | bh_agent_common: Shared code between the agent client and server
+  | bh_agent_server: The agent server program
+| python: Main directory for python code
+  | binharness: Main python module
+  | tests: Test code
+```

crates/bh_agent_client/src/bindings.rs

@@ -177,7 +177,7 @@ impl BhAgentClient {
         )
     }
 
-    fn file_read(&self, env_id: EnvironmentId, fd: FileId, num_bytes: u32) -> PyResult<Vec<u8>> {
+    fn file_read(&self, env_id: EnvironmentId, fd: FileId, num_bytes: Option<u32>) -> PyResult<Vec<u8>> {
         run_in_runtime(
             self,
             self.client

crates/bh_agent_common/src/service.rs

@@ -45,7 +45,7 @@ pub trait BhAgentService {
     async fn file_read(
         env_id: EnvironmentId,
         fd: FileId,
-        num_bytes: u32,
+        num_bytes: Option<u32>,
     ) -> Result<Vec<u8>, AgentError>;
 
     async fn file_read_lines(

crates/bh_agent_server/Cargo.toml

@@ -13,3 +13,6 @@ tarpc = { version = "0.33.0", features = ["full"] }
 tokio = { version = "1.32.0", features = ["full"] }
 futures-util = "0.3.28"
 futures = "0.3.28"
+log = "0.4.20"
+env_logger = "0.10.0"
+bimap = "0.6.3"

crates/bh_agent_server/src/main.rs

@@ -25,6 +25,8 @@ fn parse_args() -> Result<(IpAddr, u16)> {
 
 #[tokio::main]
 async fn main() -> anyhow::Result<()> {
+    env_logger::init();
+
     let server_addr = parse_args().or_else(|e| -> Result<(IpAddr, u16)> {
         eprintln!("{}", e);
         std::process::exit(1);

crates/bh_agent_server/src/server.rs

@@ -132,7 +132,7 @@ impl BhAgentService for BhAgentServer {
         _: Context,
         env_id: EnvironmentId,
         fd: FileId,
-        num_bytes: u32,
+        num_bytes: Option<u32>,
     ) -> Self::FileReadFut {
         check_env_id!(env_id);
 

crates/bh_agent_server/src/state.rs

@@ -1,8 +1,10 @@
+use bimap::BiMap;
 use std::collections::HashMap;
 use std::ffi::OsStr;
+use std::fmt::Display;
 use std::fs::{File, OpenOptions};
 use std::sync::{Arc, RwLock};
-use futures_util::TryFutureExt;
+use log::{debug, trace};
 
 use subprocess::{Popen, PopenConfig};
 
@@ -20,9 +22,9 @@ pub struct BhAgentState {
     file_modes: RwLock<HashMap<FileId, FileOpenMode>>,
     file_types: RwLock<HashMap<FileId, FileOpenType>>,
     processes: RwLock<HashMap<ProcessId, Arc<RwLock<Popen>>>>,
-    proc_stdin_ids: RwLock<HashMap<FileId, ProcessId>>,
-    proc_stdout_ids: RwLock<HashMap<FileId, ProcessId>>,
-    proc_stderr_ids: RwLock<HashMap<FileId, ProcessId>>,
+    proc_stdin_ids: RwLock<BiMap<ProcessId, FileId>>,
+    proc_stdout_ids: RwLock<BiMap<ProcessId, FileId>>,
+    proc_stderr_ids: RwLock<BiMap<ProcessId, FileId>>,
 
     next_file_id: RwLock<FileId>,
     next_process_id: RwLock<ProcessId>,
@@ -35,9 +37,9 @@ impl BhAgentState {
             file_modes: RwLock::new(HashMap::new()),
             file_types: RwLock::new(HashMap::new()),
             processes: RwLock::new(HashMap::new()),
-            proc_stdin_ids: RwLock::new(HashMap::new()),
-            proc_stdout_ids: RwLock::new(HashMap::new()),
-            proc_stderr_ids: RwLock::new(HashMap::new()),
+            proc_stdin_ids: RwLock::new(BiMap::new()),
+            proc_stdout_ids: RwLock::new(BiMap::new()),
+            proc_stderr_ids: RwLock::new(BiMap::new()),
 
             next_file_id: RwLock::new(0),
             next_process_id: RwLock::new(0),
@@ -63,6 +65,7 @@ impl BhAgentState {
         fd: &FileId,
         modes: &Vec<FileOpenMode>,
     ) -> Result<bool, AgentError> {
+        trace!("Checking file {} for modes {:?}", fd, modes);
         Ok(modes.contains(
             self.file_modes
                 .read()?
@@ -72,6 +75,7 @@ impl BhAgentState {
     }
 
     pub fn file_type(&self, fd: &FileId) -> Result<FileOpenType, AgentError> {
+        trace!("Getting file type for {}", fd);
         Ok(self
             .file_types
             .read()?
@@ -154,16 +158,31 @@ impl BhAgentState {
 
         // Stick the process channels into the file map
         if proc.stdin.is_some() {
+            trace!("Saving stdin for process {}", proc_id);
             let file_id = self.take_file_id()?;
-            self.proc_stdin_ids.write()?.insert(file_id, proc_id);
+            self.proc_stdin_ids.write()?.insert(proc_id, file_id);
+            self.file_modes.write()?.insert(file_id, FileOpenMode::Write);
+            self.file_types.write()?.insert(file_id, FileOpenType::Binary);
+        } else {
+            trace!("Process {} has no stdin", proc_id);
         }
         if proc.stdout.is_some() {
+            trace!("Saving stdout for process {}", proc_id);
             let file_id = self.take_file_id()?;
-            self.proc_stdout_ids.write()?.insert(file_id, proc_id);
+            self.proc_stdout_ids.write()?.insert(proc_id, file_id);
+            self.file_modes.write()?.insert(file_id, FileOpenMode::Read);
+            self.file_types.write()?.insert(file_id, FileOpenType::Binary);
+        } else {
+            trace!("Process {} has no stdout", proc_id);
         }
         if proc.stderr.is_some() {
+            trace!("Saving stderr for process {}", proc_id);
             let file_id = self.take_file_id()?;
-            self.proc_stdout_ids.write()?.insert(file_id, proc_id);
+            self.proc_stdout_ids.write()?.insert(proc_id, file_id);
+            self.file_modes.write()?.insert(file_id, FileOpenMode::Read);
+            self.file_types.write()?.insert(file_id, FileOpenType::Binary);
+        } else {
+            trace!("Process {} has no stderr", proc_id);
         }
 
         // Move the proc to the process map
@@ -179,18 +198,30 @@ impl BhAgentState {
         proc_id: &ProcessId,
         channel: ProcessChannel,
     ) -> Result<FileId, AgentError> {
-        match channel {
+        let channel_ids = match channel {
             ProcessChannel::Stdin => &self.proc_stdin_ids,
             ProcessChannel::Stdout => &self.proc_stdout_ids,
             ProcessChannel::Stderr => &self.proc_stderr_ids,
-        }
+        };
+
+        channel_ids
         .read()?
-        .get(&proc_id)
+        .get_by_left(&proc_id)
         .map(|i| i.clone())
-        .ok_or(InvalidProcessId)
+        .ok_or((|| {
+            debug!("Failed to get process channel");
+            debug!("Process ID: {}", proc_id);
+            debug!("Channel: {:?}", channel);
+            let proc_is_valid = self.processes.read().unwrap().contains_key(&proc_id);
+            debug!("Process is valid: {}", proc_is_valid);
+            debug!("Process with valid channels: {:?}", channel_ids.read().unwrap().left_values());
+
+            InvalidProcessId
+        })())
     }
 
     pub fn close_file(&self, fd: &FileId) -> Result<(), AgentError> {
+        trace!("Closing file {}", fd);
         Ok(drop(
             self.files
                 .write()?
@@ -208,29 +239,39 @@ impl BhAgentState {
         fd: &FileId,
         op: impl Fn(&mut File) -> R,
     ) -> Result<R, AgentError> {
+        trace!("Doing mut operation on file {}", fd);
+
         // Get file logic
         if let Some(file_lock) = self.files.read()?.get(fd) {
             return Ok(op(&mut *file_lock.write()?));
+        } else {
+            trace!("File id map: {:?}", self.files.read()?);
         }
 
         // If these unwraps fail, the state is bad
-        if let Some(pid) = self.proc_stdin_ids.read()?.get(&fd) {
+        if let Some(pid) = self.proc_stdin_ids.read()?.get_by_right(&fd) {
             let procs_binding = self.processes.read()?;
             let mut proc_binding = procs_binding.get(pid).unwrap().write()?;
             let file = proc_binding.stdin.as_mut().unwrap();
             return Ok(op(file));
+        } else {
+            trace!("Process stdin id map: {:?}", self.proc_stdin_ids.read()?);
         }
-        if let Some(pid) = self.proc_stdout_ids.read()?.get(&fd) {
+        if let Some(pid) = self.proc_stdout_ids.read()?.get_by_right(&fd) {
             let procs_binding = self.processes.read()?;
             let mut proc_binding = procs_binding.get(pid).unwrap().write()?;
             let file = proc_binding.stdout.as_mut().unwrap();
             return Ok(op(file));
+        } else {
+            trace!("Process stdout id map: {:?}", self.proc_stdout_ids.read()?);
         }
-        if let Some(pid) = self.proc_stderr_ids.read()?.get(&fd) {
+        if let Some(pid) = self.proc_stderr_ids.read()?.get_by_right(&fd) {
             let procs_binding = self.processes.read()?;
             let mut proc_binding = procs_binding.get(pid).unwrap().write()?;
             let file = proc_binding.stderr.as_mut().unwrap();
             return Ok(op(file));
+        } else {
+            trace!("Process stderr id map: {:?}", self.proc_stderr_ids.read()?);
         }
 
         Err(InvalidFileDescriptor)

crates/bh_agent_server/src/util/read_chars.rs

@@ -2,17 +2,26 @@ use std::fs::File;
 use std::io::Read;
 
 use anyhow::Result;
+use log::trace;
 
 use bh_agent_common::FileOpenType;
 
-pub fn read_generic(mut file: &File, n: u32, file_type: FileOpenType) -> Result<Vec<u8>> {
-    match file_type {
-        FileOpenType::Binary => {
-            let mut buffer = vec![0u8; n as usize];
-            file.read(&mut buffer)?;
-            Ok(buffer)
+pub fn read_generic(mut file: &File, n: Option<u32>, file_type: FileOpenType) -> Result<Vec<u8>> {
+    trace!("Entering read_generic");
+    if let Some(num_bytes) = n {
+        match file_type {
+            FileOpenType::Binary => {
+                let mut buffer = vec![0u8; num_bytes as usize];
+                file.read(&mut buffer)?;
+                Ok(buffer)
+            }
+            FileOpenType::Text => Ok(read_chars(&mut file, num_bytes as usize)?),
         }
-        FileOpenType::Text => Ok(read_chars(&mut file, n as usize)?),
+    } else {
+        // if n is None, we just read the whole file, text parsing happens on the client
+        let mut buffer = Vec::new();
+        file.read_to_end(&mut buffer)?;
+        Ok(buffer)
     }
 }
 

pyproject.toml

@@ -55,10 +55,15 @@ ignore = [
     "COM",
     "D203",
     "D213",
+    "FBT003",
+    "FIX",
     "S108",
     "S603",
+    "TD002",
+    "TD003",
 ]
 ignore-init-module-imports = true
+line-length = 88
 
 [tool.ruff.isort]
 required-imports = ["from __future__ import annotations"]
@@ -68,6 +73,11 @@ required-imports = ["from __future__ import annotations"]
     "D",
     "S101",
 ]
+"**/*.pyi" = [
+    "ANN101",
+    "FBT",
+    "PLR0913",
+]
 
 [tool.mypy]
 python_version = "3.8"

python/bh_agent_client.pyi

@@ -0,0 +1,31 @@
+class BhAgentClient:
+    @staticmethod
+    def initialize_client(ip_addr: str, port: int) -> BhAgentClient: ...
+    def get_environments(self) -> list[int]: ...
+    def get_tempdir(self, env_id: int) -> str: ...
+    def run_process(
+        self,
+        env_id: int,
+        argv: list[str],
+        stdin: bool,
+        stdout: bool,
+        stderr: bool,
+        executable: str | None,
+        env: dict[str, str] | None,
+        cwd: str | None,
+        setuid: int | None,
+        setgid: int | None,
+        setpgid: int | None,
+    ) -> int: ...
+    def get_process_channel(self, env_id: int, proc_id: int, channel: int) -> int: ...
+    def file_open(self, env_id: int, path: str, mode_and_type: str) -> int: ...
+    def file_close(self, env_id: int, fd: int) -> None: ...
+    def file_is_closed(self, env_id: int, fd: int) -> bool: ...
+    def file_is_readable(self, env_id: int, fd: int) -> bool: ...
+    def file_read(self, env_id: int, fd: int, size: int | None) -> bytes: ...
+    def file_read_lines(self, env_id: int, fd: int, hint: int) -> list[bytes]: ...
+    def file_is_seekable(self, env_id: int, fd: int) -> bool: ...
+    def file_seek(self, env_id: int, fd: int, offset: int, whence: int) -> int: ...
+    def file_tell(self, env_id: int, fd: int) -> int: ...
+    def file_is_writable(self, env_id: int, fd: int) -> bool: ...
+    def file_write(self, env_id: int, fd: int, data: bytes) -> int: ...