Program to migrate LDAP data to hashbang/userdb

test.py

# pylint: disable-all
# Search LDAP for all People of hashbang.sh
import json
import sys

import ldap
import ldap.resiter
import psycopg2


class ResiterLDAPObject(ldap.ldapobject.LDAPObject,
                        ldap.resiter.ResultProcessor):
    pass


ldap_con = ResiterLDAPObject("ldap://ldap.hashbang.sh")
ldap_con.start_tls_s()
ldap_con.simple_bind_s('cn=provisor,ou=Admin,dc=hashbang,dc=sh',
                       sys.argv[1])

userdb_con = psycopg2.connect("dbname=userdb")

msg_id = ldap_con.search("dc=hashbang,dc=sh", ldap.SCOPE_SUBTREE,
                         "(uid=*)")

insert_users = """
INSERT INTO passwd (uid, name, host, data)
VALUES (%s, %s, %s, %s);"""
insert_hosts = "INSERT INTO hosts (name, maxusers) VALUES (%s, %s)"

hosts = set()
for res_type, res_data, res_msgid, res_controls in ldap_con.allresults(msg_id):
    for dn, entry in res_data:
        uid, = entry["uidNumber"]
        name, = entry["uid"]
        host, = entry["host"]
        shell, = entry["loginShell"]
        keys = [key.decode("ascii") for key in entry["sshPublicKey"]]
        data = json.dumps({"shell": shell.decode("ascii"), "ssh_keys": keys})
        values = [value.decode("ascii") for value in [uid, name, host]]
        cursor = userdb_con.cursor()
        if host not in hosts:
            print("adding host:", host)
            try:
                cursor.execute(insert_hosts, [host.decode("ascii"), 1000])
                userdb_con.commit()
            except psycopg2.IntegrityError:
                userdb_con.rollback()
            hosts.add(host)
        try:
            cursor.execute(insert_users, values + [data])
            userdb_con.commit()
        except psycopg2.IntegrityError as e:
            print("Error adding user:", name, e)
            print(values + [data])
            userdb_con.rollback()
        finally:
            cursor.close()