pyceo/ceod/model/LDAPService.py

import copy
import grp
import pwd
from typing import Union, List

import ldap
import ldap.modlist
from zope import component
from zope.interface import implementer

from ceo_common.errors import UserNotFoundError, GroupNotFoundError, \
    UserAlreadyExistsError, GroupAlreadyExistsError
from ceo_common.interfaces import ILDAPService, IKerberosService, IConfig, \
    IUser, IGroup, IUWLDAPService
from .User import User
from .Group import Group
from .SudoRole import SudoRole
from .utils import dn_to_uid, bytes_to_strings


@implementer(ILDAPService)
class LDAPService:
    def __init__(self):
        cfg = component.getUtility(IConfig)
        self.ldap_admin_principal = cfg.get('ldap_admin_principal')
        self.ldap_server_url = cfg.get('ldap_server_url')
        self.ldap_users_base = cfg.get('ldap_users_base')
        self.ldap_groups_base = cfg.get('ldap_groups_base')
        self.member_min_id = cfg.get('members_min_id')
        self.member_max_id = cfg.get('members_max_id')
        self.club_min_id = cfg.get('clubs_min_id')
        self.club_max_id = cfg.get('clubs_max_id')

    def _get_ldap_conn(self, gssapi_bind: bool = True) -> ldap.ldapobject.LDAPObject:
        # TODO: cache the connection
        conn = ldap.initialize(self.ldap_server_url)
        if gssapi_bind:
            self._gssapi_bind(conn)
        return conn

    def _gssapi_bind(self, conn: ldap.ldapobject.LDAPObject):
        krb_srv = component.getUtility(IKerberosService)
        for i in range(2):
            try:
                conn.sasl_gssapi_bind_s()
                return
            except ldap.LOCAL_ERROR as err:
                if 'Ticket expired' in err.args[0]['info']:
                    krb_srv.kinit()
                    continue
                raise err
        raise Exception('could not perform GSSAPI bind')

    def get_user(self, username: str) -> IUser:
        conn = self._get_ldap_conn(False)
        base = self.uid_to_dn(username)
        try:
            _, result = conn.search_s(base, ldap.SCOPE_BASE)[0]
            return User.deserialize_from_ldap(result)
        except ldap.NO_SUCH_OBJECT:
            raise UserNotFoundError()

    def get_group(self, cn: str) -> IGroup:
        conn = self._get_ldap_conn(False)
        base = self.group_cn_to_dn(cn)
        try:
            _, result = conn.search_s(base, ldap.SCOPE_BASE)[0]
            return Group.deserialize_from_ldap(result)
        except ldap.NO_SUCH_OBJECT:
            raise GroupNotFoundError()

    def uid_to_dn(self, uid: str):
        return f'uid={uid},{self.ldap_users_base}'

    def group_cn_to_dn(self, cn: str):
        return f'cn={cn},{self.ldap_groups_base}'

    def _get_next_uid(self, conn: ldap.ldapobject.LDAPObject, min_id: int, max_id: int) -> int:
        """Gets the next available UID number between min_id and max_id, inclusive."""
        def uid_exists(uid: int) -> bool:
            try:
                pwd.getpwuid(uid)
                return True
            except KeyError:
                return False

        def gid_exists(gid: int) -> bool:
            try:
                grp.getgrgid(gid)
                return True
            except KeyError:
                return False

        def ldap_uid_or_gid_exists(uid: int) -> bool:
            results = conn.search_s(
                self.ldap_users_base, ldap.SCOPE_ONELEVEL,
                f'(|(uidNumber={uid})(gidNumber={uid}))')
            return len(results) > 0

        # TODO: replace this with binary search
        for uid in range(min_id, max_id + 1):
            if uid_exists(uid) or gid_exists(uid) or ldap_uid_or_gid_exists(uid):
                continue
            return uid
        raise Exception('no UIDs remaining')

    def add_sudo_role(self, uid: str):
        conn = self._get_ldap_conn()
        sudo_role = SudoRole(uid)
        modlist = ldap.modlist.addModlist(sudo_role.serialize_for_ldap())
        conn.add_s(sudo_role.dn, modlist)

    def remove_sudo_role(self, uid: str):
        conn = self._get_ldap_conn()
        sudo_role = SudoRole(uid)
        conn.delete_s(sudo_role.dn)

    def add_user(self, user: IUser) -> IUser:
        if user.is_club():
            min_id, max_id = self.club_min_id, self.club_max_id
        else:
            min_id, max_id = self.member_min_id, self.member_max_id
        conn = self._get_ldap_conn()
        uid_number = self._get_next_uid(conn, min_id, max_id)
        new_user = copy.copy(user)
        new_user.uid_number = uid_number
        new_user.gid_number = uid_number

        modlist = ldap.modlist.addModlist(new_user.serialize_for_ldap())
        try:
            conn.add_s(new_user.dn, modlist)
        except ldap.ALREADY_EXISTS:
            raise UserAlreadyExistsError()
        return new_user

    def modify_user(self, old_user: IUser, new_user: IUser):
        conn = self._get_ldap_conn()
        modlist = ldap.modlist.modifyModlist(
            old_user.serialize_for_ldap(),
            new_user.serialize_for_ldap(),
        )
        conn.modify_s(old_user.dn, modlist)

    def remove_user(self, user: IUser):
        conn = self._get_ldap_conn()
        conn.delete_s(user.dn)

    def add_group(self, group: IGroup) -> IGroup:
        conn = self._get_ldap_conn()
        # make sure that the caller initialized the GID number
        assert group.gid_number
        modlist = ldap.modlist.addModlist(group.serialize_for_ldap())
        try:
            conn.add_s(group.dn, modlist)
        except ldap.ALREADY_EXISTS:
            raise GroupAlreadyExistsError()
        return group

    def modify_group(self, old_group: IGroup, new_group: IGroup):
        conn = self._get_ldap_conn()
        modlist = ldap.modlist.modifyModlist(
            old_group.serialize_for_ldap(),
            new_group.serialize_for_ldap(),
        )
        conn.modify_s(old_group.dn, modlist)

    def remove_group(self, group: IGroup):
        conn = self._get_ldap_conn()
        conn.delete_s(group.dn)

    def update_programs(
        self,
        dry_run: bool = False,
        members: Union[List[str], None] = None,
        uwldap_batch_size: int = 100,
    ):
        if members:
            filter_str = '(|' + ''.join([
                f'(uid={uid})' for uid in members
            ]) + ')'
        else:
            filter_str = None
        conn = self._get_ldap_conn()
        raw_csc_records = conn.search_s(
            self.ldap_users_base, ldap.SCOPE_SUBTREE, filter_str,
            attrlist=['program'])
        uids = [
            dn_to_uid(dn) for dn, _ in raw_csc_records
        ]
        csc_programs = [
            bytes_to_strings(data).get('program', [None])[0]
            for _, data in raw_csc_records
        ]

        uwldap_srv = component.getUtility(IUWLDAPService)
        uw_programs = []
        # send queries in small batches
        for i in range(0, len(csc_programs), uwldap_batch_size):
            batch_uids = uids[i:i + uwldap_batch_size]
            batch_uw_programs = uwldap_srv.get_programs_for_users(batch_uids)
            uw_programs.extend(batch_uw_programs)
        users_to_change = [
            (uids[i], csc_programs[i], uw_programs[i])
            for i in range(len(uids))
            if csc_programs[i] != uw_programs[i] and (
                uw_programs[i] not in (None, 'expired', 'orphaned')
            )
        ]
        if dry_run:
            return users_to_change

        for uid, old_program, new_program in users_to_change:
            old_entry = {'program': [old_program.encode()]}
            new_entry = {'program': [new_program.encode()]}
            modlist = ldap.modlist.modifyModlist(old_entry, new_entry)
            conn.modify_s(self.uid_to_dn(uid), modlist)

        return users_to_change