#!/usr/bin/env python2

from qumulo.rest_client import RestClient

import argparse
import getpass

def parse_group_file(group_file):
    groups = {}
    gid_to_name_map = {}
    with open(group_file) as f:
        for line in f.readlines():
            line = line.rstrip()

            chunks = line.split(':')
            gid = int(chunks[2])
            members = set()
            if chunks[3]:
                members = set(chunks[3].split(','))

            # Duplicate gids indicate that we should merge groups, otherwise
            # create a new group
            if gid in gid_to_name_map:
                group_name = gid_to_name_map[gid]
            else:
                group_name = chunks[0]
                groups[group_name] = {
                    'gid': gid,
                    'members': set()
                }

            groups[group_name]['members'].update(members)
            gid_to_name_map[gid] = group_name

    return groups

def sync_groups(client, nis_groups, dry_run):
    # XXX Does not cleanly handle the case where two groups swap gid
    nis_groups = {
        'cp_group_' + name:info
        for name,info in nis_groups.items()
    }

    qumulo_groups = {
        u['name'].lower(): u
        for u in client.groups.list_groups()
    }

    new_groups = {
        name:info['gid'] for name,info in nis_groups.items()
        if name not in qumulo_groups
    }

    existing_groups = set([
        name for name in nis_groups.keys()
        if name in qumulo_groups
    ])

    dead_groups = set([
        name for name in qumulo_groups.keys()
        if name not in nis_groups and
           name.startswith('cp_group_')
    ])

    # XXX What happens if there are still users in the group?
    for name in dead_groups:
        if dry_run:
            print "Delete group " + name
            continue

        client.groups.delete_group(qumulo_groups[name]['id'])

    for name, gid in new_groups.items():
        if dry_run:
            print "Add group {} with gid {}".format(name, gid)
            continue

        client.groups.add_group(name, gid)

    for name in existing_groups:
        nis_gid = nis_groups[name]['gid']
        qumulo_gid = int(qumulo_groups[name]['gid'])
        if nis_gid != qumulo_gid:
            if dry_run:
                print "Change group {} gid from {} to {}".format(
                    name, qumulo_gid, nis_gid)
                continue

            client.groups.modify_group(
                qumulo_groups[name]['id'],
                qumulo_groups[name],
                nis_gid)

def parse_passwd_file(passwd_file):
    users = {}
    with open(passwd_file) as f:
        for line in f.readlines():
            chunks = line.split(':')
            users[chunks[0]] = int(chunks[2])

    return users

def sync_users(client, nis_users, dry_run):
    # XXX Does not cleanly handle the case where two users swap uid

    nis_users = {
        'cp_user_' + name:uid
        for name,uid in nis_users.items()
    }

    qumulo_users = {
        u['name'].lower(): u
        for u in client.users.list_users()
    }

    new_users = {
        name:uid for name, uid in nis_users.items()
        if name not in qumulo_users
    }

    existing_users = set([
        name for name in nis_users.keys()
        if name in qumulo_users
    ])

    dead_users = set([
        name for name in qumulo_users.keys()
        if name not in nis_users and
           name.startswith('cp_user_')
    ])

    for name in dead_users:
        if dry_run:
            print "Remove user " + name
            continue

        client.users.delete_user(qumulo_users[name]['id'])

    default_users_group_id = 513
    for name, uid in new_users.items():
        if dry_run:
            print "Add user {} with uid {}".format(name, uid)
            continue

        client.users.add_user(
            name, default_users_group_id, uid)

    for name in existing_users:
        nis_uid = nis_users[name]
        qumulo_uid = int(qumulo_users[name]['uid'])
        if nis_uid != qumulo_uid:
            if dry_run:
                print "Change user {} uid from {} to {}".format(
                    name, qumulo_uid, nis_uid)
                continue

            client.users.modify_user(
                qumulo_users[name]['id'],
                qumulo_users[name],
                default_users_group_id,
                nis_uid)

def sync_group_membership(client, nis_groups, dry_run):
    nis_groups = {
        'cp_group_' + name:info
        for name,info in nis_groups.items()
    }

    qumulo_groups = {
        u['name'].lower(): u
        for u in client.groups.list_groups()
    }

    qumulo_users = {
        u['name'].lower(): u
        for u in client.users.list_users()
    }

    for group_name, info in qumulo_groups.items():
        if not group_name.startswith('cp_group_'):
            continue

        qumulo_members = set([
            member['name'] for member in
            client.groups.group_get_members(info['id'])
        ])

        # Check if the group has been deleted
        if group_name not in nis_groups:
            continue

        nis_members = [
            'cp_user_' + name
            for name in nis_groups[group_name]['members']
        ]

        new_members = [
            name for name in nis_members
            if name not in qumulo_members
        ]

        dead_members = [
            name for name in qumulo_members
            if name not in nis_members
        ]

        for name in new_members:
            if dry_run:
                print 'Add user {} to group {}'.format(
                    name, group_name)
                continue

            qumulo_user = qumulo_users[name]
            client.groups.group_add_member(
                info['id'], qumulo_user['id'])

        for name in dead_members:
            if dry_run:
                print 'Remove user {} from group {}'.format(
                    name, group_name)
                continue

            qumulo_user = qumulo_users[name]
            client.groups.group_remove_member(
                info['id'], qumulo_user['id'])

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--host', required=True)
    parser.add_argument('--port', default=8000)
    parser.add_argument('--passwd-file', required=True)
    parser.add_argument('--group-file', required=True)
    parser.add_argument('--dry-run', action='store_true')

    args = parser.parse_args()

    client = RestClient(
        args.host, args.port, reuse_connection=True)

    admin_password = getpass.getpass()
    client.login('admin', admin_password)

    groups = parse_group_file(args.group_file)
    sync_groups(client, groups, args.dry_run)

    users = parse_passwd_file(args.passwd_file)
    sync_users(client, users, args.dry_run)

    sync_group_membership(client, groups, args.dry_run)

if __name__ == '__main__':
    main()