Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
570 changes: 556 additions & 14 deletions Cargo.lock

Large diffs are not rendered by default.

1 change: 1 addition & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ members = [
"crates/shadi_memory",
"crates/shadi_py",
"crates/shadi_sandbox",
"crates/shadi_telemetry",
"crates/shadictl"
]
resolver = "2"
1 change: 1 addition & 0 deletions agents/secops/telemetry.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
_resource = Resource(
attributes={
RES_SERVICE_NAME: SERVICE_NAME,
"service.namespace": "shadi",
"telemetry.sdk.language": "python",
}
)
Expand Down
2 changes: 2 additions & 0 deletions crates/shadi_py/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,8 @@ doctest = false
agent_secrets = { path = "../agent_secrets", features = ["onepassword"] }
shadi_memory = { path = "../shadi_memory" }
shadi_sandbox = { path = "../shadi_sandbox" }
shadi_telemetry = { path = "../shadi_telemetry" }
tracing = "0.1"

[dependencies.pyo3]
version = "0.21"
Expand Down
41 changes: 41 additions & 0 deletions crates/shadi_py/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ use pyo3::prelude::*;
use pyo3::types::{PyBytes, PyModule};
use shadi_memory::{MemoryEntry as ShadiMemoryEntry, SqlCipherStore};
use shadi_sandbox::{spawn_sandboxed, SandboxError, SandboxPolicy};
use tracing::{field, info_span};

struct SessionFlagVerifier;

Expand Down Expand Up @@ -127,6 +128,8 @@ impl ShadiStore {
}

