def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--config-env', help="config environment", dest="config_env", default="production")
    parser.add_argument('--force', help="force bootstrap of account even if files exist", dest="force", action="store_true", default=False)
    parser.add_argument('--debug', help="set debug level (0-4)", dest="debug", nargs="?", const=0, type=int)
    parser.add_argument('--noop', help="only print actions, make no changes", dest="noop", action="store_true", default=False)
    parser.add_argument('--report', help="generate report of what will be done", dest="report", action="store_true", default=False)
    parser.add_argument('--report-space', help="report on space that can be removed", dest="report_space", action="store_true", default=False)
    parser.add_argument('--account', help="account to create", dest="account", default=None)
    parser.add_argument('--exclude-accounts', nargs="+", help="accounts to exclude", dest="exclude_accounts", default=[])
    args = parser.parse_args()
    options = vars(args)

    # Set values based on loaded config
    config = load_config()
    _auth_token = config[args.config_env].get("api_auth_token")
    _account_home_config = config[args.config_env].get("account_home")
    _cleanup_exclude = _account_home_config.get("cleanup_exclude", []) + args.exclude_accounts
    _host = config[args.config_env].get("host")
    _port = config[args.config_env].get("port")
    _https = config[args.config_env].get("https")
    _protocol = 'https' if _https else 'http'
    _url = "%s://%s:%s/" % (_protocol, _host, _port) if _port else "%s://%s/" % (_protocol, _host)
    _json_headers = {
        "Accept": "application/json",
        "Content-Type": "application/json",
        "Authorization": "Token token=%s" % _auth_token,
    }

    # Setup logging
    setup_logging(debug=args.debug, noop=args.noop)

    logger.debug4("OPTIONS: %s" % options)
    logger.debug4("CONFIG: %s" % config)

    # Get status ID
    status = actmgr_api.get_status(url=_url, headers=_json_headers, name='CLOSED')
    logger.debug1("STATUS: %s", status)
    status_id = status.get("id")

    # Get accounts and perform account cleanup steps
    if args.account:
        accounts = actmgr_api.get_accounts(url=_url, headers=_json_headers, params={"username": args.account, "status_id": status_id})
    else:
        accounts = actmgr_api.get_accounts(url=_url, headers=_json_headers, params={"status_id": status_id})
    logger.debug4("Number of accounts returned: %s", len(accounts))

    _report = []
    for account in accounts:
        logger.debug4("Account data: %s", json.dumps(account))
        _username = account["username"]
        if _username in _cleanup_exclude:
            logger.info("EXCLUDED: %s", _username)
            continue
        try:
            _shell = getpwnam(_username).pw_shell
        except KeyError:
            logger.warn("Unable to get shell for %s", _username)
            _shell = None
        if _shell != '/sbin/nologin':
            logger.warn("User %s shell %s != /sbin/nologin", _username, _shell)
            continue

        _account_home = AccountHome(username=_username, config=_account_home_config, options=options)
        _slurm_account = SlurmAccount(username=_username, options=options)

        if args.report:
            _account_home.check_path_owner(_account_home.home)
            _account_home.check_path_owner(_account_home.scratch)
            for _dir in _account_home.extra_directories:
                _account_home.check_path_owner(_dir)
            _data = {}
            _data["username"] = _username
            _data["HOME"] = _account_home.home_exists()
            _data["SCRATCH"] = _account_home.scratch_exists()
            _data["EXTRA"] = _account_home.extra_directories
            _data["SLURM"] = _slurm_account.exists()
            if args.report_space:
                _data["HOME_USED"] = get_space_used(host=_account_home_config["server"], path=_account_home.home)
                _data["SCRATCH_USED"] = get_space_used(path=_account_home.scratch)
                _data["EXTRA_USED"] = 0
                for _dir in _account_home.extra_directories:
                    _data["EXTRA_USED"] += get_space_used(path=_dir)
            _report.append(_data)
        else:
            _account_home.cleanup()
            _slurm_account.delete()
    if args.report:
        if args.report_space:
            table = prettytable.PrettyTable(["Username", "HOME", "HOME-USED", "SCRATCH", "SCRATCH-USED", "EXTRA", "EXTRA-USED", "SLURM"])
        else:
            table = prettytable.PrettyTable(["Username", "HOME", "SCRATCH", "EXTRA", "SLURM"])
        table.hrules = prettytable.FRAME
        _home_total = 0
        _home_used_total = 0
        _scratch_total = 0
        _scratch_used_total = 0
        _extra_total = 0
        _extra_used_total = 0
        _slurm_total = 0
        for r in sorted(_report, key=lambda k: k["username"]):
            _home = r["HOME"]
            _scratch = r["SCRATCH"]
            _extra = r["EXTRA"]
            _slurm = r["SLURM"]
            if _home:
                _home_total += 1
            if _scratch:
                _scratch_total += 1
            if _extra:
                _extra_total += len(_extra)
            if _slurm:
                _slurm_total += 1
            if args.report_space:
                _home_used = bytes2human(r["HOME_USED"])
                _home_used_total += r["HOME_USED"]
                _scratch_used = bytes2human(r["SCRATCH_USED"])
                _scratch_used_total += r["SCRATCH_USED"]
                _extra_used = bytes2human(r["EXTRA_USED"])
                _extra_used_total += r["EXTRA_USED"]
                table.add_row([r["username"], _home, _home_used, _scratch, _scratch_used, "\n".join(_extra), _extra_used, _slurm])
            else:
                table.add_row([r["username"], _home, _scratch, "\n".join(_extra), _slurm])
        if args.report_space:
            table.add_row(["", "", "", "", "", "", "", ""])
            table.add_row(["Total", _home_total, bytes2human(_home_used_total), _scratch_total, bytes2human(_scratch_used_total), _extra_total, bytes2human(_extra_used_total), _slurm_total])
        else:
            table.add_row(["", "", "", "", ""])
            table.add_row(["Total", _home_total, _scratch_total, _extra_total, _slurm_total])
        print table
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--username', help="User's username", required=True)
    parser.add_argument('--new-group', help="New group to assign", required=True)
    parser.add_argument('--old-group', help="Old group", required=True)
    parser.add_argument('--config-env', help="config environment", dest="config_env", default="production")
    parser.add_argument('--debug', help="set debug level (0-4)", dest="debug", nargs="?", const=0, type=int)
    args = parser.parse_args()
    options = vars(args)

    config = load_config()
    config_env = config[args.config_env]["ldap"]
    _account_home_config = config[args.config_env].get("account_home")
    _auth_token = config[args.config_env].get("api_auth_token")
    _host = config[args.config_env].get("host")
    _port = config[args.config_env].get("port")
    _https = config[args.config_env].get("https")
    _protocol = 'https' if _https else 'http'
    _url = "%s://%s:%s/" % (_protocol, _host, _port) if _port else "%s://%s/" % (_protocol, _host)
    _json_headers = {
        "Accept": "application/json",
        "Content-Type": "application/json",
        "Authorization": "Token token=%s" % _auth_token,
    }

    # Setup logging
    setup_logging(debug=args.debug, noop=False)

    logger.debug4("OPTIONS: %s" % options)
    logger.debug4("CONFIG: %s" % config_env)

    _ldap_url = config_env.get("url")
    _use_tls = config_env.get("tls")
    _bind_dn = config_env.get("bind_dn", None)
    _bind_pass = config_env.get("bind_pass", None)

    group_search_base = "ou=Groups,dc=brazos,dc=tamu,dc=edu"
    user_search_base = "ou=People,dc=brazos,dc=tamu,dc=edu"
    new_group_filter = "cn=%s" % args.new_group
    old_group_filter = "cn=%s" % args.old_group
    user_filter = "uid=%s" % args.username
    group_attribs = [
        "dn",
        "cn",
        "gidNumber",
        "uniqueMember",
        "slurmAccountName",
    ]
    user_attribs = [
        "dn",
        "uid",
        "gidNumber",
    ]
    scope = "one"

    local_ldap = LocalLdap(url=_ldap_url[0], use_tls=_use_tls, bind_dn=_bind_dn, bind_pass=_bind_pass, log_level=None)
    new_group_results = local_ldap.paged_search(base=group_search_base, sfilter=new_group_filter, attrlist=group_attribs, scope=scope)
    old_group_results = local_ldap.paged_search(base=group_search_base, sfilter=old_group_filter, attrlist=group_attribs, scope=scope)
    user_results = local_ldap.paged_search(base=user_search_base, sfilter=user_filter, attrlist=user_attribs, scope=scope)

    logger.debug("LDAP new group: %s", json.dumps(new_group_results))
    logger.debug("LDAP old group: %s", json.dumps(old_group_results))
    logger.debug("LDAP user: %s", json.dumps(user_results))

    if len(new_group_results) != 1 or len(old_group_results) != 1:
        logger.error("Incorrect number of LDAP group results returned")
        sys.exit(1)
    if len(user_results) != 1:
        logger.error("Incorrect number of LDAP user results returned")
        sys.exit(1)

    ldap_group_new = LdapGroup()
    ldap_group_new.setattrs(data=new_group_results[0], listvals=["uniqueMember"])
    ldap_group_old = LdapGroup()
    ldap_group_old.setattrs(data=old_group_results[0], listvals=["uniqueMember"])
    ldap_user = LdapUser()
    ldap_user.setattrs(data=user_results[0])

    # Not all LDAP groups have slurmAccountName attribute
    if not hasattr(ldap_group_new, 'slurmAccountName'):
        ldap_group_new.slurmAccountName = ldap_group_new.cn
    if not hasattr(ldap_group_old, 'slurmAccountName'):
        ldap_group_old.slurmAccountName = ldap_group_old.cn

    # Check certain things exist to avoid sending None to LDAP which could delete more than we want
    _ldap_group_new_valid = True
    _ldap_group_old_valid = True
    for a in ['dn', 'gidNumber', 'slurmAccountName', 'uniqueMember', 'cn']:
        if not hasattr(ldap_group_new, a):
            _ldap_group_new_valid = False
        elif getattr(ldap_group_new, a) is None:
            _ldap_group_new_valid = False
        if not hasattr(ldap_group_old, a):
            _ldap_group_old_valid = False
        elif getattr(ldap_group_old, a) is None:
            _ldap_group_old_valid = False
    if not _ldap_group_new_valid:
        logger.error("LDAP group %s does not have all necessary information", args.new_group)
    if not _ldap_group_old_valid:
        logger.error("LDAP group %s does not have all necessary information", args.old_group)
    if ldap_user.dn is None or ldap_user.uid is None or ldap_user.gidNumber is None:
        logger.error("LDAP user %s does not have all necessary information", args.username)

    ## Update account management database
    get_old_group_params = {
        "name": args.old_group,
    }
    old_group_data = actmgr_api.get_groups(_url, _json_headers, get_old_group_params)
    old_group = old_group_data[0]
    logger.debug("Old Group API data: %s", json.dumps(old_group))

    get_new_group_params = {
        "name": args.new_group,
    }
    new_group_data = actmgr_api.get_groups(_url, _json_headers, get_new_group_params)
    new_group = new_group_data[0]
    logger.debug("New Group API data: %s", json.dumps(new_group))

    get_account_params = {
        "username": args.username,
    }
    account_data = actmgr_api.get_accounts(_url, _json_headers, get_account_params)
    account = account_data[0]
    logger.debug("Account API data: %s", json.dumps(account))

    update_account_data = {
        "primary_group_id": new_group['id'],
    }
    # Handle cases where person is assigned old group in both primary and auxiliary groups
    group_ids = []
    for g in account["groups"]:
        if g["id"] == old_group["id"]:
            group_ids.append(new_group["id"])
        else:
            group_ids.append(g["id"])
    if group_ids:
        update_account_data['group_ids'] = group_ids

    account = actmgr_api.update_account(_url, _json_headers, account['id'], update_account_data)
    if not account or not account["account"]:
        logger.error("Failed to update account management data")
        sys.exit(1)
    logger.debug("Account updated API data: %s", json.dumps(account))
    account = account["account"]

    ## Update LDAP
    if ldap_user.gidNumber != ldap_group_new.gidNumber:
        logger.info("LDAP replace %s gidNumber=%s", ldap_user.dn, ldap_group_new.gidNumber)
        local_ldap.modify(ldap_user.dn, [(ldap.MOD_REPLACE, 'gidNumber', ldap_group_new.gidNumber)])
    else:
        logger.warn("Skipping LDAP update of user gidNumber - already updated")

    if ldap_user.dn not in ldap_group_new.uniqueMember:
        logger.info("LDAP add to %s uniqueMember=%s", ldap_group_new.dn, ldap_user.dn)
        local_ldap.modify(ldap_group_new.dn, [(ldap.MOD_ADD, "uniqueMember", ldap_user.dn)])
    else:
        logger.warn("Skipping LDAP update of group add uniqueMember - already updated")

    if ldap_user.dn in ldap_group_old.uniqueMember:
        logger.info("LDAP delete from %s uniqueMember=%s", ldap_group_old.dn, ldap_user.dn)
        local_ldap.modify(ldap_group_old.dn, [(ldap.MOD_DELETE, "uniqueMember", ldap_user.dn)])
    else:
        logger.warn("Skipping LDAP update of group delete uniqueMember - already updated")

    ## Update SLURM
    _slurm_account = account["primary_group"]["alias"]
    _slurm_accounts = [g["alias"] for g in account["groups"] if "alias" in g]
    if _slurm_account not in _slurm_accounts:
        _slurm_accounts.append(_slurm_account)

    if not _slurm_account or not _slurm_accounts:
        logger.error("SLURM accounts not correctly determined")
        sys.exit(1)

    sacctmgr_check_args = [
        "--parsable2", "--noheader", "show", "user",
        "name=%s" % args.username, "account=%s" % _slurm_account,
        "format=User,DefaultAccount,Account", "WithAssoc",
    ]
    logger.debug("Executing: sacctmgr %s", " ".join(sacctmgr_check_args))
    try:
        output = sacctmgr(sacctmgr_check_args)
    except ErrorReturnCode:
        logger.error("FAILED to check if SLURM account already exists.")
        sys.exit(1)
    expected_output = "%s|%s|%s" % (args.username, _slurm_account, _slurm_account)
    existing_slurm_accounts = output.split(os.linesep)
    if expected_output not in existing_slurm_accounts:
        sacctmgr_delete_args = ["-i", "delete", "user","where", "name=%s" % args.username, "account=%s" % ldap_group_old.slurmAccountName]
        logger.debug("Executing: sacctmgr %s", " ".join(sacctmgr_delete_args))
        try:
            output = sacctmgr(sacctmgr_delete_args)
        except ErrorReturnCode:
            logger.error("FAILED to delete user from SLURM.")
            sys.exit(1)

        sacctmgr_create_args = [
            "-i", "create", "user", args.username,
            "account=%s" % ",".join(_slurm_accounts),
            "defaultaccount=%s" % _slurm_account,
        ]
        logger.debug("Executing: sacctmgr %s", " ".join(sacctmgr_create_args))
        try:
            output = sacctmgr(sacctmgr_create_args)
        except ErrorReturnCode:
            logger.error("FAILED to retrieve all user names from SLURM.")
            sys.exit(1)
    else:
        logger.warn("Skipping SLURM account modifications - record already exists")

    ## Update permissions of $HOME and $SCRATCH
    home_path = os.path.join(_account_home_config.get("base_dir"), args.username)
    scratch_path = os.path.join(_account_home_config.get("scratch_base"), args.username)
    find_home_args = [
        home_path, "-group", args.old_group, "-exec", "chgrp", args.new_group, '{}', ';'
    ]
    logger.info("Changing group ownership of files under %s", home_path)
    logger.debug("Executing: find %s", " ".join(find_home_args))
    try:
        find(find_home_args)
    except ErrorReturnCode, e:
        logger.error("Failed to fix permissions of %s: %s", home_path, e.stderr)
        sys.exit(1)
