<commit_before>from gevent import monkey, sleep, spawn
monkey.patch_all()  # NOQA

import sys
import time
import yaml
import logging
import ldap

from oncall import metrics
from ldap.controls import SimplePagedResultsControl
from datetime import datetime
from pytz import timezone
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from sqlalchemy.exc import SQLAlchemyError, IntegrityError
from phonenumbers import format_number, parse, PhoneNumberFormat
from phonenumbers.phonenumberutil import NumberParseException


logger = logging.getLogger()
formatter = logging.Formatter('%(asctime)s %(levelname)s %(name)s %(message)s')
ch = logging.StreamHandler()
ch.setLevel(logging.DEBUG)
ch.setFormatter(formatter)
logger.setLevel(logging.INFO)
logger.addHandler(ch)

stats = {
  'ldap_found': 0,
  'sql_errors': 0,
  'users_added': 0,
  'users_failed_to_add': 0,
  'users_failed_to_update': 0,
  'users_purged': 0,
  'user_contacts_updated': 0,
  'user_names_updated': 0,
  'user_photos_updated': 0,
  'users_reactivated': 0,
  'users_failed_to_reactivate': 0,
}

LDAP_SETTINGS = {}


def normalize_phone_number(num):
    return format_number(parse(num.decode('utf-8'), 'US'), PhoneNumberFormat.INTERNATIONAL)


def get_predefined_users(config):
    users = {}

    try:
        config_users = config['sync_script']['preset_users']
    except KeyError:
        return {}

    for user in config_users:
        users[user['name']] = user
        for key in ['sms', 'call']:
            try:
                users[user['name']][key] = normalize_phone_number(users[user['name']][key])
            except (NumberParseException, KeyError, AttributeError):
                users[user['name']][key] = None

    return users


def timestamp_to_human_str(timestamp, tz):
    dt = datetime.fromtimestamp(timestamp, timezone(tz))
    return ' '.join([dt.strftime('%Y-%m-%d %H:%M:%S'), tz])


def prune_user(engine, username):
    global stats
    stats['users_purged'] += 1

    try:
        engine.execute('DELETE FROM `user` WHERE `name` = %s', username)
        logger.info('Deleted inactive user %s', username)

    # The user has messages or some other user data which should be preserved. Just mark as inactive.
    except IntegrityError:
        logger.info('Marking user %s inactive', username)
        engine.execute('UPDATE `user` SET `active` = FALSE WHERE `name` = %s', username)

    except SQLAlchemyError as e:
        logger.error('Deleting user %s failed: %s', username, e)
        stats['sql_errors'] += 1


def fetch_ldap():
    ldap.set_option(ldap.OPT_X_TLS_REQUIRE_CERT, ldap.OPT_X_TLS_ALLOW)
    l = ldap.initialize(LDAP_SETTINGS['url'])
    l.set_option(ldap.OPT_X_TLS_CACERTFILE, LDAP_SETTINGS['cert_path'])
    l.simple_bind_s(LDAP_SETTINGS['user'], LDAP_SETTINGS['password'])

    req_ctrl = SimplePagedResultsControl(True, size=1000, cookie='')

    known_ldap_resp_ctrls = {
        SimplePagedResultsControl.controlType: SimplePagedResultsControl,
    }

    base = LDAP_SETTINGS['base']
    attrs = ['distinguishedName'] + LDAP_SETTINGS['attrs'].values()
    query = LDAP_SETTINGS['query']

    users = {}
    dn_map = {}

    while True:
        msgid = l.search_ext(base, ldap.SCOPE_SUBTREE, query, attrs, serverctrls=[req_ctrl])
        rtype, rdata, rmsgid, serverctrls = l.result3(msgid, resp_ctrl_classes=known_ldap_resp_ctrls)
        logger.info('Loaded %d entries from ldap.' % len(rdata))
        for dn, ldap_dict in rdata:
            if LDAP_SETTINGS['attrs']['mail'] not in ldap_dict:
                logger.error('ERROR: invalid ldap entry for dn: %s' % dn)
                continue

            username = ldap_dict['sAMAccountName'][0]

            mobile = ldap_dict.get(LDAP_SETTINGS['attrs']['mobile'])
            mail = ldap_dict.get(LDAP_SETTINGS['attrs']['mail'])
            name = ldap_dict.get(LDAP_SETTINGS['attrs']['full_name'])[0]

            if mobile:
                try:
                    mobile = normalize_phone_number(mobile[0])
                except NumberParseException:
                    mobile = None
                except UnicodeEncodeError:
                    mobile = None

            if mail:
                mail = mail[0]
                slack = mail.split('@')[0]
            else:
                slack = None

            contacts = {'call': mobile, 'sms': mobile, 'email': mail, 'slack': slack, 'name': name}
            dn_map[dn] = username
            users[username] = contacts

        pctrls = [
            c for c in serverctrls if c.controlType == SimplePagedResultsControl.controlType
            ]

        cookie = pctrls[0].cookie
        if not cookie:
            break
        req_ctrl.cookie = cookie

    return users


