#!/usr/bin/python3

"""
Maps tool users to precise tools
Run like tail -500000 /data/project/.system/accounting | ./precise-owners.py
from tools bastions
"""

import collections
import datetime
import fileinput
import http.client
import json
import ldap3
import operator
import yaml


def ldap_conn(config):
    """
    Return a ldap connection

    Return value can be used as a context manager
    """
    servers = ldap3.ServerPool([
        ldap3.Server(host)
        for host in config['servers']
    ], ldap3.POOLING_STRATEGY_ROUND_ROBIN, active=True, exhaust=True)
    return ldap3.Connection(servers, read_only=True,
        user=config['user'],
        auto_bind=True,
        password=config['password'])


def uid_from_dn(dn):
    keys = dn.split(',')
    uid_key = keys[0]
    uid = uid_key.split('=')[1]
    return uid


def tools_members(config, tools):
    """
    Return a dict that has members of a tool associated with each tool
    Ex:
    {'tools.musikbot': ['musikanimal'],
     'tools.ifttt': ['slaporte', 'mahmoud', 'madhuvishy', 'ori']}
    """
    tool_to_members = collections.defaultdict(list)
    with ldap_conn(config) as conn:
        for tool in tools:
            conn.search(
                'ou=servicegroups,dc=wikimedia,dc=org',
                '(cn={})'.format(tool),
                ldap3.SEARCH_SCOPE_WHOLE_SUBTREE,
                attributes=['member', 'cn'],
                time_limit=5
            )
            for resp in conn.response:
                attributes = resp.get('attributes')
                members = attributes['member'] if attributes.get('member') else []
                tool_to_members[tool].extend([uid_from_dn(member) for member in members])
    return tool_to_members

def members_tools(tool_to_members):
    all_members = set().union(*tool_to_members.values())
    member_to_tools = collections.defaultdict(list)
    for tool, members in tool_to_members.items():
        for m in members:
            if tool not in member_to_tools[m]:
                member_to_tools[m].append(tool)
    return member_to_tools

def is_precise_host(hostname):
    if hostname[-4:].startswith('12'):
        return True

def grid_precise_tools():
    all_precise_tools = []
    conn = http.client.HTTPConnection('tools.wmflabs.org')
    conn.request("GET", "/gridengine-status",
                 headers={"User-Agent": "Precise tools finder|labs-admin@lists.wikimedia.org"})
    res = conn.getresponse().read().decode('utf-8')
    if res:
        grid_info = json.loads(res)["data"]["attributes"]
    for hostname, info in grid_info.items():
        if is_precise_host(hostname):
            if info["jobs"]:
                all_precise_tools.extend([job["job_owner"] for job in info["jobs"].values()])
    return all_precise_tools

def accounting_tools():
    DAYS=7
    FIELD_NAMES = [
        'qname', 'hostname', 'group', 'owner', 'job_name', 'job_number', 'account',
        'priority', 'submission_time', 'start_time', 'end_time', 'failed',
        'exit_status', 'ru_wallclock', 'ru_utime', 'ru_stime', 'ru_maxrss',
        'ru_ixrss', 'ru_ismrss', 'ru_idrss', 'ru_isrss', 'ru_minflt', 'ru_majflt',
        'ru_nswap', 'ru_inblock', 'ru_oublock', 'ru_msgsnd', 'ru_msgrcv',
        'ru_nsignals', 'ru_nvcsw', 'ru_nivcsw', 'project', 'department',
        'granted_pe', 'slots', 'task_number', 'cpu', 'mem', 'io', 'category',
        'iow', 'pe_taskid', 'maxvemem', 'arid', 'ar_submission_time',
    ]

    cutoff = (datetime.datetime.now() - datetime.timedelta(days=DAYS)).timestamp()
    precise_tools = []

    for line in fileinput.input():
        parts = line.split(':')
        job = dict(zip(FIELD_NAMES, parts))
        if int(job['end_time']) < cutoff:
            continue
        if 'release=precise' in job['category'] and job['owner'] not in precise_tools:
            precise_tools.append(job['owner'])

    return precise_tools

with open('/etc/ldap.yaml') as f:
    config = yaml.safe_load(f)

tool_to_members = tools_members(config, accounting_tools() + grid_precise_tools())
mt = members_tools(tool_to_members)
print json.dumps(mt, sort_keys=True, indent=4, separators=(',', ': '))