From 69906aba7825c311569a08edbdbc8e8f8515a25c Mon Sep 17 00:00:00 2001 From: Kevin Phoenix Date: Wed, 18 Oct 2023 17:09:15 -0700 Subject: [PATCH] Merge agent into main repo --- .gitignore | 2 + Cargo.toml | 3 + crates/bh_agent_client/.gitignore | 72 +++++ crates/bh_agent_client/Cargo.toml | 16 ++ crates/bh_agent_client/src/bindings.rs | 245 +++++++++++++++++ crates/bh_agent_client/src/client.rs | 15 ++ crates/bh_agent_client/src/lib.rs | 4 + crates/bh_agent_common/Cargo.toml | 12 + crates/bh_agent_common/src/agent_error.rs | 46 ++++ crates/bh_agent_common/src/lib.rs | 7 + crates/bh_agent_common/src/service.rs | 72 +++++ crates/bh_agent_common/src/types.rs | 48 ++++ crates/bh_agent_server/Cargo.toml | 15 ++ crates/bh_agent_server/src/lib.rs | 5 + crates/bh_agent_server/src/main.rs | 51 ++++ crates/bh_agent_server/src/server.rs | 246 ++++++++++++++++++ crates/bh_agent_server/src/state.rs | 238 +++++++++++++++++ crates/bh_agent_server/src/util/mod.rs | 5 + crates/bh_agent_server/src/util/read_chars.rs | 87 +++++++ crates/bh_agent_server/src/util/read_lines.rs | 39 +++ pyproject.toml | 18 +- {binharness => python/binharness}/__init__.py | 0 .../binharness}/common/__init__.py | 0 .../binharness}/common/busybox.py | 0 .../binharness}/common/qemu.py | 0 .../binharness}/localenvironment.py | 0 .../binharness}/serialize.py | 0 .../binharness}/types/__init__.py | 0 .../binharness}/types/environment.py | 0 .../binharness}/types/executor.py | 0 .../binharness}/types/injection.py | 0 {binharness => python/binharness}/types/io.py | 0 .../binharness}/types/process.py | 0 .../binharness}/types/target.py | 0 {binharness => python/binharness}/util.py | 0 {binharness => python}/tests/__init__.py | 0 {binharness => python}/tests/ssh_keys/test | 0 .../tests/ssh_keys/test.pub | 0 {binharness => python}/tests/test_busybox.py | 0 {binharness => python}/tests/test_executor.py | 1 - {binharness => python}/tests/test_inject.py | 1 - .../tests/test_localenvironment.py | 0 .../tests/test_qemu_executor.py | 0 .../tests/test_target_local.py | 0 .../tests/test_target_serialization.py | 1 - {binharness => python}/tests/test_util.py | 0 46 files changed, 1238 insertions(+), 11 deletions(-) create mode 100644 Cargo.toml create mode 100644 crates/bh_agent_client/.gitignore create mode 100644 crates/bh_agent_client/Cargo.toml create mode 100644 crates/bh_agent_client/src/bindings.rs create mode 100644 crates/bh_agent_client/src/client.rs create mode 100644 crates/bh_agent_client/src/lib.rs create mode 100644 crates/bh_agent_common/Cargo.toml create mode 100644 crates/bh_agent_common/src/agent_error.rs create mode 100644 crates/bh_agent_common/src/lib.rs create mode 100644 crates/bh_agent_common/src/service.rs create mode 100644 crates/bh_agent_common/src/types.rs create mode 100644 crates/bh_agent_server/Cargo.toml create mode 100644 crates/bh_agent_server/src/lib.rs create mode 100644 crates/bh_agent_server/src/main.rs create mode 100644 crates/bh_agent_server/src/server.rs create mode 100644 crates/bh_agent_server/src/state.rs create mode 100644 crates/bh_agent_server/src/util/mod.rs create mode 100644 crates/bh_agent_server/src/util/read_chars.rs create mode 100644 crates/bh_agent_server/src/util/read_lines.rs rename {binharness => python/binharness}/__init__.py (100%) rename {binharness => python/binharness}/common/__init__.py (100%) rename {binharness => python/binharness}/common/busybox.py (100%) rename {binharness => python/binharness}/common/qemu.py (100%) rename {binharness => python/binharness}/localenvironment.py (100%) rename {binharness => python/binharness}/serialize.py (100%) rename {binharness => python/binharness}/types/__init__.py (100%) rename {binharness => python/binharness}/types/environment.py (100%) rename {binharness => python/binharness}/types/executor.py (100%) rename {binharness => python/binharness}/types/injection.py (100%) rename {binharness => python/binharness}/types/io.py (100%) rename {binharness => python/binharness}/types/process.py (100%) rename {binharness => python/binharness}/types/target.py (100%) rename {binharness => python/binharness}/util.py (100%) rename {binharness => python}/tests/__init__.py (100%) rename {binharness => python}/tests/ssh_keys/test (100%) rename {binharness => python}/tests/ssh_keys/test.pub (100%) rename {binharness => python}/tests/test_busybox.py (100%) rename {binharness => python}/tests/test_executor.py (99%) rename {binharness => python}/tests/test_inject.py (99%) rename {binharness => python}/tests/test_localenvironment.py (100%) rename {binharness => python}/tests/test_qemu_executor.py (100%) rename {binharness => python}/tests/test_target_local.py (100%) rename {binharness => python}/tests/test_target_serialization.py (99%) rename {binharness => python}/tests/test_util.py (100%) diff --git a/.gitignore b/.gitignore index f18df2c..c4e2be4 100644 --- a/.gitignore +++ b/.gitignore @@ -10,3 +10,5 @@ build *.egg-info *.pyc .idea +target +Cargo.lock \ No newline at end of file diff --git a/Cargo.toml b/Cargo.toml new file mode 100644 index 0000000..d767148 --- /dev/null +++ b/Cargo.toml @@ -0,0 +1,3 @@ +[workspace] +members = ["crates/*"] +resolver = "2" diff --git a/crates/bh_agent_client/.gitignore b/crates/bh_agent_client/.gitignore new file mode 100644 index 0000000..af3ca5e --- /dev/null +++ b/crates/bh_agent_client/.gitignore @@ -0,0 +1,72 @@ +/target + +# Byte-compiled / optimized / DLL files +__pycache__/ +.pytest_cache/ +*.py[cod] + +# C extensions +*.so + +# Distribution / packaging +.Python +.venv/ +env/ +bin/ +build/ +develop-eggs/ +dist/ +eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +include/ +man/ +venv/ +*.egg-info/ +.installed.cfg +*.egg + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt +pip-selfcheck.json + +# Unit test / coverage reports +htmlcov/ +.tox/ +.coverage +.cache +nosetests.xml +coverage.xml + +# Translations +*.mo + +# Mr Developer +.mr.developer.cfg +.project +.pydevproject + +# Rope +.ropeproject + +# Django stuff: +*.log +*.pot + +.DS_Store + +# Sphinx documentation +docs/_build/ + +# PyCharm +.idea/ + +# VSCode +.vscode/ + +# Pyenv +.python-version \ No newline at end of file diff --git a/crates/bh_agent_client/Cargo.toml b/crates/bh_agent_client/Cargo.toml new file mode 100644 index 0000000..a3018a9 --- /dev/null +++ b/crates/bh_agent_client/Cargo.toml @@ -0,0 +1,16 @@ +[package] +name = "bh_agent_client" +version = "0.1.0" +edition = "2021" + +# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html +[lib] +name = "bh_agent_client" +crate-type = ["cdylib"] + +[dependencies] +pyo3 = "0.19.0" +bh_agent_common = { path = "../bh_agent_common" } +tokio = "1.32.0" +anyhow = "1.0.75" +tarpc = { version = "0.33.0", features = ["full"] } diff --git a/crates/bh_agent_client/src/bindings.rs b/crates/bh_agent_client/src/bindings.rs new file mode 100644 index 0000000..07f418e --- /dev/null +++ b/crates/bh_agent_client/src/bindings.rs @@ -0,0 +1,245 @@ +use crate::client::build_client; +use anyhow::Result; +use bh_agent_common::{ + AgentError, BhAgentServiceClient, EnvironmentId, FileId, FileOpenMode, FileOpenType, + ProcessChannel, ProcessId, Redirection, RemotePOpenConfig, +}; +use pyo3::exceptions::PyRuntimeError; +use pyo3::prelude::*; +use pyo3::{pyclass, pymethods, pymodule, PyResult, Python}; +use std::future::Future; +use std::net::{IpAddr, SocketAddr}; +use std::str::FromStr; +use tarpc::client::RpcError; +use tarpc::context; +use tokio::runtime; + +#[pyclass] +struct BhAgentClient { + tokio_runtime: runtime::Runtime, + client: BhAgentServiceClient, +} + +fn run_in_runtime(client: &BhAgentClient, fut: F) -> PyResult +where + F: Future, RpcError>> + Sized, +{ + client + .tokio_runtime + .block_on(fut) + .map_err(|e| PyRuntimeError::new_err(e.to_string())) + .map(|r| r.map_err(|e| PyRuntimeError::new_err(e.to_string()))) + .and_then(|r| r) +} + +#[pymethods] +impl BhAgentClient { + #[staticmethod] + fn initialize_client(ip_addr: String, port: u16) -> PyResult { + let ip_addr = IpAddr::from_str(&ip_addr)?; + let socket_addr = SocketAddr::new(ip_addr, port); + + let tokio_runtime = runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + match tokio_runtime.block_on(build_client(socket_addr)) { + Ok(client) => Ok(Self { + tokio_runtime, + client, + }), + Err(e) => Err(PyRuntimeError::new_err(format!( + "Failed to initialize client: {}", + e + ))), + } + } + + fn get_environments(&self) -> PyResult> { + self.tokio_runtime + .block_on(self.client.get_environments(context::current())) + .map_err(|e| PyRuntimeError::new_err(e.to_string())) + } + + fn get_tempdir(&self, env_id: EnvironmentId) -> PyResult { + run_in_runtime(self, self.client.get_tempdir(context::current(), env_id)) + } + + fn run_process( + &self, + env_id: EnvironmentId, + argv: Vec, + stdin: bool, + stdout: bool, + stderr: bool, + executable: Option, + env: Option>, + cwd: Option, + setuid: Option, + setgid: Option, + setpgid: bool, + ) -> PyResult { + let config = RemotePOpenConfig { + argv, + stdin: match stdin { + true => Redirection::Save, + false => Redirection::None, + }, + stdout: match stdout { + true => Redirection::Save, + false => Redirection::None, + }, + stderr: match stderr { + true => Redirection::Save, + false => Redirection::None, + }, + executable, + env, + cwd, + setuid, + setgid, + setpgid, + }; + run_in_runtime( + self, + self.client.run_command(context::current(), env_id, config), + ) + } + + fn get_process_channel( + &self, + env_id: EnvironmentId, + proc_id: ProcessId, + channel: i32, // TODO: This is just 0, 1, 2 for now + ) -> PyResult { + run_in_runtime( + self, + self.client.get_process_channel( + context::current(), + env_id, + proc_id, + match channel { + 0 => ProcessChannel::Stdin, + 1 => ProcessChannel::Stdout, + 2 => ProcessChannel::Stderr, + _ => return Err(PyRuntimeError::new_err("Invalid channel")), + }, + ), + ) + } + + // File IO + fn file_open( + &self, + env_id: EnvironmentId, + path: String, + mode_and_type: String, + ) -> PyResult { + // Mode parsing + let mut mode = FileOpenMode::Read; + mode_and_type.chars().for_each(|c| match c { + 'r' => mode = FileOpenMode::Read, + 'w' => mode = FileOpenMode::Write, + 'x' => mode = FileOpenMode::ExclusiveWrite, + 'a' => mode = FileOpenMode::Append, + '+' => mode = FileOpenMode::Update, + _ => {} + }); + + // Type parsing + let mut type_ = FileOpenType::Text; + if mode_and_type.contains("b") { + type_ = FileOpenType::Binary; + } + + run_in_runtime( + self, + self.client + .file_open(context::current(), env_id, path, mode, type_), + ) + } + + fn file_close(&self, env_id: EnvironmentId, fd: FileId) -> PyResult<()> { + run_in_runtime(self, self.client.file_close(context::current(), env_id, fd)) + } + + fn file_is_closed(&self, env_id: EnvironmentId, fd: FileId) -> PyResult { + run_in_runtime( + self, + self.client.file_is_closed(context::current(), env_id, fd), + ) + } + + fn file_is_readable(&self, env_id: EnvironmentId, fd: FileId) -> PyResult { + run_in_runtime( + self, + self.client.file_is_readable(context::current(), env_id, fd), + ) + } + + fn file_read(&self, env_id: EnvironmentId, fd: FileId, num_bytes: u32) -> PyResult> { + run_in_runtime( + self, + self.client + .file_read(context::current(), env_id, fd, num_bytes), + ) + } + + fn file_read_lines( + &self, + env_id: EnvironmentId, + fd: FileId, + hint: u32, + ) -> PyResult>> { + run_in_runtime( + self, + self.client + .file_read_lines(context::current(), env_id, fd, hint), + ) + } + + fn file_is_seekable(&self, env_id: EnvironmentId, fd: FileId) -> PyResult { + run_in_runtime( + self, + self.client.file_is_seekable(context::current(), env_id, fd), + ) + } + + fn file_seek( + &self, + env_id: EnvironmentId, + fd: FileId, + offset: i32, + whence: i32, + ) -> PyResult<()> { + run_in_runtime( + self, + self.client + .file_seek(context::current(), env_id, fd, offset, whence), + ) + } + + fn file_tell(&self, env_id: EnvironmentId, fd: FileId) -> PyResult { + run_in_runtime(self, self.client.file_tell(context::current(), env_id, fd)) + } + + fn file_is_writable(&self, env_id: EnvironmentId, fd: FileId) -> PyResult { + run_in_runtime( + self, + self.client.file_is_writable(context::current(), env_id, fd), + ) + } + + fn file_write(&self, env_id: EnvironmentId, fd: FileId, data: Vec) -> PyResult<()> { + run_in_runtime( + self, + self.client.file_write(context::current(), env_id, fd, data), + ) + } +} + +#[pymodule] +pub fn bh_agent_client(_py: Python, m: &PyModule) -> PyResult<()> { + m.add_class::()?; + Ok(()) +} diff --git a/crates/bh_agent_client/src/client.rs b/crates/bh_agent_client/src/client.rs new file mode 100644 index 0000000..8b32579 --- /dev/null +++ b/crates/bh_agent_client/src/client.rs @@ -0,0 +1,15 @@ +use bh_agent_common::BhAgentServiceClient; +use tarpc::{client, tokio_serde::formats::Json}; +use tokio::net::ToSocketAddrs; + +pub async fn build_client(socket_addr: A) -> anyhow::Result +where + A: ToSocketAddrs, +{ + let mut transport = tarpc::serde_transport::tcp::connect(socket_addr, Json::default); + transport.config_mut().max_frame_length(usize::MAX); + + let client = BhAgentServiceClient::new(client::Config::default(), transport.await?).spawn(); + + Ok(client) +} diff --git a/crates/bh_agent_client/src/lib.rs b/crates/bh_agent_client/src/lib.rs new file mode 100644 index 0000000..6b48809 --- /dev/null +++ b/crates/bh_agent_client/src/lib.rs @@ -0,0 +1,4 @@ +mod bindings; +mod client; + +pub use bindings::bh_agent_client; diff --git a/crates/bh_agent_common/Cargo.toml b/crates/bh_agent_common/Cargo.toml new file mode 100644 index 0000000..2f84189 --- /dev/null +++ b/crates/bh_agent_common/Cargo.toml @@ -0,0 +1,12 @@ +[package] +name = "bh_agent_common" +version = "0.1.0" +edition = "2021" + +# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html + +[dependencies] +anyhow = { version = "1.0.75", features = [] } +tarpc = { version = "0.33.0", features = ["tokio1"] } +serde = { version = "1.0.188", features = ["derive"] } +thiserror = "1.0.48" diff --git a/crates/bh_agent_common/src/agent_error.rs b/crates/bh_agent_common/src/agent_error.rs new file mode 100644 index 0000000..1d38bc1 --- /dev/null +++ b/crates/bh_agent_common/src/agent_error.rs @@ -0,0 +1,46 @@ +use crate::AgentError::LockError; +use serde::{Deserialize, Serialize}; +use std::sync::{PoisonError, RwLockReadGuard, RwLockWriteGuard}; +use thiserror::Error; + +#[derive(Error, Debug, Serialize, Deserialize)] +pub enum AgentError { + #[error("Invalid environment ID")] + InvalidEnvironmentId, + #[error("IO Error")] + IoError, + #[error("Invalid file ID")] + InvalidFileDescriptor, + #[error("Invalid seek whence")] + InvalidSeekWhence, + #[error("Lock Error")] + LockError, + #[error("Failed to start process")] + ProcessStartFailure, + #[error("Invalid process ID")] + InvalidProcessId, + #[error("Process channel not piped")] + ProcessChannelNotPiped, + #[error("The server state is inconsistent")] + Inconsistent, + #[error("Unknown Error")] + Unknown, +} + +impl From> for AgentError { + fn from(_: PoisonError) -> Self { + LockError + } +} + +impl From> for AgentError { + fn from(_: RwLockReadGuard) -> Self { + LockError + } +} + +impl From> for AgentError { + fn from(_: RwLockWriteGuard) -> Self { + LockError + } +} diff --git a/crates/bh_agent_common/src/lib.rs b/crates/bh_agent_common/src/lib.rs new file mode 100644 index 0000000..954cd37 --- /dev/null +++ b/crates/bh_agent_common/src/lib.rs @@ -0,0 +1,7 @@ +mod agent_error; +mod service; +mod types; + +pub use agent_error::*; +pub use service::*; +pub use types::*; diff --git a/crates/bh_agent_common/src/service.rs b/crates/bh_agent_common/src/service.rs new file mode 100644 index 0000000..78c50aa --- /dev/null +++ b/crates/bh_agent_common/src/service.rs @@ -0,0 +1,72 @@ +use crate::agent_error::AgentError; +use crate::{ + EnvironmentId, FileId, FileOpenMode, FileOpenType, ProcessChannel, ProcessId, RemotePOpenConfig, +}; +use anyhow::Result; + +#[tarpc::service] +pub trait BhAgentService { + // Environment enumeration + async fn get_environments() -> Vec; + + async fn get_tempdir(env_id: EnvironmentId) -> Result; + + // Process management + async fn run_command( + env_id: EnvironmentId, + config: RemotePOpenConfig, + ) -> Result; + + async fn get_process_channel( + env_id: EnvironmentId, + proc_id: ProcessId, + channel: ProcessChannel, + ) -> Result; + + // File IO + // Implement most of the methods in binharness.IO, but omit ones that there can just be + // replicated on the client side without a performance hit. + // Data is represented as a Vec instead of a String because if we're in text mode, we can + // just decode it on the client side. The server still needs to know the mode however, as an + // N length read in text mode will be N chars, not N bytes. + async fn file_open( + env_id: EnvironmentId, + path: String, + mode: FileOpenMode, + type_: FileOpenType, + ) -> Result; + + async fn file_close(env_id: EnvironmentId, fd: FileId) -> Result<(), AgentError>; + + async fn file_is_closed(env_id: EnvironmentId, fd: FileId) -> Result; + + async fn file_is_readable(env_id: EnvironmentId, fd: FileId) -> Result; + + async fn file_read( + env_id: EnvironmentId, + fd: FileId, + num_bytes: u32, + ) -> Result, AgentError>; + + async fn file_read_lines( + env_id: EnvironmentId, + fd: FileId, + hint: u32, + ) -> Result>, AgentError>; + + async fn file_is_seekable(env_id: EnvironmentId, fd: FileId) -> Result; + + async fn file_seek( + env_id: EnvironmentId, + fd: FileId, + offset: i32, + whence: i32, + ) -> Result<(), AgentError>; + + async fn file_tell(env_id: EnvironmentId, fd: FileId) -> Result; + + async fn file_is_writable(env_id: EnvironmentId, fd: FileId) -> Result; + + async fn file_write(env_id: EnvironmentId, fd: FileId, data: Vec) + -> Result<(), AgentError>; +} diff --git a/crates/bh_agent_common/src/types.rs b/crates/bh_agent_common/src/types.rs new file mode 100644 index 0000000..93c9090 --- /dev/null +++ b/crates/bh_agent_common/src/types.rs @@ -0,0 +1,48 @@ +use serde::{Deserialize, Serialize}; + +pub type EnvironmentId = u64; +pub type ProcessId = u64; +pub type FileId = u64; + +#[derive(Copy, Clone, Debug, Serialize, Deserialize)] +pub enum ProcessChannel { + Stdin, + Stdout, + Stderr, +} + +#[derive(Copy, Clone, Debug, Default, Serialize, Deserialize)] +pub enum Redirection { + #[default] + None, + Save, +} + +#[derive(Clone, Debug, Default, Serialize, Deserialize)] +pub struct RemotePOpenConfig { + pub argv: Vec, + pub stdin: Redirection, + pub stdout: Redirection, + pub stderr: Redirection, + pub executable: Option, + pub env: Option>, + pub cwd: Option, + pub setuid: Option, + pub setgid: Option, + pub setpgid: bool, +} + +#[derive(Copy, Clone, Debug, Serialize, Deserialize, PartialEq)] +pub enum FileOpenMode { + Read, + Write, + ExclusiveWrite, + Append, + Update, +} + +#[derive(Copy, Clone, Debug, Serialize, Deserialize, PartialEq)] +pub enum FileOpenType { + Binary, + Text, +} diff --git a/crates/bh_agent_server/Cargo.toml b/crates/bh_agent_server/Cargo.toml new file mode 100644 index 0000000..ac68d93 --- /dev/null +++ b/crates/bh_agent_server/Cargo.toml @@ -0,0 +1,15 @@ +[package] +name = "bh_agent_server" +version = "0.1.0" +edition = "2021" + +# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html + +[dependencies] +anyhow = "1.0.75" +bh_agent_common = { path = "../bh_agent_common" } +subprocess = "0.2.9" +tarpc = { version = "0.33.0", features = ["full"] } +tokio = { version = "1.32.0", features = ["full"] } +futures-util = "0.3.28" +futures = "0.3.28" diff --git a/crates/bh_agent_server/src/lib.rs b/crates/bh_agent_server/src/lib.rs new file mode 100644 index 0000000..44fefb5 --- /dev/null +++ b/crates/bh_agent_server/src/lib.rs @@ -0,0 +1,5 @@ +pub mod server; +mod state; +pub mod util; + +pub use server::BhAgentServer; diff --git a/crates/bh_agent_server/src/main.rs b/crates/bh_agent_server/src/main.rs new file mode 100644 index 0000000..be840de --- /dev/null +++ b/crates/bh_agent_server/src/main.rs @@ -0,0 +1,51 @@ +use std::net::IpAddr; +use std::str::FromStr; + +use anyhow::Result; +use futures::{future, prelude::*}; +use tarpc::{ + server::{self, Channel}, + tokio_serde::formats::Json, +}; + +use bh_agent_common::BhAgentService; +use bh_agent_server::BhAgentServer; + +fn parse_args() -> Result<(IpAddr, u16)> { + let args: Vec = std::env::args().collect(); + if args.len() != 3 { + return Err(anyhow::anyhow!("Usage: {} ", args[0])); + } + + let ip_addr = IpAddr::from_str(&args[1])?; + let port = args[2].parse::()?; + + Ok((ip_addr, port)) +} + +#[tokio::main] +async fn main() -> anyhow::Result<()> { + let server_addr = parse_args().or_else(|e| -> Result<(IpAddr, u16)> { + eprintln!("{}", e); + std::process::exit(1); + })?; + + let mut listener = tarpc::serde_transport::tcp::listen(&server_addr, Json::default).await?; + listener.config_mut().max_frame_length(usize::MAX); + listener + // Ignore accept errors. + .filter_map(|r| future::ready(r.ok())) + .map(server::BaseChannel::with_defaults) + // serve is generated by the service attribute. It takes as input any type implementing + // the generated World trait. + .map(|channel| { + let server = BhAgentServer::new(channel.transport().peer_addr().unwrap()); + channel.execute(server.serve()) + }) + // Max 10 channels. + .buffer_unordered(10) + .for_each(|_| async {}) + .await; + + Ok(()) +} diff --git a/crates/bh_agent_server/src/server.rs b/crates/bh_agent_server/src/server.rs new file mode 100644 index 0000000..05265ed --- /dev/null +++ b/crates/bh_agent_server/src/server.rs @@ -0,0 +1,246 @@ +use std::future::{ready, Ready}; +use std::io::{Seek, SeekFrom, Write}; +use std::net::SocketAddr; +use std::sync::Arc; + +use anyhow::Result; +use tarpc::context::Context; + +use bh_agent_common::AgentError::*; +use bh_agent_common::{ + AgentError, BhAgentService, EnvironmentId, FileId, FileOpenMode, FileOpenType, ProcessChannel, + ProcessId, RemotePOpenConfig, +}; + +use crate::state::BhAgentState; +use crate::util::{read_generic, read_lines}; + +macro_rules! check_env_id { + ($env_id:expr) => { + if $env_id != 0 { + return ready(Err(AgentError::InvalidEnvironmentId)); + } + }; +} + +#[derive(Clone)] +pub struct BhAgentServer { + sockaddr: SocketAddr, + state: Arc, +} + +impl BhAgentServer { + pub fn new(socket_addr: SocketAddr) -> Self { + Self { + sockaddr: socket_addr, + state: Arc::new(BhAgentState::new()), + } + } +} + +#[tarpc::server] +impl BhAgentService for BhAgentServer { + type GetEnvironmentsFut = Ready>; + fn get_environments(self, _: Context) -> Self::GetEnvironmentsFut { + // Our implementation currently only supports the default environment + ready(vec![0]) + } + + type GetTempdirFut = Ready>; + fn get_tempdir(self, _: Context, env_id: EnvironmentId) -> Self::GetTempdirFut { + check_env_id!(env_id); + + ready(Ok("/tmp".to_string())) // TODO: make configurable + } + + type RunCommandFut = Ready>; + fn run_command( + self, + _: Context, + env_id: EnvironmentId, + config: RemotePOpenConfig, + ) -> Self::RunCommandFut { + check_env_id!(env_id); + + ready(self.state.run_command(config)) + } + + type GetProcessChannelFut = Ready>; + fn get_process_channel( + self, + _: Context, + env_id: EnvironmentId, + proc_id: ProcessId, + channel: ProcessChannel, + ) -> Self::GetProcessChannelFut { + check_env_id!(env_id); + + ready(self.state.get_process_channel(&proc_id, channel)) + } + + type FileOpenFut = Ready>; + fn file_open( + self, + _: Context, + env_id: EnvironmentId, + path: String, + mode: FileOpenMode, + type_: FileOpenType, + ) -> Self::FileOpenFut { + check_env_id!(env_id); + + ready(self.state.open_path(path, mode, type_)) + } + + type FileCloseFut = Ready>; + fn file_close(self, _: Context, env_id: EnvironmentId, fd: FileId) -> Self::FileCloseFut { + check_env_id!(env_id); + + ready(self.state.close_file(&fd)) + } + + type FileIsClosedFut = Ready>; + fn file_is_closed( + self, + _: Context, + env_id: EnvironmentId, + fd: FileId, + ) -> Self::FileIsClosedFut { + check_env_id!(env_id); + + ready(self.state.is_file_closed(&fd)) + } + + type FileIsReadableFut = Ready>; + fn file_is_readable( + self, + _: Context, + env_id: EnvironmentId, + fd: FileId, + ) -> Self::FileIsReadableFut { + check_env_id!(env_id); + + ready( + self.state + .file_has_any_mode(&fd, &vec![FileOpenMode::Read, FileOpenMode::Update]), + ) + } + + type FileReadFut = Ready, AgentError>>; + fn file_read( + self, + _: Context, + env_id: EnvironmentId, + fd: FileId, + num_bytes: u32, + ) -> Self::FileReadFut { + check_env_id!(env_id); + + ready( + self.state + .do_mut_operation(&fd, |file| { + read_generic(file, num_bytes, self.state.file_type(&fd)?) + }) + .and_then(|v| v.map_err(|_| IoError)), + ) + } + + type FileReadLinesFut = Ready>, AgentError>>; + fn file_read_lines( + self, + _: Context, + env_id: EnvironmentId, + fd: FileId, + hint: u32, + ) -> Self::FileReadLinesFut { + check_env_id!(env_id); + + // TODO: support hint + ready( + self.state + .do_mut_operation(&fd, |file| read_lines(file).map_err(|_| IoError)) + .and_then(|r| r), + ) + } + + type FileIsSeekableFut = Ready>; + fn file_is_seekable( + self, + _: Context, + env_id: EnvironmentId, + fd: FileId, + ) -> Self::FileIsSeekableFut { + check_env_id!(env_id); + + todo!() + } + + type FileSeekFut = Ready>; + fn file_seek( + self, + _: Context, + env_id: EnvironmentId, + fd: FileId, + offset: i32, + whence: i32, + ) -> Self::FileSeekFut { + check_env_id!(env_id); + + let from = match whence { + 0 => SeekFrom::Start(offset as u64), + 1 => SeekFrom::Current(offset as i64), + 2 => SeekFrom::End(offset as i64), + _ => return ready(Err(AgentError::InvalidSeekWhence)), + }; + + ready( + self.state + .do_mut_operation(&fd, |file| file.seek(from)) + .map(|_| ()), + ) + } + + type FileTellFut = Ready>; + fn file_tell(self, _: Context, env_id: EnvironmentId, fd: FileId) -> Self::FileTellFut { + check_env_id!(env_id); + + todo!() + } + + type FileIsWritableFut = Ready>; + fn file_is_writable( + self, + _: Context, + env_id: EnvironmentId, + fd: FileId, + ) -> Self::FileIsWritableFut { + check_env_id!(env_id); + + ready(self.state.file_has_any_mode( + &fd, + &vec![ + FileOpenMode::Write, + FileOpenMode::ExclusiveWrite, + FileOpenMode::Update, + FileOpenMode::Append, + ], + )) + } + + type FileWriteFut = Ready>; + fn file_write( + self, + _: Context, + env_id: EnvironmentId, + fd: FileId, + data: Vec, + ) -> Self::FileWriteFut { + check_env_id!(env_id); + + ready( + self.state + .do_mut_operation(&fd, |file| file.write(&data)) + .map(|_| ()), + ) + } +} diff --git a/crates/bh_agent_server/src/state.rs b/crates/bh_agent_server/src/state.rs new file mode 100644 index 0000000..5f60ae9 --- /dev/null +++ b/crates/bh_agent_server/src/state.rs @@ -0,0 +1,238 @@ +use std::collections::HashMap; +use std::ffi::OsStr; +use std::fs::{File, OpenOptions}; +use std::sync::{Arc, RwLock}; +use futures_util::TryFutureExt; + +use subprocess::{Popen, PopenConfig}; + +use bh_agent_common::AgentError::{ + InvalidFileDescriptor, InvalidProcessId, IoError, ProcessStartFailure, +}; +use bh_agent_common::{ + AgentError, FileId, FileOpenMode, FileOpenType, ProcessChannel, ProcessId, Redirection, + RemotePOpenConfig, +}; + +// TODO: Someday a simple in-memory key value store might be a good idea +pub struct BhAgentState { + files: RwLock>>>, + file_modes: RwLock>, + file_types: RwLock>, + processes: RwLock>>>, + proc_stdin_ids: RwLock>, + proc_stdout_ids: RwLock>, + proc_stderr_ids: RwLock>, + + next_file_id: RwLock, + next_process_id: RwLock, +} + +impl BhAgentState { + pub fn new() -> BhAgentState { + Self { + files: RwLock::new(HashMap::new()), + 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()), + + next_file_id: RwLock::new(0), + next_process_id: RwLock::new(0), + } + } + + fn take_file_id(&self) -> Result { + let mut next_file_id = self.next_file_id.write()?; + let file_id = *next_file_id; + *next_file_id += 1; + Ok(file_id) + } + + fn take_proc_id(&self) -> Result { + let mut next_process_id = self.next_process_id.write()?; + let process_id = *next_process_id; + *next_process_id += 1; + Ok(process_id) + } + + pub fn file_has_any_mode( + &self, + fd: &FileId, + modes: &Vec, + ) -> Result { + Ok(modes.contains( + self.file_modes + .read()? + .get(&fd) + .ok_or(InvalidFileDescriptor)?, + )) + } + + pub fn file_type(&self, fd: &FileId) -> Result { + Ok(self + .file_types + .read()? + .get(&fd) + .ok_or(InvalidFileDescriptor) + .and_then(|t| Ok(t.clone()))?) + } + + pub fn open_path( + &self, + path: String, + mode: FileOpenMode, + type_: FileOpenType, + ) -> Result { + let mut open_opts = OpenOptions::new(); + match mode { + FileOpenMode::Read => open_opts.read(true), + FileOpenMode::Write => open_opts.write(true).create(true), + FileOpenMode::ExclusiveWrite => open_opts.write(true).create_new(true), + FileOpenMode::Append => open_opts.append(true), + FileOpenMode::Update => open_opts.read(true).write(true), + }; + let file = open_opts.open(&path).map_err(|e| { + eprintln!("Path: {}", path); + eprintln!("Error opening file: {}", e); + IoError + })?; + let file_id = self.take_file_id()?; + self.files + .write()? + .insert(file_id, Arc::new(RwLock::new(file))); + self.file_modes.write()?.insert(file_id, mode); + self.file_types.write()?.insert(file_id, type_); + Ok(file_id) + } + + pub fn run_command(&self, config: RemotePOpenConfig) -> Result { + let mut popenconfig = PopenConfig { + stdin: match config.stdin { + Redirection::None => subprocess::Redirection::None, + Redirection::Save => subprocess::Redirection::Pipe, + }, + stdout: match config.stdout { + Redirection::None => subprocess::Redirection::None, + Redirection::Save => subprocess::Redirection::Pipe, + }, + stderr: match config.stderr { + Redirection::None => subprocess::Redirection::None, + Redirection::Save => subprocess::Redirection::Pipe, + }, + detached: false, + executable: config.executable.map(|s| s.into()), + env: config.env.map(|v| { + v.iter() + .map(|t| (t.0.clone().into(), t.1.clone().into())) + .collect() + }), + cwd: config.cwd.map(|s| s.into()), + ..PopenConfig::default() + }; + #[cfg(unix)] + { + popenconfig.setuid = config.setuid.or(popenconfig.setuid); + popenconfig.setgid = config.setuid.or(popenconfig.setgid); + popenconfig.setpgid = config.setpgid || popenconfig.setpgid; + } + + let proc = Popen::create( + config + .argv + .iter() + .map(|s| OsStr::new(s)) + .collect::>() + .as_slice(), + popenconfig, + ) + .map_err(|_| ProcessStartFailure)?; + + let proc_id = self.take_proc_id()?; + + // Stick the process channels into the file map + if proc.stdin.is_some() { + let file_id = self.take_file_id()?; + self.proc_stdin_ids.write()?.insert(file_id, proc_id); + } + if proc.stdout.is_some() { + let file_id = self.take_file_id()?; + self.proc_stdout_ids.write()?.insert(file_id, proc_id); + } + if proc.stderr.is_some() { + let file_id = self.take_file_id()?; + self.proc_stdout_ids.write()?.insert(file_id, proc_id); + } + + // Move the proc to the process map + self.processes + .write()? + .insert(proc_id, Arc::new(RwLock::new(proc))); + + Ok(proc_id) + } + + pub fn get_process_channel( + &self, + proc_id: &ProcessId, + channel: ProcessChannel, + ) -> Result { + match channel { + ProcessChannel::Stdin => &self.proc_stdin_ids, + ProcessChannel::Stdout => &self.proc_stdout_ids, + ProcessChannel::Stderr => &self.proc_stderr_ids, + } + .read()? + .get(&proc_id) + .map(|i| i.clone()) + .ok_or(InvalidProcessId) + } + + pub fn close_file(&self, fd: &FileId) -> Result<(), AgentError> { + Ok(drop( + self.files + .write()? + .remove(&fd) + .ok_or(InvalidFileDescriptor)?, + )) + } + + pub fn is_file_closed(&self, fd: &FileId) -> Result { + Ok(self.files.read()?.contains_key(&fd)) + } + + pub fn do_mut_operation( + &self, + fd: &FileId, + op: impl Fn(&mut File) -> R, + ) -> Result { + // Get file logic + if let Some(file_lock) = self.files.read()?.get(fd) { + return Ok(op(&mut *file_lock.write()?)); + } + + // If these unwraps fail, the state is bad + if let Some(pid) = self.proc_stdin_ids.read()?.get(&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)); + } + if let Some(pid) = self.proc_stdout_ids.read()?.get(&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)); + } + if let Some(pid) = self.proc_stderr_ids.read()?.get(&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)); + } + + Err(InvalidFileDescriptor) + } +} diff --git a/crates/bh_agent_server/src/util/mod.rs b/crates/bh_agent_server/src/util/mod.rs new file mode 100644 index 0000000..a9b02c1 --- /dev/null +++ b/crates/bh_agent_server/src/util/mod.rs @@ -0,0 +1,5 @@ +mod read_chars; +mod read_lines; + +pub use read_chars::*; +pub use read_lines::read_lines; diff --git a/crates/bh_agent_server/src/util/read_chars.rs b/crates/bh_agent_server/src/util/read_chars.rs new file mode 100644 index 0000000..2a077cf --- /dev/null +++ b/crates/bh_agent_server/src/util/read_chars.rs @@ -0,0 +1,87 @@ +use std::fs::File; +use std::io::Read; + +use anyhow::Result; + +use bh_agent_common::FileOpenType; + +pub fn read_generic(mut file: &File, n: u32, file_type: FileOpenType) -> Result> { + match file_type { + FileOpenType::Binary => { + let mut buffer = vec![0u8; n as usize]; + file.read(&mut buffer)?; + Ok(buffer) + } + FileOpenType::Text => Ok(read_chars(&mut file, n as usize)?), + } +} + +pub fn read_chars(reader: &mut R, n: usize) -> Result> { + let mut buffer = vec![0u8; n]; + let mut result = String::new(); + + let bytes_read = reader.read(&mut buffer)?; + buffer.truncate(bytes_read); // Truncate buffer to actual bytes read + + while result.chars().count() < n && !buffer.is_empty() { + let utf8_str = std::str::from_utf8(&buffer); + + match utf8_str { + Ok(s) => { + result.push_str(s); + break; + } + Err(err) if err.valid_up_to() > 0 => { + let valid_str = std::str::from_utf8(&buffer[0..err.valid_up_to()]).unwrap(); + result.push_str(valid_str); + buffer.drain(0..err.valid_up_to()); + } + _ => {} + } + + if result.chars().count() < n { + let mut additional_buffer = vec![0u8; n - result.chars().count()]; + let additional_bytes = reader.read(&mut additional_buffer)?; + if additional_bytes == 0 { + break; + } + buffer.extend_from_slice(&additional_buffer[0..additional_bytes]); + } + } + + Ok(result.into_bytes()) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::io; + use std::io::Cursor; + + #[test] + fn test_single_byte_chars() -> io::Result<()> { + let data = "abcdef"; + let mut cursor = Cursor::new(data.as_bytes()); + let result = read_chars(&mut cursor, 3); + assert_eq!(result.unwrap(), b"abc"); + Ok(()) + } + + #[test] + fn test_multi_byte_chars() -> io::Result<()> { + let data = "a😀b"; + let mut cursor = Cursor::new(data.as_bytes()); + let result = read_chars(&mut cursor, 2); + assert_eq!(result.unwrap(), "a😀".as_bytes()); + Ok(()) + } + + #[test] + fn test_mixed_chars() -> io::Result<()> { + let data = "a😀b😂c"; + let mut cursor = Cursor::new(data.as_bytes()); + let result = read_chars(&mut cursor, 4); + assert_eq!(result.unwrap(), "a😀b😂".as_bytes()); + Ok(()) + } +} diff --git a/crates/bh_agent_server/src/util/read_lines.rs b/crates/bh_agent_server/src/util/read_lines.rs new file mode 100644 index 0000000..2438b23 --- /dev/null +++ b/crates/bh_agent_server/src/util/read_lines.rs @@ -0,0 +1,39 @@ +use std::io::Read; + +use anyhow::Result; + +// Function to split Vec into lines +fn split_lines(buffer: Vec) -> Vec> { + buffer + .split(|x: &u8| *x == b'\n' || *x == 0x0D || *x == 0x0A) // '\n', '\r' + .filter(|x| !x.is_empty()) + .map(|x| x.to_vec()) + .collect() +} + +// Function to read File and split lines +pub fn read_lines(file: &mut T) -> Result>> { + let mut buffer = Vec::new(); + file.read_to_end(&mut buffer)?; + Ok(split_lines(buffer)) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_split_lines() { + let buffer = b"line1\nline2\r\nline3\rline4\n".to_vec(); + let result = split_lines(buffer); + assert_eq!( + result, + vec![ + b"line1".to_vec(), + b"line2".to_vec(), + b"line3".to_vec(), + b"line4".to_vec() + ] + ); + } +} diff --git a/pyproject.toml b/pyproject.toml index a200c1c..9c997eb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [build-system] -requires = ["setuptools>=67", "wheel"] -build-backend = "setuptools.build_meta" +requires = ["maturin>=1.0,<2.0"] +build-backend = "maturin" [project] name = "binharness" @@ -28,11 +28,13 @@ dev = [ homepage = "https://github.com/twizmwazin/binharness" repository = "https://github.com/twizmwazin/binarness.git" -[tool.setuptools.dynamic] -version = {attr = "binharness.__version__"} - -[tool.setuptools.packages.find] -exclude = ["binharness.tests*"] +[tool.maturin] +bindings = "pyo3" +manifest-path = "crates/bh_agent_client/Cargo.toml" +python-source = "python" +python-packages = ["binharness"] +strip = true +sdist-generator = "cargo" [tool.pytest.ini_options] addopts = "--tb=native --cov=binharness --cov-report lcov:.lcov --cov-report term-missing" @@ -62,7 +64,7 @@ ignore-init-module-imports = true required-imports = ["from __future__ import annotations"] [tool.ruff.per-file-ignores] -"binharness/tests/*" = [ +"python/tests/*" = [ "D", "S101", ] diff --git a/binharness/__init__.py b/python/binharness/__init__.py similarity index 100% rename from binharness/__init__.py rename to python/binharness/__init__.py diff --git a/binharness/common/__init__.py b/python/binharness/common/__init__.py similarity index 100% rename from binharness/common/__init__.py rename to python/binharness/common/__init__.py diff --git a/binharness/common/busybox.py b/python/binharness/common/busybox.py similarity index 100% rename from binharness/common/busybox.py rename to python/binharness/common/busybox.py diff --git a/binharness/common/qemu.py b/python/binharness/common/qemu.py similarity index 100% rename from binharness/common/qemu.py rename to python/binharness/common/qemu.py diff --git a/binharness/localenvironment.py b/python/binharness/localenvironment.py similarity index 100% rename from binharness/localenvironment.py rename to python/binharness/localenvironment.py diff --git a/binharness/serialize.py b/python/binharness/serialize.py similarity index 100% rename from binharness/serialize.py rename to python/binharness/serialize.py diff --git a/binharness/types/__init__.py b/python/binharness/types/__init__.py similarity index 100% rename from binharness/types/__init__.py rename to python/binharness/types/__init__.py diff --git a/binharness/types/environment.py b/python/binharness/types/environment.py similarity index 100% rename from binharness/types/environment.py rename to python/binharness/types/environment.py diff --git a/binharness/types/executor.py b/python/binharness/types/executor.py similarity index 100% rename from binharness/types/executor.py rename to python/binharness/types/executor.py diff --git a/binharness/types/injection.py b/python/binharness/types/injection.py similarity index 100% rename from binharness/types/injection.py rename to python/binharness/types/injection.py diff --git a/binharness/types/io.py b/python/binharness/types/io.py similarity index 100% rename from binharness/types/io.py rename to python/binharness/types/io.py diff --git a/binharness/types/process.py b/python/binharness/types/process.py similarity index 100% rename from binharness/types/process.py rename to python/binharness/types/process.py diff --git a/binharness/types/target.py b/python/binharness/types/target.py similarity index 100% rename from binharness/types/target.py rename to python/binharness/types/target.py diff --git a/binharness/util.py b/python/binharness/util.py similarity index 100% rename from binharness/util.py rename to python/binharness/util.py diff --git a/binharness/tests/__init__.py b/python/tests/__init__.py similarity index 100% rename from binharness/tests/__init__.py rename to python/tests/__init__.py diff --git a/binharness/tests/ssh_keys/test b/python/tests/ssh_keys/test similarity index 100% rename from binharness/tests/ssh_keys/test rename to python/tests/ssh_keys/test diff --git a/binharness/tests/ssh_keys/test.pub b/python/tests/ssh_keys/test.pub similarity index 100% rename from binharness/tests/ssh_keys/test.pub rename to python/tests/ssh_keys/test.pub diff --git a/binharness/tests/test_busybox.py b/python/tests/test_busybox.py similarity index 100% rename from binharness/tests/test_busybox.py rename to python/tests/test_busybox.py diff --git a/binharness/tests/test_executor.py b/python/tests/test_executor.py similarity index 99% rename from binharness/tests/test_executor.py rename to python/tests/test_executor.py index 7da0ef9..88021f8 100644 --- a/binharness/tests/test_executor.py +++ b/python/tests/test_executor.py @@ -3,7 +3,6 @@ from pathlib import Path import pytest - from binharness.common.busybox import BusyboxShellExecutor from binharness.localenvironment import LocalEnvironment from binharness.types.executor import ( diff --git a/binharness/tests/test_inject.py b/python/tests/test_inject.py similarity index 99% rename from binharness/tests/test_inject.py rename to python/tests/test_inject.py index bace783..84db91b 100644 --- a/binharness/tests/test_inject.py +++ b/python/tests/test_inject.py @@ -3,7 +3,6 @@ from pathlib import Path import pytest - from binharness.localenvironment import LocalEnvironment from binharness.types.injection import ( ExecutableInjection, diff --git a/binharness/tests/test_localenvironment.py b/python/tests/test_localenvironment.py similarity index 100% rename from binharness/tests/test_localenvironment.py rename to python/tests/test_localenvironment.py diff --git a/binharness/tests/test_qemu_executor.py b/python/tests/test_qemu_executor.py similarity index 100% rename from binharness/tests/test_qemu_executor.py rename to python/tests/test_qemu_executor.py diff --git a/binharness/tests/test_target_local.py b/python/tests/test_target_local.py similarity index 100% rename from binharness/tests/test_target_local.py rename to python/tests/test_target_local.py diff --git a/binharness/tests/test_target_serialization.py b/python/tests/test_target_serialization.py similarity index 99% rename from binharness/tests/test_target_serialization.py rename to python/tests/test_target_serialization.py index 77936aa..addb1ce 100644 --- a/binharness/tests/test_target_serialization.py +++ b/python/tests/test_target_serialization.py @@ -5,7 +5,6 @@ from pathlib import Path import pytest - from binharness.localenvironment import LocalEnvironment from binharness.serialize import TargetImportError, export_target, import_target from binharness.types.executor import NullExecutor diff --git a/binharness/tests/test_util.py b/python/tests/test_util.py similarity index 100% rename from binharness/tests/test_util.py rename to python/tests/test_util.py