azure_pim_cli/
graph.rs

1use crate::{az_cli::TokenScope, PimClient};
2use anyhow::{bail, Context, Result};
3use futures::future::join_all;
4use reqwest::Method;
5use serde::{Deserialize, Serialize};
6use serde_json::Value;
7use std::collections::{BTreeMap, BTreeSet};
8use tracing::info;
9
10#[derive(Deserialize, Serialize, PartialOrd, Ord, PartialEq, Eq, Debug, Clone)]
11pub struct Object {
12    pub id: String,
13    pub display_name: String,
14    #[serde(skip_serializing_if = "Option::is_none")]
15    pub upn: Option<String>,
16    pub object_type: PrincipalType,
17}
18
19#[derive(Deserialize, Serialize, PartialOrd, Ord, PartialEq, Eq, Debug, Clone)]
20pub enum PrincipalType {
21    User,
22    Group,
23    ServicePrincipal,
24}
25
26fn parse_objects(value: &Value) -> Result<BTreeSet<Object>> {
27    let mut results = BTreeSet::new();
28    if let Some(values) = value.get("value").and_then(|x| x.as_array()) {
29        for value in values {
30            let Some(id) = value
31                .get("id")
32                .map(|v| v.as_str().unwrap_or(""))
33                .map(ToString::to_string)
34            else {
35                continue;
36            };
37
38            let Some(display_name) = value
39                .get("displayName")
40                .map(|v| v.as_str().unwrap_or(""))
41                .map(ToString::to_string)
42            else {
43                continue;
44            };
45
46            let upn = value
47                .get("userPrincipalName")
48                .and_then(|v| v.as_str())
49                .map(ToString::to_string);
50
51            let data_type = value
52                .get("@odata.type")
53                .map(|x| x.as_str().unwrap_or(""))
54                .context("missing @odata.type")?;
55
56            let object_type = match data_type {
57                "#microsoft.graph.user" => PrincipalType::User,
58                "#microsoft.graph.group" => PrincipalType::Group,
59                "#microsoft.graph.servicePrincipal" => PrincipalType::ServicePrincipal,
60                _ => {
61                    bail!("unknown object type: {data_type} - {value:#?}");
62                }
63            };
64            results.insert(Object {
65                id,
66                display_name,
67                upn,
68                object_type,
69            });
70        }
71    }
72
73    Ok(results)
74}
75
76async fn get_objects_by_ids_small(
77    pim_client: &PimClient,
78    ids: &[&&str],
79) -> Result<BTreeSet<Object>> {
80    info!("checking {} objects", ids.len());
81    let builder = pim_client
82        .backend
83        .client
84        .request(
85            Method::POST,
86            "https://graph.microsoft.com/v1.0/directoryObjects/getByIds",
87        )
88        .bearer_auth(pim_client.backend.get_token(TokenScope::Graph).await?);
89
90    let body = serde_json::json!({ "ids": ids });
91    let request = builder.json(&body).build()?;
92    let value = pim_client.backend.retry_request(&request, None).await?;
93
94    parse_objects(&value)
95}
96
97pub(crate) async fn get_objects_by_ids(
98    pim_client: &PimClient,
99    ids: BTreeSet<&str>,
100) -> Result<BTreeMap<String, Object>> {
101    let mut cache = pim_client.object_cache.lock().await;
102    let to_update = ids
103        .iter()
104        .filter(|id| !cache.contains_key(**id))
105        .collect::<Vec<_>>();
106
107    let chunks = to_update.chunks(50).collect::<Vec<_>>();
108
109    let results = join_all(
110        chunks
111            .iter()
112            .map(|chunk| get_objects_by_ids_small(pim_client, chunk)),
113    )
114    .await;
115
116    for entry in results {
117        for entry in entry? {
118            cache.insert(entry.id.clone(), Some(entry));
119        }
120    }
121
122    let mut result = BTreeMap::new();
123    for id in ids {
124        if let Some(entry) = cache.get(id).cloned() {
125            if let Some(entry) = entry {
126                result.insert(entry.id.clone(), entry);
127            }
128        } else {
129            cache.insert(id.to_string(), None);
130        }
131    }
132
133    Ok(result)
134}
135
136pub(crate) async fn group_members(pim_client: &PimClient, id: &str) -> Result<BTreeSet<Object>> {
137    let mut group_cache = pim_client.group_cache.lock().await;
138    if let Some(entries) = group_cache.get(id) {
139        return Ok(entries.clone());
140    }
141
142    let mut cache = pim_client.object_cache.lock().await;
143
144    let url = format!("https://graph.microsoft.com/v1.0/groups/{id}/members");
145    let request = pim_client
146        .backend
147        .client
148        .request(Method::GET, &url)
149        .bearer_auth(pim_client.backend.get_token(TokenScope::Graph).await?)
150        .build()?;
151    let value = pim_client.backend.retry_request(&request, None).await?;
152    let results = parse_objects(&value)?;
153
154    for object in &results {
155        if cache.get(&object.id).is_none() {
156            cache.insert(object.id.clone(), Some(object.clone()));
157        }
158    }
159
160    group_cache.insert(id.to_string(), results.clone());
161
162    Ok(results)
163}
azure_pim_cli/graph.rs

azure_pim_cli/
graph.rs
