#!/usr/bin/env python3
"""
This is a script which adds 'sn' and 'givenName' attributes to each user record.
It pulls first/last name information from UWLDAP.

GSSAPI is used for LDAP authentication, so make sure to run `kinit` first.
"""
import sys
import traceback

import ldap3

# modify as necessary
LDAP_URI = "ldaps://auth1.csclub.uwaterloo.ca"
LDAP_MEMBERS_BASE = "ou=People,dc=csclub,dc=uwaterloo,dc=ca"
UWLDAP_URI = "ldaps://uwldap.uwaterloo.ca"
UWLDAP_MEMBERS_BASE = "dc=uwaterloo,dc=ca"

csc_conn = ldap3.Connection(
    LDAP_URI, authentication=ldap3.SASL, sasl_mechanism=ldap3.KERBEROS,
    auto_bind=True, raise_exceptions=True)
uw_conn = ldap3.Connection(UWLDAP_URI, auto_bind=True, raise_exceptions=True)
csc_conn.search(
    LDAP_MEMBERS_BASE,
    '(&(objectClass=member)(!(|(sn=*)(givenName=*))))',
    attributes=['uid', 'cn'])
total_records_updated = 0
for csc_entry in csc_conn.entries:
    uid = csc_entry.uid.value
    cn = csc_entry.cn.value
    sn = None
    given_name = None
    try:
        uw_conn.search(
            f'uid={uid},{UWLDAP_MEMBERS_BASE}', '(objectClass=*)',
            attributes=['sn', 'givenName'], search_scope=ldap3.BASE)
        uw_entry = uw_conn.entries[0]
        sn = uw_entry.sn.value
        given_name = uw_entry.givenName.value
    except ldap3.core.exceptions.LDAPNoSuchObjectResult:
        pass
    if given_name is None or sn is None:
        print(f'WARNING: could not retrieve first and last names for {uid}; inferring from whitespace instead')
        words = cn.split()
        given_name, sn = words[0], words[-1]
    changes = {
        'givenName': [(ldap3.MODIFY_ADD, [given_name])],
        'sn': [(ldap3.MODIFY_ADD, [sn])],
    }
    csc_conn.modify(csc_entry.entry_dn, changes)
    print(f'Updated record for {uid}')
    total_records_updated += 1
print(f'Total records updated: {total_records_updated}')