def sync(config, engine):
    Session = sessionmaker(bind=engine)
    session = Session()
    oncall_users = {}
    users_query = '''SELECT `user`.`name` as `name`, `contact_mode`.`name` as `mode`, `user_contact`.`destination`,
                            `user`.`full_name`, `user`.`photo_url`
                     FROM `user`
                     LEFT OUTER JOIN `user_contact` ON `user`.`id` = `user_contact`.`user_id`
                     LEFT OUTER JOIN `contact_mode` ON `user_contact`.`mode_id` = `contact_mode`.`id`
                     ORDER BY `user`.`name`'''
    for row in engine.execute(users_query):
        contacts = oncall_users.setdefault(row.name, {})
        if row.mode is None or row.destination is None:
            continue
        contacts[row.mode] = row.destination
        contacts['full_name'] = row.full_name
        contacts['photo_url'] = row.photo_url

    oncall_usernames = set(oncall_users)

    # users from ldap and config file
    ldap_users = fetch_ldap()
    stats['ldap_found'] += len(ldap_users)
    ldap_users.update(get_predefined_users(config))
    ldap_usernames = set(ldap_users)

    # set of ldap users not in oncall
    users_to_insert = ldap_usernames - oncall_usernames
    # set of existing oncall users that are in ldap
    users_to_update = oncall_usernames & ldap_usernames
    # set of users in oncall but not ldap, assumed to be inactive
    inactive_users = oncall_usernames - ldap_usernames
    # users who need to be deactivated
    rows = engine.execute('SELECT name FROM user WHERE active = TRUE AND name IN %s', inactive_users)
    users_to_purge = (user.name for user in rows)

    # set of inactive oncall users who appear in ldap
    rows = engine.execute('SELECT name FROM user WHERE active = FALSE AND name IN %s', ldap_usernames)
    users_to_reactivate = (user.name for user in rows)

    # get objects needed for insertion
    modes = dict(list(session.execute('SELECT `name`, `id` FROM `contact_mode`')))

    user_add_sql = 'INSERT INTO `user` (`name`, `full_name`, `photo_url`) VALUES (%s, %s, %s)'

    # insert users that need to be
    logger.debug('Users to insert:')
    for username in users_to_insert:
        logger.debug('Inserting %s' % username)
        full_name = ldap_users[username].pop('name')
        try:
            user_id = engine.execute(user_add_sql, (username, full_name, LDAP_SETTINGS['image_url'] % username)).lastrowid
        except SQLAlchemyError:
            stats['users_failed_to_add'] += 1
            stats['sql_errors'] += 1
            logger.exception('Failed to add user %s' % username)
            continue
        stats['users_added'] += 1
        for key, value in ldap_users[username].iteritems():
            if value and key in modes:
                logger.debug('\t%s -> %s' % (key, value))
                user_contact_add_sql = 'INSERT INTO `user_contact` (`user_id`, `mode_id`, `destination`) VALUES (%s, %s, %s)'
                engine.execute(user_contact_add_sql, (user_id, modes[key], value))<commit_msg>user_sync/ldap_sync.py: Handle empty sets more gracefully when the set of ldap_usernames and oncall_usernames are identical<commit_after>from gevent import monkey, sleep, spawn
monkey.patch_all()  # NOQA