fn put(&self, session: &PySessionContext, key: &str, secret: &[u8]) -> PyResult<()> {
let span = info_span!("shadi.secret.put", secret.key = %key);
let _guard = span.enter();
let ctx = session.to_context();
let guard = self.store.lock().map_err(|_| PyRuntimeError::new_err("lock poisoned"))?;
let access = AgentSecretAccess::new(guard.as_ref(), &self.verifier);
Expand All @@ -141,6 +144,8 @@ impl ShadiStore {
session: &PySessionContext,
key: &str,
) -> PyResult<Bound<'py, PyBytes>> {
let span = info_span!("shadi.secret.get", secret.key = %key);
let _guard = span.enter();
let ctx = session.to_context();
let guard = self.store.lock().map_err(|_| PyRuntimeError::new_err("lock poisoned"))?;
let access = AgentSecretAccess::new(guard.as_ref(), &self.verifier);
Expand All @@ -150,6 +155,8 @@ impl ShadiStore {
}

fn delete(&self, session: &PySessionContext, key: &str) -> PyResult<()> {
let span = info_span!("shadi.secret.delete", secret.key = %key);
let _guard = span.enter();
let ctx = session.to_context();
let guard = self.store.lock().map_err(|_| PyRuntimeError::new_err("lock poisoned"))?;
let access = AgentSecretAccess::new(guard.as_ref(), &self.verifier);
Expand All @@ -159,6 +166,8 @@ impl ShadiStore {
}

fn list_keys(&self, session: &PySessionContext) -> PyResult<Vec<String>> {
let span = info_span!("shadi.secret.list_keys");
let _guard = span.enter();
let ctx = session.to_context();
AgentSecretAccess::require_verified(&ctx).map_err(map_secret_error)?;
let guard = self.store.lock().map_err(|_| PyRuntimeError::new_err("lock poisoned"))?;
Expand All @@ -178,12 +187,16 @@ impl SqlCipherMemoryStore {
}

fn put(&self, scope: &str, entry_key: &str, payload: &str) -> PyResult<i64> {
let span = info_span!("shadi.memory.put", memory.scope = %scope, memory.entry_key = %entry_key);
let _guard = span.enter();
self.store
.put(scope, entry_key, payload)
.map_err(|err| PyRuntimeError::new_err(err.to_string()))
}

fn get_latest(&self, scope: &str, entry_key: &str) -> PyResult<Option<MemoryEntry>> {
let span = info_span!("shadi.memory.get_latest", memory.scope = %scope, memory.entry_key = %entry_key);
let _guard = span.enter();
let entry = self
.store
.get_latest(scope, entry_key)
Expand All @@ -193,6 +206,13 @@ impl SqlCipherMemoryStore {

#[pyo3(signature = (query, scope=None, limit=10))]
fn search(&self, query: &str, scope: Option<String>, limit: usize) -> PyResult<Vec<MemoryEntry>> {
let span = info_span!(
"shadi.memory.search",
memory.query = %query,
memory.scope = %scope.as_deref().unwrap_or(""),
memory.limit = limit as i64,
);
let _guard = span.enter();
let entries = self
.store
.search(scope.as_deref(), query, limit)
Expand All @@ -205,6 +225,12 @@ impl SqlCipherMemoryStore {

#[pyo3(signature = (scope=None, limit=50))]
fn list(&self, scope: Option<String>, limit: usize) -> PyResult<Vec<MemoryEntry>> {
let span = info_span!(
"shadi.memory.list",
memory.scope = %scope.as_deref().unwrap_or(""),
memory.limit = limit as i64,
);
let _guard = span.enter();
let entries = self
.store
.list(scope.as_deref(), limit)
Expand All @@ -216,6 +242,8 @@ impl SqlCipherMemoryStore {
}

fn delete(&self, scope: &str, entry_key: &str) -> PyResult<usize> {
let span = info_span!("shadi.memory.delete", memory.scope = %scope, memory.entry_key = %entry_key);
let _guard = span.enter();
self.store
.delete(scope, entry_key)
.map_err(|err| PyRuntimeError::new_err(err.to_string()))
Expand Down Expand Up @@ -286,6 +314,7 @@ impl PySessionContext {

#[pymodule]
fn shadi(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
shadi_telemetry::init("shadi-runtime");
m.add_class::<ShadiStore>()?;
m.add_class::<PySessionContext>()?;
m.add_class::<SqlCipherMemoryStore>()?;
Expand Down Expand Up @@ -325,6 +354,8 @@ fn inject_keychain_with_store(
command: &mut Command,
mappings: &[String],
) -> Result<(), String> {
let span = info_span!("shadi.secrets.inject", secret.count = mappings.len() as i64);
let _guard = span.enter();
for mapping in mappings {
let (key, env) = parse_key_env(mapping)?;
let secret = store
Expand Down Expand Up @@ -361,6 +392,15 @@ fn run_sandboxed(
return Err(PyRuntimeError::new_err("command must not be empty"));
}

let cwd_value = cwd.as_deref().unwrap_or("");
let span = info_span!(
"shadi.sandbox.run",
command = %command[0],
cwd = %cwd_value,
exit.code = field::Empty,
);
let _guard = span.enter();

let mut cmd = Command::new(&command[0]);
if command.len() > 1 {
cmd.args(&command[1..]);
Expand All @@ -381,6 +421,7 @@ fn run_sandboxed(
let status = child
.wait()
.map_err(|err| PyRuntimeError::new_err(err.to_string()))?;
span.record("exit.code", &status.code().unwrap_or(-1));
Ok(status.code().unwrap_or(1))
}

Expand Down
1 change: 1 addition & 0 deletions crates/shadi_sandbox/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ coverage = []

[dependencies]
thiserror = "1.0"
tracing = "0.1"

[target.'cfg(target_os = "macos")'.dependencies]
libc = "0.2"
Expand Down
38 changes: 37 additions & 1 deletion crates/shadi_sandbox/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +7,32 @@ mod platform;
pub use policy::SandboxPolicy;
use std::process::{Command, ExitStatus};
use std::io;
use tracing::{field, info_span};

pub fn spawn_sandboxed(command: &mut Command, policy: &SandboxPolicy) -> Result<SandboxedChild, SandboxError> {
let program = command.get_program().to_string_lossy().to_string();
let args = command
.get_args()
.map(|arg| arg.to_string_lossy())
.collect::<Vec<_>>()
.join(" ");
let cwd = command
.get_current_dir()
.map(|path| path.display().to_string())
.unwrap_or_else(|| "".to_string());
let allowed_paths = policy.allow_read().len() + policy.allow_write().len();
let network_mode = if policy.net_blocked() { "blocked" } else { "allowed" };

let span = info_span!(
"shadi.sandbox.spawn",
command = %program,
args = %args,
cwd = %cwd,
policy.allowed_paths = allowed_paths as i64,
network.mode = %network_mode,
);
let _guard = span.enter();

platform::spawn_sandboxed(command, policy)
}

Expand Down Expand Up @@ -37,14 +61,26 @@ impl SandboxedChild {
}

pub fn wait(&mut self) -> io::Result<ExitStatus> {
match &mut self.inner {
let span = info_span!("shadi.sandbox.wait", pid = self.id(), exit.code = field::Empty);
let _guard = span.enter();

let status = match &mut self.inner {
SandboxedChildInner::Std(child) => child.wait(),
#[cfg(target_os = "windows")]
SandboxedChildInner::Windows(child) => child.wait(),
};

if let Ok(ref status) = status {
span.record("exit.code", &status.code().unwrap_or(-1));
}

status
}

pub fn kill(&mut self) -> io::Result<()> {
let span = info_span!("shadi.sandbox.kill", pid = self.id());
let _guard = span.enter();

match &mut self.inner {
SandboxedChildInner::Std(child) => child.kill(),
#[cfg(target_os = "windows")]
Expand Down
13 changes: 13 additions & 0 deletions crates/shadi_telemetry/Cargo.toml
Original file line number Diff line number Diff line change
@@ -0,0 +1,13 @@
[package]
name = "shadi_telemetry"
version = "0.1.0"
edition = "2021"

[dependencies]
opentelemetry = "0.24"
opentelemetry_sdk = "0.24"
opentelemetry-otlp = { version = "0.17", features = ["http-proto", "reqwest-client"] }
tracing = "0.1"
tracing-opentelemetry = "0.25"
tracing-subscriber = { version = "0.3", features = ["json"] }
tracing-appender = "0.2"
Loading
Loading