# python populator.py --ldap ldap://localhost:389/ --ldap-username cn=admin,dc=example,dc=org --ldap-password admin --ldap-next-uid 1000 --ldap-gid 100 --password pass --key-file ~/.ssh/id_rsa.pub users.csv

import argparse
import csv

try:
    import ldap
    import ldap.modlist as modlist
except ImportError:
    pass

try:
    from pyad import pyad
except ImportError:
    pass


class AdPopulator(object):
    def __init__(self, args):
        pyad.set_defaults(ldap_server=args.ad, username=args.ad_username, password=args.ad_password)

    def populate(self, username, password=None, key=None):
        ou = pyad.adcontainer.ADContainer.from_dn("ou=SshUsers, dc=example, dc=org")

        kwargs = {}
        if password:
            kwargs["password"] = password
        elif key:
            kwargs["optional_attributes"] = {"sshPublicKey": key}
        else:
            raise ValueError("Neither password nor key passed")

        pyad.aduser.ADUser.create(username, ou, **kwargs)


class LdapPopulator(object):
    def __init__(self, args):
        self.ldap = ldap.initialize(args.ldap)
        self.ldap.simple_bind_s(args.ldap_username, args.ldap_password)

        self.next_uid = args.ldap_next_uid
        self.gid = args.ldap_gid

    def populate(self, username, password=None, key=None):
        dn = "uid=%s,dc=example,dc=org" % username

        attrs = {
            "cn": username,
            "sn": username,
            "objectClass": [
                "top",
                "posixAccount",
                "shadowAccount",
                "inetOrgPerson",
                "organizationalPerson",
                "person"
            ],
            "uidNumber": str(self.next_uid),
            "gidNumber": str(self.gid),
            "homeDirectory": "/home/%s" % username,
        }
        if password:
            attrs["userPassword"] = password
        elif key:
            attrs["objectClass"].append("ldapPublicKey")
            attrs["sshPublicKey"] = key
        else:
            raise ValueError("Neither password nor key passed")

        ldif = modlist.addModlist(attrs)
        self.ldap.add_s(dn, ldif)

        self.next_uid += 1

if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument("--ldap")
    parser.add_argument("--ldap-username")
    parser.add_argument("--ldap-password")
    parser.add_argument("--ldap-next-uid", type=int)
    parser.add_argument("--ldap-gid", type=int)

    parser.add_argument("--ad")
    parser.add_argument("--ad-username")
    parser.add_argument("--ad-password")

    parser.add_argument("--password")
    parser.add_argument("--key-file")

    parser.add_argument("input")

    args = parser.parse_args()

    if args.ldap:
        populator = LdapPopulator(args)
    elif args.ad:
        populator = AdPopulator(args)
    else:
        raise ValueError("No populator specified")

    password = args.password

    key = None
    if args.key_file:
        with open(args.key_file, "r") as f:
            key = f.read()

    with open(args.input, "r") as f:
        reader = csv.reader(f)
        for username, use_key in reader:
            kwargs = {"username": username}
            if use_key == "1":
                kwargs["key"] = key
            else:
                kwargs["password"] = password

            populator.populate(**kwargs)