refactor: move functions to repository type

4 files changed, 108 insertions, 129 deletions

diff --git a/src/authorization/entities.rs b/src/authorization/entities.rs
index 7af5fa4a..8c3f7955 100644
--- a/src/authorization/entities.rs
+++ b/src/authorization/entities.rs
@@ -1,6 +1,5 @@
 use serde::Serialize;
 use std::collections::HashSet;
-use std::fs;
 
 // Cedar entity structures
 #[derive(Debug, Serialize)]
@@ -61,117 +60,118 @@ pub struct Group {
     pub parent_id: Option<u64>,
 }
 
-pub async fn generate_entities_from_api(
-    token: String,
-    host: String,
-    project: String,
-) -> Result<Vec<CedarEntity>, Box<dyn std::error::Error>> {
-    let client = reqwest::Client::new();
-    let mut entities = Vec::new();
-    let mut processed_groups = HashSet::new();
+pub struct EntitiesRepository {
+    pub token: String,
+    pub host: String,
+    pub project: String,
+}
 
-    // Fetch project information
-    let project_url = format!(
-        "{}/api/v4/projects/{}",
-        host.trim_end_matches('/'),
-        urlencoding::encode(&project)
-    );
+impl EntitiesRepository {
+    pub fn new(token: String, host: String, project: String) -> EntitiesRepository {
+        EntitiesRepository {
+            token: token,
+            host: host,
+            project: project,
+        }
+    }
 
-    println!("Fetching project information...");
-    let project: Project = client
-        .get(&project_url)
-        .header("PRIVATE-TOKEN", &token)
-        .send()
-        .await?
-        .error_for_status()?
-        .json()
-        .await?;
+    pub async fn generate(&self) -> Result<Vec<CedarEntity>, Box<dyn std::error::Error>> {
+        let client = reqwest::Client::new();
+        let mut entities = Vec::new();
+        let mut processed_groups = HashSet::new();
 
-    // Add organization
-    entities.push(CedarEntity {
-        uid: CedarUid {
-            entity_type: "Organization".to_string(),
-            id: "1".to_string(),
-        },
-        attrs: serde_json::json!({
-            "name": "gitlab",
-        }),
-        parents: vec![],
-    });
+        let project_url = format!(
+            "{}/api/v4/projects/{}",
+            self.host.trim_end_matches('/'),
+            urlencoding::encode(&self.project)
+        );
 
-    // Add project entity
-    entities.push(CedarEntity {
-        uid: CedarUid {
-            entity_type: "Project".to_string(),
-            id: project.id.to_string(),
-        },
-        attrs: serde_json::json!({
-            "name": project.name,
-            "path": project.path,
-            "full_path": format!("{}/{}", project.namespace.full_path, project.path),
-        }),
-        parents: if project.namespace.kind == "group" {
-            vec![CedarParent {
-                parent_type: "Group".to_string(),
-                id: project.namespace.id.to_string(),
-            }]
-        } else {
-            vec![]
-        },
-    });
+        let project: Project = client
+            .get(&project_url)
+            .header("PRIVATE-TOKEN", &self.token)
+            .send()
+            .await?
+            .error_for_status()?
+            .json()
+            .await?;
 
-    // Fetch project members
-    let members_url = format!(
-        "{}/api/v4/projects/{}/members/all",
-        host.trim_end_matches('/'),
-        project.id
-    );
+        entities.push(CedarEntity {
+            uid: CedarUid {
+                entity_type: "Organization".to_string(),
+                id: "1".to_string(),
+            },
+            attrs: serde_json::json!({
+                "name": "gitlab",
+            }),
+            parents: vec![],
+        });
 
-    println!("Fetching project members...");
-    let members: Vec<Member> = client
-        .get(&members_url)
-        .header("PRIVATE-TOKEN", &token)
-        .send()
-        .await?
-        .error_for_status()?
-        .json()
-        .await?;
+        entities.push(CedarEntity {
+            uid: CedarUid {
+                entity_type: "Project".to_string(),
+                id: project.id.to_string(),
+            },
+            attrs: serde_json::json!({
+                "name": project.name,
+                "path": project.path,
+                "full_path": format!("{}/{}", project.namespace.full_path, project.path),
+            }),
+            parents: if project.namespace.kind == "group" {
+                vec![CedarParent {
+                    parent_type: "Group".to_string(),
+                    id: project.namespace.id.to_string(),
+                }]
+            } else {
+                vec![]
+            },
+        });
+
+        let members_url = format!(
+            "{}/api/v4/projects/{}/members/all",
+            self.host.trim_end_matches('/'),
+            project.id
+        );
+
+        let members: Vec<Member> = client
+            .get(&members_url)
+            .header("PRIVATE-TOKEN", &self.token)
+            .send()
+            .await?
+            .error_for_status()?
+            .json()
+            .await?;
 
-    println!("Found {} members", members.len());
+        for member in members {
+            if member.state == "active" {
+                entities.push(CedarEntity {
+                    uid: CedarUid {
+                        entity_type: "User".to_string(),
+                        id: member.id.to_string(),
+                    },
+                    attrs: serde_json::json!({
+                        "username": member.username,
+                        "name": member.name,
+                        "access_level": member.access_level,
+                    }),
+                    parents: vec![],
+                });
+            }
+        }
 
-    // Add user entities
-    for member in members {
-        if member.state == "active" {
-            entities.push(CedarEntity {
-                uid: CedarUid {
-                    entity_type: "User".to_string(),
-                    id: member.id.to_string(),
-                },
-                attrs: serde_json::json!({
-                    "username": member.username,
-                    "name": member.name,
-                    "access_level": member.access_level,
-                }),
-                parents: vec![],
-            });
+        if project.namespace.kind == "group" {
+            fetch_group_hierarchy(
+                &client,
+                &self.host,
+                &self.token,
+                project.namespace.id,
+                &mut entities,
+                &mut processed_groups,
+            )
+            .await?;
         }
-    }
 
-    // Fetch group hierarchy if project belongs to a group
-    if project.namespace.kind == "group" {
-        println!("Fetching group hierarchy...");
-        fetch_group_hierarchy(
-            &client,
-            &host,
-            &token,
-            project.namespace.id,
-            &mut entities,
-            &mut processed_groups,
-        )
-        .await?;
+        Ok(entities)
     }
-
-    Ok(entities)
 }
 
 pub fn fetch_group_hierarchy<'a>(
@@ -244,24 +244,3 @@ pub fn fetch_group_hierarchy<'a>(
         Ok(())
     })
 }
