cognito_user_reader/
cognito.rs

1#![allow(clippy::missing_errors_doc)]
2
3use crate::{
4    emojis::{ERROR, ROCKET},
5    UserTypeExt,
6};
7
8use chrono::prelude::*;
9use console::style;
10use rusoto_cognito_idp::{
11    CognitoIdentityProvider, CognitoIdentityProviderClient, ListUsersRequest, ListUsersResponse,
12    UserType,
13};
14use rusoto_core::Region;
15use std::str::FromStr;
16
17pub struct UserReader {
18    pub aws_pool_id: String,
19    pub aws_region: String,
20    cognito_provider: CognitoIdentityProviderClient,
21}
22
23pub struct UserReaderOptions<'a> {
24    pub attributes_to_get: &'a Option<Vec<String>>,
25    pub limit_of_users: Option<i64>,
26    pub show_unconfirmed_users: bool,
27    pub filtered_user_ids: &'a Option<Vec<String>>,
28    pub include_user_ids: bool,
29    pub filtered_user_emails: &'a Option<Vec<String>>,
30    pub include_user_emails: bool,
31    pub created_at: Option<DateTime<Utc>>,
32}
33
34impl UserReader {
35    #[must_use]
36    pub fn new(aws_pool_id: String) -> Self {
37        let raw_aws_region = Self::extract_region(&aws_pool_id);
38        let aws_region: Region =
39            Region::from_str(&raw_aws_region).expect("Wrong format for this pool id.");
40        let cognito_provider = CognitoIdentityProviderClient::new(aws_region);
41
42        Self {
43            aws_pool_id,
44            aws_region: raw_aws_region,
45            cognito_provider,
46        }
47    }
48
49    #[must_use]
50    pub fn extract_region(pool_id: &str) -> String {
51        pool_id
52            .split('_')
53            .next()
54            .expect("Impossible to get the region from the pool id")
55            .to_owned()
56    }
57
58    pub async fn get_users(
59        &self,
60        options: UserReaderOptions<'_>,
61        show_messages: bool,
62    ) -> Vec<UserType> {
63        let mut users: Vec<UserType> = Vec::new();
64        let mut pending_users: i64 = 0;
65        let mut limit: Option<i64> = None;
66
67        if let Some(max_users) = options.limit_of_users {
68            println!("------ max users {}", max_users);
69            if max_users <= 60 {
70                limit = Some(max_users);
71            } else {
72                limit = Some(60);
73                pending_users = max_users - 60;
74            }
75        }
76
77        let mut req = ListUsersRequest {
78            user_pool_id: self.aws_pool_id.clone(),
79            attributes_to_get: options.attributes_to_get.clone(),
80            filter: if options.show_unconfirmed_users {
81                None
82            } else {
83                Some("cognito:user_status = 'CONFIRMED'".to_string())
84            },
85            limit,
86            pagination_token: None,
87        };
88
89        // loop until we get all the users that we want
90        loop {
91            match self.cognito_provider.list_users(req.clone()).await {
92                Ok(ListUsersResponse {
93                    pagination_token,
94                    users: Some(mut response_users),
95                }) => {
96                    if show_messages {
97                        println!(
98                            "{} {} {}",
99                            ROCKET,
100                            style(format!("We got a batch of {} users", response_users.len()))
101                                .bold()
102                                .green(),
103                            ROCKET
104                        );
105                    }
106                    req.pagination_token = pagination_token;
107                    users.append(&mut response_users);
108                }
109                Err(e) => {
110                    if show_messages {
111                        println!(
112                            "{} {} {}\n{}",
113                            ERROR,
114                            style("SOMETHING WENT WRONG!").bold().red(),
115                            ERROR,
116                            style(e).red(),
117                        );
118                    }
119                    req.pagination_token = None;
120                }
121                Ok(_x) => (),
122            }
123
124            if req.limit.is_some() {
125                if pending_users == 0 {
126                    break;
127                } else if pending_users <= 60 {
128                    req.limit = Some(pending_users);
129                    pending_users = 0;
130                } else {
131                    req.limit = Some(60);
132                    pending_users -= 60;
133                }
134            }
135
136            if req.pagination_token.is_none() {
137                break;
138            }
139        }
140
141        Self::order_users(users, &options)
142    }
143
144    fn order_users(mut users: Vec<UserType>, options: &UserReaderOptions<'_>) -> Vec<UserType> {
145        // order by creation date
146        users.sort_by(|a, b| {
147            a.user_create_date
148                .partial_cmp(&b.user_create_date)
149                .unwrap_or(std::cmp::Ordering::Less)
150        });
151
152        // apply filters
153        users
154            .into_iter()
155            .filter(|u| {
156                if let Some(ref avoid) = options.filtered_user_ids {
157                    let username = u.username.as_deref().unwrap_or("");
158                    let is_in = avoid.contains(&username.to_string());
159                    return if options.include_user_ids {
160                        is_in
161                    } else {
162                        !is_in
163                    };
164                }
165                true
166            })
167            .filter(|u| {
168                if let Some(ref avoid) = options.filtered_user_emails {
169                    let is_in = avoid
170                        .iter()
171                        .any(|e| e.to_lowercase() == (&u.get_email()).to_lowercase());
172                    return if options.include_user_emails {
173                        is_in
174                    } else {
175                        !is_in
176                    };
177                }
178                true
179            })
180            .filter(|u| {
181                if let Some(created_at) = options.created_at {
182                    let duration = &u.creation_date().signed_duration_since(created_at);
183                    return duration.num_days() >= 0;
184                }
185                true
186            })
187            .collect()
188    }
189}
190
191#[must_use]
192pub fn users_to_csv(users: &[UserType], print_screen: bool) -> (String, i32) {
193    let mut filtered_len = 0;
194
195    let content = users.iter().fold(String::new(), |acc, u| {
196        let creation_date = u.creation_date();
197        let username = u.username.as_deref().unwrap_or("No username");
198        let user_status = u.user_status.as_deref().unwrap_or("No user status");
199        if print_screen {
200            println!(
201                "{} | {} | {} | {}",
202                style(creation_date).red(),
203                style(username).green(),
204                style(user_status).yellow(),
205                u.attributes_values_to_string(" | "),
206            );
207        }
208        filtered_len += 1;
209        format!(
210            "{}\n{}",
211            if acc.is_empty() {
212                format!(
213                    "createdAt,username,status,{}",
214                    u.attributes_keys_to_string(",")
215                )
216            } else {
217                acc
218            },
219            format!(
220                "{},{},{},{}",
221                creation_date,
222                username,
223                user_status,
224                u.attributes_values_to_string(","),
225            )
226        )
227    });
228    (content, filtered_len)
229}