-
Notifications
You must be signed in to change notification settings - Fork 6
/
add_users_and_groups.py
71 lines (66 loc) · 2.47 KB
/
add_users_and_groups.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
import boto3
import os
import getpass
cognito_client = boto3.client('cognito-idp')
user_pools = cognito_client.list_user_pools(MaxResults=60)['UserPools']
user_pool_id = [user_pool['Id'] for user_pool in user_pools if user_pool['Name']=='mlflow-user-pool'][0]
groups = ['admins', 'readers', 'model-approvers']
list_groups = cognito_client.list_groups(UserPoolId=user_pool_id)['Groups']
existing_group_names = [group['GroupName'] for group in list_groups]
users_groups = [
{
'username': 'mlflow-admin@example.com',
'group': 'admins'
},
{
'username': 'mlflow-reader@example.com',
'group': 'readers',
},
{
'username': 'mlflow-model-approver@example.com',
'group': 'model-approvers'
}
]
list_users = cognito_client.list_users(UserPoolId=user_pool_id)['Users']
existing_email_list = []
for user in list_users:
attributes = user['Attributes']
email = [attribute['Value'] for attribute in attributes if attribute['Name']=='email'][0]
existing_email_list.append(email)
if __name__=="__main__":
# Create groups
for group in groups:
if group in existing_group_names:
print(f"group {group} already exists")
else:
print(f"create group {group} for cognito user pool {user_pool_id}")
cognito_client.create_group(
GroupName=group,
UserPoolId=user_pool_id
)
# Create users and associate them with a group
for user_group in users_groups:
username = user_group['username']
group = user_group['group']
if username in existing_email_list:
print(f"user {username} already exist. skip it")
else:
print(f"create user {username}")
cognito_client.admin_create_user(
UserPoolId=user_pool_id,
Username=username,
#TemporaryPassword=args.password
)
pwd = getpass.getpass(prompt = f"Enter the password for {username}: ")
cognito_client.admin_set_user_password(
UserPoolId=user_pool_id,
Username=username,
Password=pwd,
Permanent=True # does not force a user to change the password
)
print(f"add user {username} to group {group}")
cognito_client.admin_add_user_to_group(
UserPoolId=user_pool_id,
Username=username,
GroupName=group
)