From 7af93a46e799b6e33184c31da6520a31bfe0fb47 Mon Sep 17 00:00:00 2001 From: Chun Ly Date: Tue, 14 Jul 2020 10:51:14 -0700 Subject: [PATCH] Factor out static function of ManualOverride #31 Update user_update script accordingly --- requiam/manual_override.py | 189 +++++++++++++++++++++---------------- requiam/user_update | 12 +-- 2 files changed, 113 insertions(+), 88 deletions(-) diff --git a/requiam/manual_override.py b/requiam/manual_override.py index a26a36a1..ab3ea3f1 100644 --- a/requiam/manual_override.py +++ b/requiam/manual_override.py @@ -42,44 +42,8 @@ def __init__(self, portal_file, quota_file): self.quota_file = quota_file # Read in as pandas DataFrame - self.portal_df = self.read_manual_file(self.portal_file) - self.quota_df = self.read_manual_file(self.quota_file) - - @staticmethod - def read_manual_file(input_file): - """Read in manual override file as pandas DataFrame""" - - try: - df = pd.read_csv(input_file, comment='#') - - return df - except FileNotFoundError: - print(f"File not found! : {input_file}") - - @staticmethod - def update_entries(ldap_set, netid, uaid, action): - """Add/remove entries from a set""" - - if action not in ['remove', 'add']: - raise ValueError("Incorrect [action] input") - - new_ldap_set = set(ldap_set) - - if action == 'remove': - if isinstance(netid, list): - print(f"Removing : {list(netid)}") - if isinstance(netid, str): - print(f"Removing : {netid}") - new_ldap_set = ldap_set - uaid - - if action == 'add': - if isinstance(netid, list): - print(f"Adding : {list(netid)}") - if isinstance(netid, str): - print(f"Adding : {netid}") - new_ldap_set = set.union(ldap_set, uaid) - - return new_ldap_set + self.portal_df = read_manual_file(self.portal_file) + self.quota_df = read_manual_file(self.quota_file) def identify_changes(self, ldap_set, group, group_type): """Identify changes to call update_entries accordingly""" @@ -100,66 +64,127 @@ def identify_changes(self, ldap_set, group, group_type): add_ldap_set = set(ldap_set) if len(add_df) > 0: # Add to ldap_set - add_ldap_set = self.update_entries(ldap_set, add_df['netid'], - add_df['uaid'], 'add') + add_ldap_set = update_entries(ldap_set, add_df['netid'], + add_df['uaid'], 'add') # Identify those that needs to be excluded in [group] outside_df = manual_df.loc[manual_df[group_type] != group] if len(outside_df) > 0: - new_ldap_set = self.update_entries(add_ldap_set, outside_df['netid'], - outside_df['uaid'], 'remove') + new_ldap_set = update_entries(add_ldap_set, outside_df['netid'], + outside_df['uaid'], 'remove') else: new_ldap_set = add_ldap_set return new_ldap_set - @staticmethod - def get_current_groups(uid, ldap_dict, log): - """Retrieve current Figshare ismemberof association""" - mo_ldc = LDAPConnection(**ldap_dict) - mo_ldc.ldap_attribs = ['ismemberof'] - user_query = f'(uid={uid})' +def read_manual_file(input_file): + """ + Purpose: + Read in manual override file as pandas DataFrame + + :param input_file: full filename + :return df: pandas DataFrame + """ + + try: + df = pd.read_csv(input_file, comment='#') + + return df + except FileNotFoundError: + print(f"File not found! : {input_file}") + + +def update_entries(ldap_set, netid, uaid, action): + """ + Purpose: + Add/remove entries from a set + + :param ldap_set: set of uaid values + :param netid: User netid + :param uaid: User uaid + :param action: str + Action to perform. Either 'remove' or 'add' + :return new_ldap_set: Updated set of uaid values + """ + + if action not in ['remove', 'add']: + raise ValueError("Incorrect [action] input") + + new_ldap_set = set(ldap_set) + + if action == 'remove': + if isinstance(netid, list): + print(f"Removing : {list(netid)}") + if isinstance(netid, str): + print(f"Removing : {netid}") + new_ldap_set = ldap_set - uaid + + if action == 'add': + if isinstance(netid, list): + print(f"Adding : {list(netid)}") + if isinstance(netid, str): + print(f"Adding : {netid}") + new_ldap_set = set.union(ldap_set, uaid) + + return new_ldap_set + - mo_ldc.ldc.search(mo_ldc.ldap_search_dn, user_query, attributes=mo_ldc.ldap_attribs) +def get_current_groups(uid, ldap_dict, log): + """ + Purpose: + Retrieve current Figshare ismemberof association + + :param uid: str containing User NetID + :param ldap_dict: dict containing ldap settings + :param log: LogClass object for logging + :return figshare_dict: dict containing current Figshare portal and quota + """ + + mo_ldc = LDAPConnection(**ldap_dict) + mo_ldc.ldap_attribs = ['ismemberof'] - membership = mo_ldc.ldc.entries[0].ismemberof.value + user_query = f'(uid={uid})' - figshare_dict = dict() + mo_ldc.ldc.search(mo_ldc.ldap_search_dn, user_query, attributes=mo_ldc.ldap_attribs) - if isinstance(membership, type(None)): - log.warning("No ismembersof attributes") + membership = mo_ldc.ldc.entries[0].ismemberof.value - figshare_dict['portal'] = '' - figshare_dict['quota'] = '' - return figshare_dict + figshare_dict = dict() - # Extract portal - portal_stem = figshare_stem('portal') - portal = [s for s in membership if ((portal_stem in s) and ('grouper' not in s))] - if len(portal) == 0: - log.info("No Grouper group found!") - figshare_dict['portal'] = '' # Initialize to use later + if isinstance(membership, type(None)): + log.warning("No ismembersof attributes") + + figshare_dict['portal'] = '' + figshare_dict['quota'] = '' + return figshare_dict + + # Extract portal + portal_stem = figshare_stem('portal') + portal = [s for s in membership if ((portal_stem in s) and ('grouper' not in s))] + if len(portal) == 0: + log.info("No Grouper group found!") + figshare_dict['portal'] = '' # Initialize to use later + else: + if len(portal) != 1: + log.warning("ERROR! Multiple Grouper portal found") + raise ValueError else: - if len(portal) != 1: - log.warning("ERROR! Multiple Grouper portal found") - raise ValueError - else: - figshare_dict['portal'] = portal[0].replace(portal_stem + ':', '') - log.info(f"Current portal is : {figshare_dict['portal']}") - - # Extract quota - quota_stem = figshare_stem('quota') - quota = [s for s in membership if ((quota_stem in s) and ('grouper' not in s))] - if len(quota) == 0: - log.info("No Grouper group found!") - figshare_dict['quota'] = '' # Initialize to use later + figshare_dict['portal'] = portal[0].replace(portal_stem + ':', '') + log.info(f"Current portal is : {figshare_dict['portal']}") + + # Extract quota + quota_stem = figshare_stem('quota') + quota = [s for s in membership if ((quota_stem in s) and ('grouper' not in s))] + if len(quota) == 0: + log.info("No Grouper group found!") + figshare_dict['quota'] = '' # Initialize to use later + else: + if len(quota) != 1: + log.warning("ERROR! Multiple Grouper quota found") + raise ValueError else: - if len(quota) != 1: - log.warning("ERROR! Multiple Grouper quota found") - raise ValueError - else: - figshare_dict['quota'] = quota[0].replace(quota_stem + ':', '') - log.info(f"Current quota is : {figshare_dict['quota']} bytes") + figshare_dict['quota'] = quota[0].replace(quota_stem + ':', '') + log.info(f"Current quota is : {figshare_dict['quota']} bytes") - return figshare_dict + return figshare_dict diff --git a/requiam/user_update b/requiam/user_update index df5e272a..43b7c33f 100644 --- a/requiam/user_update +++ b/requiam/user_update @@ -14,7 +14,7 @@ from requiam import delta from requiam import quota from requiam.logger import LogClass from requiam import TimerClass -from requiam import manual_override +from requiam.manual_override import ManualOverride, update_entries, get_current_groups today = date.today() @@ -111,7 +111,7 @@ if __name__ == '__main__': sync_max=int(vargs['sync_max'])) # Manual override class - mo = manual_override.ManualOverride(vargs['portal_file'], vargs['quota_file']) + mo = ManualOverride(vargs['portal_file'], vargs['quota_file']) # Initiate LDAP connection ldc = ldap_query.LDAPConnection(**ldap_dict) @@ -122,7 +122,7 @@ if __name__ == '__main__': log.info(f" uaid for {vargs['netid']} : {list(user_uaid)[0]}") # Retrieve ismemberof figshare information - current_dict = mo.get_current_groups(vargs['netid'], ldap_dict, log) + current_dict = get_current_groups(vargs['netid'], ldap_dict, log) # Portal update if vargs['portal'] != '(unset)': @@ -136,7 +136,7 @@ if __name__ == '__main__': gq = grouper_query.GrouperQuery(**grouper_dict, grouper_group=portal_query) portal_set = gq.members - portal_set = mo.update_entries(portal_set, vargs['netid'], user_uaid, 'remove') + portal_set = update_entries(portal_set, vargs['netid'], user_uaid, 'remove') log.info(" Grouper size {}".format(len(gq.members))) @@ -161,7 +161,7 @@ if __name__ == '__main__': gq = grouper_query.GrouperQuery(**grouper_dict, grouper_group=portal_query) portal_set = gq.members - portal_set = mo.update_entries(portal_set, vargs['netid'], user_uaid, 'add') + portal_set = update_entries(portal_set, vargs['netid'], user_uaid, 'add') # Grouper query log.info(" Grouper size {}".format(len(gq.members))) @@ -188,7 +188,7 @@ if __name__ == '__main__': gq = grouper_query.GrouperQuery(**grouper_dict, grouper_group=quota_query) quota_set = gq.members - quota_set = mo.update_entries(quota_set, vargs['netid'], user_uaid, 'add') + quota_set = update_entries(quota_set, vargs['netid'], user_uaid, 'add') # Grouper query log.info(" Grouper size {}".format(len(gq.members)))