import sys
import time
import yaml
import logging
import ldap

from oncall import metrics
from ldap.controls import SimplePagedResultsControl
from datetime import datetime
from pytz import timezone
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from sqlalchemy.exc import SQLAlchemyError, IntegrityError
from phonenumbers import format_number, parse, PhoneNumberFormat
from phonenumbers.phonenumberutil import NumberParseException


logger = logging.getLogger()
formatter = logging.Formatter('%(asctime)s %(levelname)s %(name)s %(message)s')
ch = logging.StreamHandler()
ch.setLevel(logging.DEBUG)
ch.setFormatter(formatter)
logger.setLevel(logging.INFO)
logger.addHandler(ch)

stats = {
  'ldap_found': 0,
  'sql_errors': 0,
  'users_added': 0,
  'users_failed_to_add': 0,
  'users_failed_to_update': 0,
  'users_purged': 0,
  'user_contacts_updated': 0,
  'user_names_updated': 0,
  'user_photos_updated': 0,
  'users_reactivated': 0,
  'users_failed_to_reactivate': 0,
}

LDAP_SETTINGS = {}


def normalize_phone_number(num):
    return format_number(parse(num.decode('utf-8'), 'US'), PhoneNumberFormat.INTERNATIONAL)


def get_predefined_users(config):
    users = {}

    try:
        config_users = config['sync_script']['preset_users']
    except KeyError:
        return {}

    for user in config_users:
        users[user['name']] = user
        for key in ['sms', 'call']:
            try:
                users[user['name']][key] = normalize_phone_number(users[user['name']][key])
            except (NumberParseException, KeyError, AttributeError):
                users[user['name']][key] = None

    return users


def timestamp_to_human_str(timestamp, tz):
    dt = datetime.fromtimestamp(timestamp, timezone(tz))
    return ' '.join([dt.strftime('%Y-%m-%d %H:%M:%S'), tz])


def prune_user(engine, username):
    global stats
    stats['users_purged'] += 1

    try:
        engine.execute('DELETE FROM `user` WHERE `name` = %s', username)
        logger.info('Deleted inactive user %s', username)

    # The user has messages or some other user data which should be preserved. Just mark as inactive.
    except IntegrityError:
        logger.info('Marking user %s inactive', username)
        engine.execute('UPDATE `user` SET `active` = FALSE WHERE `name` = %s', username)

    except SQLAlchemyError as e:
        logger.error('Deleting user %s failed: %s', username, e)
        stats['sql_errors'] += 1


def fetch_ldap():
    ldap.set_option(ldap.OPT_X_TLS_REQUIRE_CERT, ldap.OPT_X_TLS_ALLOW)
    l = ldap.initialize(LDAP_SETTINGS['url'])
    l.set_option(ldap.OPT_X_TLS_CACERTFILE, LDAP_SETTINGS['cert_path'])
    l.simple_bind_s(LDAP_SETTINGS['user'], LDAP_SETTINGS['password'])

    req_ctrl = SimplePagedResultsControl(True, size=1000, cookie='')

    known_ldap_resp_ctrls = {
        SimplePagedResultsControl.controlType: SimplePagedResultsControl,
    }

    base = LDAP_SETTINGS['base']
    attrs = ['distinguishedName'] + LDAP_SETTINGS['attrs'].values()
    query = LDAP_SETTINGS['query']

    users = {}
    dn_map = {}

    while True:
        msgid = l.search_ext(base, ldap.SCOPE_SUBTREE, query, attrs, serverctrls=[req_ctrl])
        rtype, rdata, rmsgid, serverctrls = l.result3(msgid, resp_ctrl_classes=known_ldap_resp_ctrls)
        logger.info('Loaded %d entries from ldap.' % len(rdata))
        for dn, ldap_dict in rdata:
            if LDAP_SETTINGS['attrs']['mail'] not in ldap_dict:
                logger.error('ERROR: invalid ldap entry for dn: %s' % dn)
                continue

            username = ldap_dict['sAMAccountName'][0]

            mobile = ldap_dict.get(LDAP_SETTINGS['attrs']['mobile'])
            mail = ldap_dict.get(LDAP_SETTINGS['attrs']['mail'])
            name = ldap_dict.get(LDAP_SETTINGS['attrs']['full_name'])[0]

            if mobile:
                try:
                    mobile = normalize_phone_number(mobile[0])
                except NumberParseException:
                    mobile = None
                except UnicodeEncodeError:
                    mobile = None

            if mail:
                mail = mail[0]
                slack = mail.split('@')[0]
            else:
                slack = None

            contacts = {'call': mobile, 'sms': mobile, 'email': mail, 'slack': slack, 'name': name}
            dn_map[dn] = username
            users[username] = contacts

        pctrls = [
            c for c in serverctrls if c.controlType == SimplePagedResultsControl.controlType
            ]

        cookie = pctrls[0].cookie
        if not cookie:
            break
        req_ctrl.cookie = cookie

    return users


