policy: tokenize command line and environment variables

src/agent/src/main.rs

@@ -75,7 +75,7 @@ use tokio::{
 mod rpc;
 mod tracer;
 
-#[cfg(feature = "agent-policy")]
+// #[cfg(feature = "agent-policy")]
 mod policy;
 
 cfg_if! {

src/agent/src/policy.rs

@@ -25,6 +25,70 @@ macro_rules! sl {
     };
 }
 
+#[cfg(test)]
+mod tests {
+    use super::*; // Import the function and structs to be tested
+
+    #[tokio::test]
+    async fn test_allow_request() {
+        let mut policy = AgentPolicy::new();
+
+        // Mock endpoint and input
+        // let request = protocols::agent::CreateContainerRequest {
+        //     ..Default::default()
+        // };
+
+        // Deserialize the JSON content into the Config struct
+        // let data =
+        //     std::fs::read_to_string("../tools/genpolicy/create.json").expect("Unable to read file");
+        // let request: protocols::agent::CreateContainerRequest =
+        //     serde_json::from_str(&data).expect("JSON was not well-formatted");
+        // let request = serde_json::to_string(&request).unwrap();
+        // let ep = "CreateContainerRequest";
+
+        let data =
+            std::fs::read_to_string("../tools/genpolicy/create.json").expect("Unable to read file");
+        // println!("Data: {}", data);
+        let request: protocols::agent::CreateContainerRequest =
+            serde_json::from_str(&data).expect("JSON was not well-formatted");
+        // println!("Request: {}", request);
+        let request = serde_json::to_string(&request).unwrap();
+        let ep = "CreateContainerRequest";
+
+        // println!("Request: {}", request);
+
+        // let request = r#"{"key": "value"}"#;
+
+        // Mock policy initialization (if needed)
+        policy
+            .engine
+            .add_policy_from_file("../tools/genpolicy/exec2.rego")
+            .unwrap();
+
+        // Call the function
+        let result = policy.allow_request(ep, &request).await;
+
+        // let x = "while true; do echo Kubernetes; echo aks-nodepool1-40948945-vmss000000; sleep 10; done";
+
+        // let split = shlex::split(x).unwrap();
+        // println!("shlex split {:?}", split);
+
+        // let split = shell_words::split(x).unwrap();
+        // println!("shell_words split {:?}", split);
+
+        // let split = shellwords::split(x).unwrap();
+        // println!("shellwords split {:?}", split);
+
+        // Assert the expected result
+        match result {
+            Ok((allowed, _)) => assert!(allowed, "Expected the request to be allowed"),
+            Err(e) => panic!("Unexpected error: {:?}", e),
+        }
+
+        // assert!(false, "fail");
+    }
+}
+
 async fn allow_request(policy: &mut AgentPolicy, ep: &str, request: &str) -> ttrpc::Result<()> {
     match policy.allow_request(ep, request).await {
         Ok((allowed, prints)) => {
@@ -120,6 +184,43 @@ pub struct AgentPolicy {
     engine: regorus::Engine,
 }
 
+#[derive(Debug, serde::Serialize)]
+struct PolicyCreateContainerRequest {
+    base: protocols::agent::CreateContainerRequest,
+    tokenized_args: Vec<Vec<String>>,
+    env_map: std::collections::HashMap<String, String>,
+}
+
+fn map_request(ep: &str, input: &str) -> (String, String) {
+    println!("Mapping request");
+    match ep {
+        "CreateContainerRequest" => {
+            println!("CreateContainerRequest detected");
+            let req: protocols::agent::CreateContainerRequest =
+                serde_json::from_str(input).expect("JSON was not well-formatted");
+            let tokenized_args = oci::get_tokenized_args(&req.OCI.Process.Args);
+
+            // println!("Tokenized args: {:?}", tokenized_args);
+
+            let env_map = oci::get_env_map(&req.OCI.Process.Env);
+            // println!("Env: {:?}", env_map);
+
+            let req_v2 = PolicyCreateContainerRequest {
+                base: req,
+                tokenized_args,
+                env_map,
+            };
+
+            (
+                "PolicyCreateContainerRequest".to_string(),
+                serde_json::to_string(&req_v2).expect("failed to serialize"),
+            )
+        }
+
+        _ => (ep.to_string(), input.to_string()),
+    }
+}
+
 #[derive(serde::Deserialize, Debug)]
 struct MetadataResponse {
     allowed: bool,
@@ -193,6 +294,9 @@ impl AgentPolicy {
     /// Ask regorus if an API call should be allowed or not.
     async fn allow_request(&mut self, ep: &str, ep_input: &str) -> Result<(bool, String)> {
         debug!(sl!(), "policy check: {ep}");
+
+        let (ep, ep_input) = &map_request(ep, ep_input);
+
         self.log_request(ep, ep_input).await;
 
         let query = format!("data.agent_policy.{ep}");
@@ -201,10 +305,13 @@ impl AgentPolicy {
         let results = self.engine.eval_query(query, false)?;
 
         let prints = match self.engine.take_prints() {
-            Ok(p) => p.join(" "),
+            Ok(p) => p.join("\n"),
             Err(e) => format!("Failed to get policy log: {e}"),
         };
 
+        println!("Policy prints: {}", prints);
+        // println!("allow_failures: {}", self.allow_failures);
+
         if results.result.len() != 1 {
             // Results are empty when AllowRequestsFailingPolicy is used to allow a Request that hasn't been defined in the policy
             if self.allow_failures {

src/libs/oci/Cargo.toml

@@ -10,3 +10,4 @@ serde = "1.0.131"
 serde_derive = "1.0.131"
 serde_json = "1.0.73"
 libc = "0.2.112"
+shell-words = "1.1.0"

src/libs/oci/src/lib.rs

@@ -81,6 +81,35 @@ impl Spec {
 
 pub type LinuxRlimit = PosixRlimit;
 
+pub fn get_tokenized_args(args: &Vec<String>) -> Vec<Vec<String>> {
+    let mut tokenized_args: Vec<Vec<String>> = Vec::new();
+    for arg in args {
+        let mut arg_split = shell_words::split(arg).expect("Failed to split");
+        for s in arg_split.iter_mut() {
+            if s.ends_with(';') {
+                s.pop(); // Remove the last character
+            }
+        }
+        tokenized_args.push(arg_split);
+    }
+    tokenized_args
+}
+
+pub fn get_env_map(env: &Vec<String>) -> std::collections::HashMap<String, String> {
+    let env_map: std::collections::HashMap<String, String> = env
+        .iter()
+        .filter_map(|v| {
+            let mut split = v.splitn(2, "=");
+            if let (Some(key), Some(value)) = (split.next(), split.next()) {
+                Some((key.to_string(), value.to_string()))
+            } else {
+                None
+            }
+        })
+        .collect();
+    env_map
+}
+
 #[derive(Serialize, Deserialize, Debug, Default, Clone, PartialEq)]
 pub struct Process {
     #[serde(default)]