Exemplo n.º 3
0
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--username', help="User's username", required=True)
    parser.add_argument('--new-group',
                        help="New group to assign",
                        required=True)
    parser.add_argument('--old-group', help="Old group", required=True)
    parser.add_argument('--config-env',
                        help="config environment",
                        dest="config_env",
                        default="production")
    parser.add_argument('--debug',
                        help="set debug level (0-4)",
                        dest="debug",
                        nargs="?",
                        const=0,
                        type=int)
    args = parser.parse_args()
    options = vars(args)

    config = load_config()
    config_env = config[args.config_env]["ldap"]
    _account_home_config = config[args.config_env].get("account_home")
    _auth_token = config[args.config_env].get("api_auth_token")
    _host = config[args.config_env].get("host")
    _port = config[args.config_env].get("port")
    _https = config[args.config_env].get("https")
    _protocol = 'https' if _https else 'http'
    _url = "%s://%s:%s/" % (
        _protocol, _host, _port) if _port else "%s://%s/" % (_protocol, _host)
    _json_headers = {
        "Accept": "application/json",
        "Content-Type": "application/json",
        "Authorization": "Token token=%s" % _auth_token,
    }

    # Setup logging
    setup_logging(debug=args.debug, noop=False)

    logger.debug4("OPTIONS: %s" % options)
    logger.debug4("CONFIG: %s" % config_env)

    _ldap_url = config_env.get("url")
    _use_tls = config_env.get("tls")
    _bind_dn = config_env.get("bind_dn", None)
    _bind_pass = config_env.get("bind_pass", None)

    group_search_base = "ou=Groups,dc=brazos,dc=tamu,dc=edu"
    user_search_base = "ou=People,dc=brazos,dc=tamu,dc=edu"
    new_group_filter = "cn=%s" % args.new_group
    old_group_filter = "cn=%s" % args.old_group
    user_filter = "uid=%s" % args.username
    group_attribs = [
        "dn",
        "cn",
        "gidNumber",
        "uniqueMember",
        "slurmAccountName",
    ]
    user_attribs = [
        "dn",
        "uid",
        "gidNumber",
    ]
    scope = "one"

    local_ldap = LocalLdap(url=_ldap_url[0],
                           use_tls=_use_tls,
                           bind_dn=_bind_dn,
                           bind_pass=_bind_pass,
                           log_level=None)
    new_group_results = local_ldap.paged_search(base=group_search_base,
                                                sfilter=new_group_filter,
                                                attrlist=group_attribs,
                                                scope=scope)
    old_group_results = local_ldap.paged_search(base=group_search_base,
                                                sfilter=old_group_filter,
                                                attrlist=group_attribs,
                                                scope=scope)
    user_results = local_ldap.paged_search(base=user_search_base,
                                           sfilter=user_filter,
                                           attrlist=user_attribs,
                                           scope=scope)

    logger.debug("LDAP new group: %s", json.dumps(new_group_results))
    logger.debug("LDAP old group: %s", json.dumps(old_group_results))
    logger.debug("LDAP user: %s", json.dumps(user_results))

    if len(new_group_results) != 1 or len(old_group_results) != 1:
        logger.error("Incorrect number of LDAP group results returned")
        sys.exit(1)
    if len(user_results) != 1:
        logger.error("Incorrect number of LDAP user results returned")
        sys.exit(1)

    ldap_group_new = LdapGroup()
    ldap_group_new.setattrs(data=new_group_results[0],
                            listvals=["uniqueMember"])
    ldap_group_old = LdapGroup()
    ldap_group_old.setattrs(data=old_group_results[0],
                            listvals=["uniqueMember"])
    ldap_user = LdapUser()
    ldap_user.setattrs(data=user_results[0])

    # Not all LDAP groups have slurmAccountName attribute
    if not hasattr(ldap_group_new, 'slurmAccountName'):
        ldap_group_new.slurmAccountName = ldap_group_new.cn
    if not hasattr(ldap_group_old, 'slurmAccountName'):
        ldap_group_old.slurmAccountName = ldap_group_old.cn

    # Check certain things exist to avoid sending None to LDAP which could delete more than we want
    _ldap_group_new_valid = True
    _ldap_group_old_valid = True
    for a in ['dn', 'gidNumber', 'slurmAccountName', 'uniqueMember', 'cn']:
        if not hasattr(ldap_group_new, a):
            _ldap_group_new_valid = False
        elif getattr(ldap_group_new, a) is None:
            _ldap_group_new_valid = False
        if not hasattr(ldap_group_old, a):
            _ldap_group_old_valid = False
        elif getattr(ldap_group_old, a) is None:
            _ldap_group_old_valid = False
    if not _ldap_group_new_valid:
        logger.error("LDAP group %s does not have all necessary information",
                     args.new_group)
    if not _ldap_group_old_valid:
        logger.error("LDAP group %s does not have all necessary information",
                     args.old_group)
    if ldap_user.dn is None or ldap_user.uid is None or ldap_user.gidNumber is None:
        logger.error("LDAP user %s does not have all necessary information",
                     args.username)

    ## Update account management database
    get_old_group_params = {
        "name": args.old_group,
    }
    old_group_data = actmgr_api.get_groups(_url, _json_headers,
                                           get_old_group_params)
    old_group = old_group_data[0]
    logger.debug("Old Group API data: %s", json.dumps(old_group))

    get_new_group_params = {
        "name": args.new_group,
    }
    new_group_data = actmgr_api.get_groups(_url, _json_headers,
                                           get_new_group_params)
    new_group = new_group_data[0]
    logger.debug("New Group API data: %s", json.dumps(new_group))

    get_account_params = {
        "username": args.username,
    }
    account_data = actmgr_api.get_accounts(_url, _json_headers,
                                           get_account_params)
    account = account_data[0]
    logger.debug("Account API data: %s", json.dumps(account))

    update_account_data = {
        "primary_group_id": new_group['id'],
    }
    # Handle cases where person is assigned old group in both primary and auxiliary groups
    group_ids = []
    for g in account["groups"]:
        if g["id"] == old_group["id"]:
            group_ids.append(new_group["id"])
        else:
            group_ids.append(g["id"])
    if group_ids:
        update_account_data['group_ids'] = group_ids

    account = actmgr_api.update_account(_url, _json_headers, account['id'],
                                        update_account_data)
    if not account or not account["account"]:
        logger.error("Failed to update account management data")
        sys.exit(1)
    logger.debug("Account updated API data: %s", json.dumps(account))
    account = account["account"]

    ## Update LDAP
    if ldap_user.gidNumber != ldap_group_new.gidNumber:
        logger.info("LDAP replace %s gidNumber=%s", ldap_user.dn,
                    ldap_group_new.gidNumber)
        local_ldap.modify(
            ldap_user.dn,
            [(ldap.MOD_REPLACE, 'gidNumber', ldap_group_new.gidNumber)])
    else:
        logger.warn("Skipping LDAP update of user gidNumber - already updated")

    if ldap_user.dn not in ldap_group_new.uniqueMember:
        logger.info("LDAP add to %s uniqueMember=%s", ldap_group_new.dn,
                    ldap_user.dn)
        local_ldap.modify(ldap_group_new.dn,
                          [(ldap.MOD_ADD, "uniqueMember", ldap_user.dn)])
    else:
        logger.warn(
            "Skipping LDAP update of group add uniqueMember - already updated")

    if ldap_user.dn in ldap_group_old.uniqueMember:
        logger.info("LDAP delete from %s uniqueMember=%s", ldap_group_old.dn,
                    ldap_user.dn)
        local_ldap.modify(ldap_group_old.dn,
                          [(ldap.MOD_DELETE, "uniqueMember", ldap_user.dn)])
    else:
        logger.warn(
            "Skipping LDAP update of group delete uniqueMember - already updated"
        )

    ## Update SLURM
    _slurm_account = account["primary_group"]["alias"]
    _slurm_accounts = [g["alias"] for g in account["groups"] if "alias" in g]
    if _slurm_account not in _slurm_accounts:
        _slurm_accounts.append(_slurm_account)

    if not _slurm_account or not _slurm_accounts:
        logger.error("SLURM accounts not correctly determined")
        sys.exit(1)

    sacctmgr_check_args = [
        "--parsable2",
        "--noheader",
        "show",
        "user",
        "name=%s" % args.username,
        "account=%s" % _slurm_account,
        "format=User,DefaultAccount,Account",
        "WithAssoc",
    ]
    logger.debug("Executing: sacctmgr %s", " ".join(sacctmgr_check_args))
    try:
        output = sacctmgr(sacctmgr_check_args)
    except ErrorReturnCode:
        logger.error("FAILED to check if SLURM account already exists.")
        sys.exit(1)
    expected_output = "%s|%s|%s" % (args.username, _slurm_account,
                                    _slurm_account)
    existing_slurm_accounts = output.split(os.linesep)
    if expected_output not in existing_slurm_accounts:
        sacctmgr_delete_args = [
            "-i", "delete", "user", "where",
            "name=%s" % args.username,
            "account=%s" % ldap_group_old.slurmAccountName
        ]
        logger.debug("Executing: sacctmgr %s", " ".join(sacctmgr_delete_args))
        try:
            output = sacctmgr(sacctmgr_delete_args)
        except ErrorReturnCode:
            logger.error("FAILED to delete user from SLURM.")
            sys.exit(1)

        sacctmgr_create_args = [
            "-i",
            "create",
            "user",
            args.username,
            "account=%s" % ",".join(_slurm_accounts),
            "defaultaccount=%s" % _slurm_account,
        ]
        logger.debug("Executing: sacctmgr %s", " ".join(sacctmgr_create_args))
        try:
            output = sacctmgr(sacctmgr_create_args)
        except ErrorReturnCode:
            logger.error("FAILED to retrieve all user names from SLURM.")
            sys.exit(1)
    else:
        logger.warn(
            "Skipping SLURM account modifications - record already exists")

    ## Update permissions of $HOME and $SCRATCH
    home_path = os.path.join(_account_home_config.get("base_dir"),
                             args.username)
    scratch_path = os.path.join(_account_home_config.get("scratch_base"),
                                args.username)
    find_home_args = [
        home_path, "-group", args.old_group, "-exec", "chgrp", args.new_group,
        '{}', ';'
    ]
    logger.info("Changing group ownership of files under %s", home_path)
    logger.debug("Executing: find %s", " ".join(find_home_args))
    try:
        find(find_home_args)
    except ErrorReturnCode, e:
        logger.error("Failed to fix permissions of %s: %s", home_path,
                     e.stderr)
        sys.exit(1)
