Merge pull request #12 from totvscloud-seginf/main

Load entities and policies from json file

Author: omer9564

README.md

@@ -95,6 +95,12 @@ Cedar Agent configuration is available using environment variables and command l
 - The log level to filter logs. Defaults to `info`.  
   `LOG_LEVEL` environment variable.  
   `--log-level`, `-l` command line argument.
+- Load data from json file. Defaults to `None`.  
+  `DATA` environment variable.
+  `--data`, `-d` command line argument.
+- Load policies from json file. Defaults to `None`.
+  `POLICIES` environment variable.
+  `--policies` command line argument.
 
 **command line arguments take precedence over environment variables when configuring the Cedar Agent**
 

src/config.rs

@@ -1,6 +1,7 @@
 use fmt::Debug;
 use std::borrow::Borrow;
 use std::fmt;
+use std::path::PathBuf;
 
 use clap::Parser;
 use log::LevelFilter;
@@ -18,6 +19,10 @@ pub struct Config {
     pub port: Option<u16>,
     #[arg(short, long, value_enum)]
     pub log_level: Option<LevelFilter>,
+    #[arg(short, long)]
+    pub data: Option<PathBuf>,
+    #[arg(long)]
+    pub policies: Option<PathBuf>,
 }
 
 impl Into<rocket::figment::Figment> for &Config {
@@ -34,6 +39,12 @@ impl Into<rocket::figment::Figment> for &Config {
         } else {
             config = config.merge(("port", 8180))
         }
+        if let Some(data) = self.data.borrow() {
+            config = config.merge(("data", data));
+        }
+        if let Some(policies) = self.policies.borrow() {
+            config = config.merge(("policies", policies));
+        }
 
         config
     }
@@ -46,6 +57,8 @@ impl Config {
             addr: None,
             port: None,
             log_level: None,
+            data: None,
+            policies: None,
         }
     }
 
@@ -56,6 +69,8 @@ impl Config {
             config.addr = c.addr.or(config.addr);
             config.port = c.port.or(config.port);
             config.log_level = c.log_level.or(config.log_level);
+            config.data = c.data.or(config.data);
+            config.policies = c.policies.or(config.policies);
         }
 
         config

src/main.rs

