atmosphere/scripts/add_new_accounts.py at f97780a6eaf5ade8fa718897d021c3a5c6fd118b · cyverse/atmosphere · GitHub

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
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
#!/usr/bin/env python

# DEPRECATION WARNING -- Will be removed in favor of periodic task (To facilitate auto-generation of `atmo-user` accounts)
# and the use of the /v2/accounts API.
# FIXME: Add 'account_user, group_name, is_leader' args to this script
import argparse
import libcloud.security

import django

django.setup()

import core.models
from core.models import AtmosphereUser as User
from core.models import Provider, Identity
from core.query import contains_credential
from core.plugins import ValidationPluginManager, ExpirationPluginManager, AccountCreationPluginManager

from django_cyverse_auth.protocol.ldap import get_members
from service.driver import get_account_driver
from threepio import logger

libcloud.security.VERIFY_SSL_CERT = False


def get_usernames(provider):
    """
    """
    return Identity.objects.filter(provider=provider).values_list('created_by__username', flat=True)

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--provider", type=int,
                        help="Atmosphere provider ID"
                        " to use when importing users.")
    parser.add_argument("--provider-id", type=int,
                        help="Atmosphere provider ID"
                        " to use when importing users."
                        " DEPRECATION WARNING -- THIS WILL BE REMOVED SOON!")
    parser.add_argument("--provider-list",
                        action="store_true",
                        help="List of provider names and IDs")
    parser.add_argument("--rebuild", action="store_true",
                        help="Rebuild all accounts that are in the provider")
    parser.add_argument("--group",
                        help="LDAP group of usernames to import.")
    parser.add_argument("--users",
                        help="usernames to add to Atmosphere. (comma separated list with no spaces)")
    parser.add_argument("--admin", action="store_true",
                        help="Users addded as admin and staff users.")
    args = parser.parse_args()

    if args.provider_list:
        print "ID\tName"
        for p in Provider.objects.all().order_by('id'):
            print "%d\t%s" % (p.id, p.location)
        return

    usernames = []
    if args.provider_id and not args.provider:
        print "WARNING: --provider-id has been *DEPRECATED*! Use --provider instead!"
        args.provider = args.provider_id
    if args.provider:
        provider = Provider.objects.get(id=args.provider)
    else:
        raise Exception("Missing required argument: --provider <id>. use --provider-list to get a list of provider ID+names")
    print "Using Provider: %s" % provider
    try:
        acct_driver = get_account_driver(provider, raise_exception=True)
    except:
        account_provider = provider.accountprovider_set.first()
        print "Could not create the account Driver for this Provider."\
              " Check the configuration of this identity:%s" % account_provider
        raise
    if args.group:
        print "Retrieving all '%s' members in LDAP." % args.group
        usernames = get_members(args.group)
    elif args.users:
        usernames = args.users.split(",")
    else: # if not args.users
        if not args.rebuild:
            print "Retrieving all 'atmo-user' members in LDAP."
            usernames = get_members('atmo-user')
        else:
            print "Rebuilding all existing users."
            usernames = get_usernames(provider)
    return run_create_accounts(acct_driver, provider, usernames,
                               args.rebuild, args.admin)


def run_create_accounts(acct_driver, provider, usernames, rebuild=False, admin=False):
    user_total = 0
    identity_total = 0
    for username in sorted(usernames):
        User.objects.get_or_create(username=username)
        new_identities = AccountCreationPluginManager.create_accounts(provider, username, force=rebuild)
        if new_identities:
            count = len(new_identities)
            print "%s new identities identity_total for %s." % (count, username)
            identity_total += count
            user_total += 1
        if admin:
            make_admin(username)
    print "%s Total identities identity_total for %s users" % (identity_total, user_total)


def make_admin(username):
    u = User.objects.get(username=username)
    u.is_superuser = True
    u.is_staff = True
    u.save()


if __name__ == "__main__":
    main()