Exemplo n.º 4
0
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--config-env',
                        help="config environment",
                        dest="config_env",
                        default="production")
    parser.add_argument('--force',
                        help="force bootstrap of account even if files exist",
                        dest="force",
                        action="store_true",
                        default=False)
    parser.add_argument('--debug',
                        help="set debug level (0-4)",
                        dest="debug",
                        nargs="?",
                        const=0,
                        type=int)
    parser.add_argument('--noop',
                        help="only print actions, make no changes",
                        dest="noop",
                        action="store_true",
                        default=False)
    parser.add_argument('--report',
                        help="generate report of what will be done",
                        dest="report",
                        action="store_true",
                        default=False)
    parser.add_argument('--report-space',
                        help="report on space that can be removed",
                        dest="report_space",
                        action="store_true",
                        default=False)
    parser.add_argument('--account',
                        help="account to create",
                        dest="account",
                        default=None)
    parser.add_argument('--exclude-accounts',
                        nargs="+",
                        help="accounts to exclude",
                        dest="exclude_accounts",
                        default=[])
    args = parser.parse_args()
    options = vars(args)

    # Set values based on loaded config
    config = load_config()
    _auth_token = config[args.config_env].get("api_auth_token")
    _account_home_config = config[args.config_env].get("account_home")
    _cleanup_exclude = _account_home_config.get("cleanup_exclude",
                                                []) + args.exclude_accounts
    _host = config[args.config_env].get("host")
    _port = config[args.config_env].get("port")
    _https = config[args.config_env].get("https")
    _protocol = 'https' if _https else 'http'
    _url = "%s://%s:%s/" % (
        _protocol, _host, _port) if _port else "%s://%s/" % (_protocol, _host)
    _json_headers = {
        "Accept": "application/json",
        "Content-Type": "application/json",
        "Authorization": "Token token=%s" % _auth_token,
    }

    # Setup logging
    setup_logging(debug=args.debug, noop=args.noop)

    logger.debug4("OPTIONS: %s" % options)
    logger.debug4("CONFIG: %s" % config)

    # Get status ID
    status = actmgr_api.get_status(url=_url,
                                   headers=_json_headers,
                                   name='CLOSED')
    logger.debug1("STATUS: %s", status)
    status_id = status.get("id")

    # Get accounts and perform account cleanup steps
    if args.account:
        accounts = actmgr_api.get_accounts(url=_url,
                                           headers=_json_headers,
                                           params={
                                               "username": args.account,
                                               "status_id": status_id
                                           })
    else:
        accounts = actmgr_api.get_accounts(url=_url,
                                           headers=_json_headers,
                                           params={"status_id": status_id})
    logger.debug4("Number of accounts returned: %s", len(accounts))

    _report = []
    for account in accounts:
        logger.debug4("Account data: %s", json.dumps(account))
        _username = account["username"]
        if _username in _cleanup_exclude:
            logger.info("EXCLUDED: %s", _username)
            continue
        try:
            _shell = getpwnam(_username).pw_shell
        except KeyError:
            logger.warn("Unable to get shell for %s", _username)
            _shell = None
        if _shell != '/sbin/nologin':
            logger.warn("User %s shell %s != /sbin/nologin", _username, _shell)
            continue

        _account_home = AccountHome(username=_username,
                                    config=_account_home_config,
                                    options=options)
        _slurm_account = SlurmAccount(username=_username, options=options)

        if args.report:
            _account_home.check_path_owner(_account_home.home)
            _account_home.check_path_owner(_account_home.scratch)
            for _dir in _account_home.extra_directories:
                _account_home.check_path_owner(_dir)
            _data = {}
            _data["username"] = _username
            _data["HOME"] = _account_home.home_exists()
            _data["SCRATCH"] = _account_home.scratch_exists()
            _data["EXTRA"] = _account_home.extra_directories
            _data["SLURM"] = _slurm_account.exists()
            if args.report_space:
                _data["HOME_USED"] = get_space_used(
                    host=_account_home_config["server"],
                    path=_account_home.home)
                _data["SCRATCH_USED"] = get_space_used(
                    path=_account_home.scratch)
                _data["EXTRA_USED"] = 0
                for _dir in _account_home.extra_directories:
                    _data["EXTRA_USED"] += get_space_used(path=_dir)
            _report.append(_data)
        else:
            _account_home.cleanup()
            _slurm_account.delete()
    if args.report:
        if args.report_space:
            table = prettytable.PrettyTable([
                "Username", "HOME", "HOME-USED", "SCRATCH", "SCRATCH-USED",
                "EXTRA", "EXTRA-USED", "SLURM"
            ])
        else:
            table = prettytable.PrettyTable(
                ["Username", "HOME", "SCRATCH", "EXTRA", "SLURM"])
        table.hrules = prettytable.FRAME
        _home_total = 0
        _home_used_total = 0
        _scratch_total = 0
        _scratch_used_total = 0
        _extra_total = 0
        _extra_used_total = 0
        _slurm_total = 0
        for r in sorted(_report, key=lambda k: k["username"]):
            _home = r["HOME"]
            _scratch = r["SCRATCH"]
            _extra = r["EXTRA"]
            _slurm = r["SLURM"]
            if _home:
                _home_total += 1
            if _scratch:
                _scratch_total += 1
            if _extra:
                _extra_total += len(_extra)
            if _slurm:
                _slurm_total += 1
            if args.report_space:
                _home_used = bytes2human(r["HOME_USED"])
                _home_used_total += r["HOME_USED"]
                _scratch_used = bytes2human(r["SCRATCH_USED"])
                _scratch_used_total += r["SCRATCH_USED"]
                _extra_used = bytes2human(r["EXTRA_USED"])
                _extra_used_total += r["EXTRA_USED"]
                table.add_row([
                    r["username"], _home, _home_used, _scratch, _scratch_used,
                    "\n".join(_extra), _extra_used, _slurm
                ])
            else:
                table.add_row([
                    r["username"], _home, _scratch, "\n".join(_extra), _slurm
                ])
        if args.report_space:
            table.add_row(["", "", "", "", "", "", "", ""])
            table.add_row([
                "Total", _home_total,
                bytes2human(_home_used_total), _scratch_total,
                bytes2human(_scratch_used_total), _extra_total,
                bytes2human(_extra_used_total), _slurm_total
            ])
        else:
            table.add_row(["", "", "", "", ""])
            table.add_row([
                "Total", _home_total, _scratch_total, _extra_total,
                _slurm_total
            ])
        print table