@@ -29,6 +29,8 @@ async fn main() {
     let server_config: rocket::figment::Figment = config.borrow().into();
     let launch_result = rocket::custom(server_config)
         .attach(common::DefaultContentType::new(ContentType::JSON))
+        .attach(services::data::load_from_file::InitDataFairing)
+        .attach(services::policies::load_from_file::InitPoliciesFairing)
         .manage(config)
         .manage(Box::new(MemoryPolicyStore::new()) as Box<dyn PolicyStore>)
         .manage(Box::new(MemoryDataStore::new()) as Box<dyn DataStore>)

src/services/data/load_from_file.rs

@@ -0,0 +1,92 @@
+use std::path::PathBuf;
+use std::error::Error;
+use std::fs::File;
+use std::io::Read;
+use log::{error, info};
+
+use rocket::fairing::{Fairing, Info, Kind};
+use rocket::Rocket;
+use rocket::Build;
+
+use crate::services::data::DataStore;
+use crate::config;
+use crate::schemas::data::Entities;
+
+pub struct InitDataFairing;
+
+pub(crate) async fn init(conf: &config::Config, data_store: &Box<dyn DataStore>) {
+
+    if conf.data.is_none() {
+        return;
+    }
+
+    let file_path = conf.data.clone().unwrap();
+    let entities_file_path = &file_path;
+    let entities = match load_entities_from_file(entities_file_path.to_path_buf()).await {
+        Ok(entities) => entities,
+        Err(err) => {
+            error!("Failed to load entities from file: {}", err);
+            return;
+        }
+    };
+
+    match data_store.update_entities(entities).await {
+        Ok(entities) => {
+            info!("Successfully updated entities from file {}: {} entities", &file_path.display(), entities.len());
+        }
+        Err(err) => {
+            error!("Failed to update entities: {}", err);
+            return;
+        }
+    };
+}
+
+pub async fn load_entities_from_file(path: PathBuf) -> Result<Entities, Box<dyn Error>> {
+    
+    if !path.try_exists().unwrap_or(false) || !path.is_file() {
+        return Err("File does not exist".into());
+    }
+    
+    if path.extension().unwrap() != "json" {
+        return Err("File is not a json file".into());
+    }
+
+    let mut file = match File::open(&path) {
+        Ok(file) => file,
+        Err(err) => return Err(format!("Failed to open file: {}", err).into()),
+    };
+
+    let mut contents = String::new();
+    if let Err(err) = file.read_to_string(&mut contents) {
+        return Err(format!("Failed to read file: {}", err).into());
+    }
+
+    let entities: Entities = match rocket::serde::json::from_str(&contents) {
+        Ok(entities) => entities,
+        Err(err) => return Err(format!("Failed to deserialize JSON: {}", err).into()),
+    };
+    
+    Ok(entities)
+}
+
+#[async_trait::async_trait]
+impl Fairing for InitDataFairing {
+    async fn on_ignite(&self, rocket: Rocket<Build>) -> Result<Rocket<Build>, Rocket<Build>> {
+        let config = rocket.state::<config::Config>();
+
+        if config.is_none() {
+            return Ok(rocket);
+        }
+
+        init(config.unwrap(), rocket.state::<Box<dyn DataStore>>().unwrap()).await;
+
+        Ok(rocket)
+    }
+
+    fn info(&self) -> Info {
+        Info {
+            name: "Init Data",
+            kind: Kind::Ignite
+        }
+    }
+}

src/services/data/mod.rs

@@ -5,6 +5,7 @@ use async_trait::async_trait;
 use crate::schemas::data as schemas;
 
 pub mod memory;
+pub mod load_from_file;
 
 #[async_trait]
 pub trait DataStore: Send + Sync {

src/services/policies/load_from_file.rs

@@ -0,0 +1,93 @@
+use std::path::PathBuf;
+use std::error::Error;
+use std::fs::File;
+use std::io::Read;
+use log::{error, info};
+
+use rocket::fairing::{Fairing, Info, Kind};
+use rocket::serde::json::Json;
+use rocket::Rocket;
+use rocket::Build;
+
+use crate::services::policies::PolicyStore;
+use crate::schemas::policies::Policy;
+use crate::config;
+
+pub struct InitPoliciesFairing;
+
+pub(crate) async fn init(conf: &config::Config, policy_store: &Box<dyn PolicyStore>) {
+
+    if conf.policies.is_none() {
+        return;
+    }
+
+    let file_path = conf.policies.clone().unwrap();
+    let policies_file_path = &file_path;
+    let policies = match load_policies_from_file(policies_file_path.to_path_buf()).await {
+        Ok(policies) => policies,
+        Err(err) => {
+            error!("Failed to load policies from file: {}", err);
+            return;
+        }
+    };
+
+    match policy_store.update_policies(policies.into_inner()).await {
+        Ok(policies) => {
+            info!("Successfully updated policies from file {}: {} policies", &file_path.display(), policies.len());
+        }
+        Err(err) => {
+            error!("Failed to update policies: {}", err);
+            return;
+        }
+    };
+}
+
+pub async fn load_policies_from_file(path: PathBuf) -> Result<Json<Vec<Policy>>, Box<dyn Error>> {
+
+    if !path.try_exists().unwrap_or(false) || !path.is_file() {
+        return Err("File does not exist".into());
+    }
+
+    if path.extension().unwrap() != "json" {
+        return Err("File is not a json file".into());
+    }
+
+    let mut file = match File::open(&path) {
+        Ok(file) => file,
+        Err(err) => return Err(format!("Failed to open file: {}", err).into()),
+    };
+
+    let mut contents = String::new();
+    if let Err(err) = file.read_to_string(&mut contents) {
+        return Err(format!("Failed to read file: {}", err).into());
+    }
+
+    let policies: Vec<Policy> = match rocket::serde::json::from_str(&contents) {
+        Ok(policies) => policies,
+        Err(err) => return Err(format!("Failed to deserialize JSON: {}", err).into()),
+    };
+    
+    Ok(Json(policies))
+}
+
+#[async_trait::async_trait]
+impl Fairing for InitPoliciesFairing {
+    async fn on_ignite(&self, rocket: Rocket<Build>) -> Result<Rocket<Build>, Rocket<Build>> {
+        let config = rocket.state::<config::Config>();
+
+        if config.is_none() {
+            return Ok(rocket);
+        }
+
+        init(config.unwrap(), rocket.state::<Box<dyn PolicyStore>>().unwrap()).await;
+
+        Ok(rocket)
+    }
+
+    fn info(&self) -> Info {
+        Info {
+            name: "Init Policies",
+            kind: Kind::Ignite
+        }
+    }
+}

src/services/policies/mod.rs

@@ -7,6 +7,7 @@ use crate::schemas::policies::{Policy, PolicyUpdate};
 
 pub(crate) mod errors;
 pub mod memory;
+pub mod load_from_file;
 
 #[async_trait]
 pub trait PolicyStore: Send + Sync {

tests/services/data_tests.rs

@@ -1,6 +1,9 @@
+use std::path::PathBuf;
+
 use crate::services::utils;
-use cedar_agent::data::memory::MemoryDataStore;
 
+use cedar_agent::data::memory::MemoryDataStore;
+use cedar_agent::data::load_from_file::load_entities_from_file;
 use cedar_agent::DataStore;
 
 #[tokio::test]
@@ -18,3 +21,9 @@ async fn memory_tests() {
     let entities = store.get_entities().await;
     assert_eq!(entities.len(), 0);
 }
+
+#[tokio::test]
+async fn test_load_entities_from_file() {
+    let entities = load_entities_from_file(PathBuf::from("./examples/data.json")).await.unwrap();
+    assert_eq!(entities.len(), 12);
+}

tests/services/policies_tests.rs

@@ -1,11 +1,14 @@
 use std::str::FromStr;
+use std::path::PathBuf;
 
 use cedar_policy::PolicyId;
 
 use crate::services::utils::*;
+
 use cedar_agent::policies::memory::MemoryPolicyStore;
 use cedar_agent::schemas::policies::PolicyUpdate;
 use cedar_agent::PolicyStore;
+use cedar_agent::policies::load_from_file::load_policies_from_file;
 
 #[tokio::test]
 async fn memory_tests() {
@@ -82,3 +85,10 @@ async fn memory_tests() {
         .policy(&PolicyId::from_str("test").unwrap())
         .is_none());
 }
+
+#[tokio::test]
+async fn test_load_policies_from_file() {
+    let policies = load_policies_from_file(PathBuf::from("./examples/policies.json")).await.unwrap();
+    assert_eq!(policies.len(), 3);
+    assert_eq!(policies[0].id, "admins-policy".to_string());
+}