-
-pub fn write_entities_file(
-    entities: &[CedarEntity],
-    output: &str,
-) -> Result<(), Box<dyn std::error::Error>> {
-    let json = serde_json::to_string_pretty(&entities)?;
-    fs::write(output, json)?;
-
-    println!(
-        "\nSuccessfully wrote {} entities to {}",
-        entities.len(),
-        output
-    );
-    println!("\nTo use these entities with Cedar:");
-    println!(
-        "  let entities = cedar_policy::Entities::from_json_file(\"{}\", None)?;",
-        output
-    );
-
-    Ok(())
-}
diff --git a/src/authorization/mod.rs b/src/authorization/mod.rs
index 7e1d69b5..d687d53f 100644
--- a/src/authorization/mod.rs
+++ b/src/authorization/mod.rs
@@ -7,5 +7,5 @@ pub mod server;
 pub use authorizer::Authorizer;
 pub use cedar_authorizer::CedarAuthorizer;
 pub use check_service::CheckService;
-pub use entities::{CedarEntity, generate_entities_from_api, write_entities_file};
+pub use entities::{CedarEntity, EntitiesRepository};
 pub use server::Server;
diff --git a/src/bin/cli.rs b/src/bin/cli.rs
index 275bd410..d0bbe989 100644
--- a/src/bin/cli.rs
+++ b/src/bin/cli.rs
@@ -1,4 +1,4 @@
-use authzd::{generate_entities_from_api, write_entities_file};
+use authzd::EntitiesRepository;
 use clap::{Parser, Subcommand};
 
 #[derive(Parser, Debug)]
@@ -50,9 +50,10 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
             token,
             host,
         } => {
-            let entities = generate_entities_from_api(token, host, project).await?;
-
-            write_entities_file(&entities, &output)?;
+            let repository = EntitiesRepository::new(token, host, project);
+            let entities = repository.generate().await?;
+            let json = serde_json::to_string_pretty(&entities)?;
+            std::fs::write(output, json)?;
         }
     }
 
diff --git a/src/lib.rs b/src/lib.rs
index cf570238..e70e64f7 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,6 +1,5 @@
 pub mod authorization;
 
 pub use authorization::{
-    Authorizer, CedarAuthorizer, CedarEntity, CheckService, Server, generate_entities_from_api,
-    write_entities_file,
+    Authorizer, CedarAuthorizer, CedarEntity, CheckService, EntitiesRepository, Server,
 };