def sync(config, engine):
    Session = sessionmaker(bind=engine)
    session = Session()
    oncall_users = {}
    users_query = '''SELECT `user`.`name` as `name`, `contact_mode`.`name` as `mode`, `user_contact`.`destination`,
                            `user`.`full_name`, `user`.`photo_url`
                     FROM `user`
                     LEFT OUTER JOIN `user_contact` ON `user`.`id` = `user_contact`.`user_id`
                     LEFT OUTER JOIN `contact_mode` ON `user_contact`.`mode_id` = `contact_mode`.`id`
                     ORDER BY `user`.`name`'''
    for row in engine.execute(users_query):
        contacts = oncall_users.setdefault(row.name, {})
        if row.mode is None or row.destination is None:
            continue
        contacts[row.mode] = row.destination
        contacts['full_name'] = row.full_name
        contacts['photo_url'] = row.photo_url

    oncall_usernames = set(oncall_users)

    # users from ldap and config file
    ldap_users = fetch_ldap()
    stats['ldap_found'] += len(ldap_users)
    ldap_users.update(get_predefined_users(config))
    ldap_usernames = set(ldap_users)

    # set of ldap users not in oncall
    users_to_insert = ldap_usernames - oncall_usernames
    # set of existing oncall users that are in ldap
    users_to_update = oncall_usernames & ldap_usernames
    # set of users in oncall but not ldap, assumed to be inactive
    inactive_users = oncall_usernames - ldap_usernames
    # users who need to be deactivated
    if inactive_users:
        rows = engine.execute('SELECT name FROM user WHERE active = TRUE AND name IN %s', inactive_users)
        users_to_purge = (user.name for user in rows)
    else:
        users_to_purge = []


    # set of inactive oncall users who appear in ldap
    rows = engine.execute('SELECT name FROM user WHERE active = FALSE AND name IN %s', ldap_usernames)
    users_to_reactivate = (user.name for user in rows)

    # get objects needed for insertion
    modes = dict(list(session.execute('SELECT `name`, `id` FROM `contact_mode`')))

    user_add_sql = 'INSERT INTO `user` (`name`, `full_name`, `photo_url`) VALUES (%s, %s, %s)'

    # insert users that need to be
    logger.debug('Users to insert:')
    for username in users_to_insert:
        logger.debug('Inserting %s' % username)
        full_name = ldap_users[username].pop('name')
        try:
            user_id = engine.execute(user_add_sql, (username, full_name, LDAP_SETTINGS['image_url'] % username)).lastrowid
        except SQLAlchemyError:
            stats['users_failed_to_add'] += 1
            stats['sql_errors'] += 1
            logger.exception('Failed to add user %s' % username)
            continue
        stats['users_added'] += 1
        for key, value in ldap_users[username].iteritems():
            if value and key in modes:
                logger.debug('\t%s -> %s' % (key, value))
                user_contact_add_sql = 'INSERT INTO `user_contact` (`user_id`, `mode_id`, `destination`) VALUES (%s, %s, %s)'
                engine.execute(user_contact_add_sql, (user_id, modes[key], value))
{'fraction_non_alphanumeric': 0.11313079954239227, 'fraction_numerical': 0.0035591712215584086, 'mean_word_length': 2.923940149625935, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7299944', 'n_tokens_mistral': 5041, 'n_tokens_neox': 4844, 'n_words': 1201}