riposte-marketplace/integrations/active-directory/scripts/ad_get_user.py

import json, os, sys

S = json.loads(os.environ.get("INTEGRATION_SECRETS", "{}"))
I = json.loads(os.environ.get("INTEGRATION_INPUTS", "{}"))


def out(value):
    print(json.dumps(value, default=str))


def fail(message, **extra):
    payload = {"error": message}
    payload.update(extra)
    print(json.dumps(payload, default=str))
    sys.exit(1)


try:
    import ssl
    from ldap3 import (
        Server, Connection, Tls, NTLM, SUBTREE, BASE, ALL_ATTRIBUTES,
        MODIFY_REPLACE, AUTO_BIND_NO_TLS, AUTO_BIND_TLS_BEFORE_BIND,
    )
    from ldap3.core.exceptions import LDAPException
    from ldap3.extend import microsoft
    from ldap3.utils.conv import escape_filter_chars
except ImportError as _e:
    print(json.dumps({
        "error": "The 'ldap3' Python library is required for the Active Directory integration. "
                 "Install it on the execution host (engine): pip install ldap3",
        "detail": str(_e),
    }))
    sys.exit(1)

SSL_VERSIONS = {
    "None": None,
    "TLS": ssl.PROTOCOL_TLS,
    "TLSv1": ssl.PROTOCOL_TLSv1,
    "TLSv1_1": ssl.PROTOCOL_TLSv1_1,
    "TLSv1_2": ssl.PROTOCOL_TLSv1_2,
    "TLS_CLIENT": ssl.PROTOCOL_TLS_CLIENT,
}
CIPHERS = (
    "@SECLEVEL=1:ECDHE+AESGCM:ECDHE+CHACHA20:DHE+AESGCM:DHE+CHACHA20:ECDH+AESGCM:"
    "DH+AESGCM:ECDH+AES:DH+AES:RSA+AESGCM:RSA+AES:!aNULL:!eNULL:!MD5:!DSS"
)

CONN = None


def _bool(v):
    if isinstance(v, bool):
        return v
    return str(v).lower() in ("1", "true", "yes")


def _list(v):
    if isinstance(v, list):
        return v
    return [x.strip() for x in str(v).split(",") if x.strip()]


def _cfg():
    return {
        "host": S.get("server_ip") or S.get("host") or "",
        "port": int(S["port"]) if S.get("port") else None,
        "user": S.get("username") or "",
        "password": S.get("password") or "",
        "base_dn": S.get("base_dn") or "",
        "secure": S.get("secure_connection") or "SSL",
        "ssl_version": S.get("ssl_version") or "None",
        "ntlm": _bool(S.get("ntlm")),
        "unsecure": _bool(S.get("unsecure")),
        "page_size": int(S["page_size"]) if S.get("page_size") else 500,
    }


def _tls(unsecure, ssl_version):
    version = SSL_VERSIONS.get(ssl_version)
    if unsecure:
        return Tls(validate=ssl.CERT_NONE, ca_certs_file=None, ciphers=CIPHERS, version=version)
    return Tls(validate=ssl.CERT_REQUIRED, ca_certs_file=os.environ.get("SSL_CERT_FILE"), version=version)


def _server(c):
    secure = c["secure"]
    if secure in ("SSL", "TLS"):
        return Server(host=c["host"], port=c["port"], use_ssl=True, tls=_tls(c["unsecure"], c["ssl_version"]), connect_timeout=120)
    if secure == "Start TLS":
        return Server(host=c["host"], port=c["port"], use_ssl=False, tls=_tls(c["unsecure"], c["ssl_version"]), connect_timeout=120)
    return Server(host=c["host"], port=c["port"], connect_timeout=120)


def _auto_bind(secure):
    return AUTO_BIND_TLS_BEFORE_BIND if secure == "Start TLS" else AUTO_BIND_NO_TLS


def connect():
    global CONN
    c = _cfg()
    server = _server(c)
    if c["ntlm"]:
        domain_user = c["user"] if "\\" in c["user"] else c["host"] + "\\" + c["user"]
        CONN = Connection(server, domain_user, password=c["password"], authentication=NTLM, auto_bind=_auto_bind(c["secure"]))
    else:
        CONN = Connection(server, user=c["user"], password=c["password"], auto_bind=_auto_bind(c["secure"]))
    return CONN, c


def base_dn_arg():
    return I.get("base-dn") or I.get("base_dn") or _cfg()["base_dn"]


def search_paging(conn, search_filter, search_base, attributes=None, page_size=100, size_limit=0, time_limit=0):
    entries = []
    cookie = None
    total = 0
    left = size_limit
    while True:
        ps = page_size
        if 0 < left < page_size:
            ps = left
        conn.search(search_base, search_filter, search_scope=SUBTREE, attributes=attributes,
                    paged_size=ps, paged_cookie=cookie, time_limit=time_limit)
        left -= len(conn.entries)
        total += len(conn.entries)
        try:
            cookie = conn.result["controls"]["1.2.840.113556.1.4.319"]["value"]["cookie"]
        except (KeyError, TypeError):
            cookie = None
        entries.extend(conn.entries)
        if (size_limit and size_limit <= total) or not cookie:
            break
    flat = []
    for entry in entries:
        obj = json.loads(entry.entry_to_json())
        flat_entry = {"dn": obj["dn"]}
        for attr in obj.get("attributes", {}):
            flat_entry[attr] = obj["attributes"][attr]
        flat.append(flat_entry)
    return flat


def user_dn(conn, sam, search_base):
    if "\\" in sam:
        sam = sam.split("\\", 1)[1]
    flat = search_paging(conn, f"(&(objectClass=user)(sAMAccountName={escape_filter_chars(sam)}))",
                         search_base, page_size=1, size_limit=1)
    if not flat:
        raise Exception(f"Could not find user '{sam}'")
    return flat[0]["dn"]


def computer_dn(conn, name, search_base):
    flat = search_paging(conn, f"(&(objectClass=user)(objectCategory=computer)(name={escape_filter_chars(name)}))",
                         search_base, page_size=1, size_limit=1)
    if not flat:
        raise Exception(f"Could not find computer '{name}'")
    return flat[0]["dn"]


def group_dn(conn, name, search_base):
    flat = search_paging(conn, f"(&(objectClass=group)(cn={escape_filter_chars(name)}))",
                         search_base, page_size=1, size_limit=1)
    if not flat:
        raise Exception(f"Could not find group '{name}'")
    return flat[0]["dn"]


def modify(conn, dn, modification):
    if not conn.modify(dn, modification):
        raise Exception(f"Failed to modify {dn}: {json.dumps(conn.result, default=str)}")


def first_uac(values):
    if isinstance(values, list):
        return values[0] if values else 0
    return values or 0


def uac_fields(value):
    value = int(value)
    flags = {
        "SCRIPT": 0x1, "ACCOUNTDISABLE": 0x2, "HOMEDIR_REQUIRED": 0x8, "LOCKOUT": 0x10,
        "PASSWD_NOTREQD": 0x20, "PASSWD_CANT_CHANGE": 0x40, "ENCRYPTED_TEXT_PWD_ALLOWED": 0x80,
        "TEMP_DUPLICATE_ACCOUNT": 0x100, "NORMAL_ACCOUNT": 0x200, "INTERDOMAIN_TRUST_ACCOUNT": 0x800,
        "WORKSTATION_TRUST_ACCOUNT": 0x1000, "SERVER_TRUST_ACCOUNT": 0x2000, "DONT_EXPIRE_PASSWORD": 0x10000,
        "MNS_LOGON_ACCOUNT": 0x20000, "SMARTCARD_REQUIRED": 0x40000, "TRUSTED_FOR_DELEGATION": 0x80000,
        "NOT_DELEGATED": 0x100000, "USE_DES_KEY_ONLY": 0x200000, "DONT_REQ_PREAUTH": 0x400000,
        "PASSWORD_EXPIRED": 0x800000, "TRUSTED_TO_AUTH_FOR_DELEGATION": 0x1000000, "PARTIAL_SECRETS_ACCOUNT": 0x4000000,
    }
    return {name: bool(value & mask) for name, mask in flags.items()}


def current_uac(conn, search_base):
    sam = I.get("username") or I.get("sAMAccountName")
    query = "(&(objectClass=User)(objectCategory=person))"
    if sam:
        query = f"(&(objectClass=User)(objectCategory=person)(sAMAccountName={escape_filter_chars(sam)}))"
    flat = search_paging(conn, query, search_base, attributes=["userAccountControl"], size_limit=1, page_size=1)
    if flat and flat[0].get("userAccountControl"):
        return first_uac(flat[0]["userAccountControl"])
    return 512



def run():
    conn, c = connect()
    defaults = ["name", "displayName", "memberOf", "mail", "sAMAccountName", "manager", "userAccountControl"]
    query = S.get("default_base_query") or "(&(objectClass=User)(objectCategory=person))"
    esc = escape_filter_chars
    if I.get("dn"):
        query = f"(&(objectClass=User)(objectCategory=person)(distinguishedName={esc(I['dn'])}))"
    elif I.get("name"):
        query = f"(&(objectClass=User)(objectCategory=person)(cn={esc(I['name'])}))"
    elif I.get("email"):
        query = f"(&(objectClass=User)(objectCategory=person)(mail={esc(I['email'])}))"
    elif I.get("username") or I.get("sAMAccountName"):
        u = esc(I.get("username") or I.get("sAMAccountName"))
        query = f"(&(objectClass=User)(objectCategory=person)(sAMAccountName={u}))"
    elif I.get("custom-field-type"):
        if not I.get("custom-field-data"):
            fail("custom-field-data is required when custom-field-type is given")
        query = f"(&(objectClass=User)(objectCategory=person)({esc(I['custom-field-type'])}={esc(I['custom-field-data'])}))"
    custom = I["attributes"].split(",") if I.get("attributes") else []
    exclude = [a.strip() for a in (I.get("attributes-to-exclude") or "").split(",") if a.strip()]
    attributes = list((set(custom) | set(defaults)) - set(exclude))
    limit = int(I.get("limit") or 20)
    if limit <= 0:
        limit = 20
    page_size = int(I.get("page-size") or c["page_size"])
    if I.get("page-size"):
        limit = int(I["page-size"])
    flat = search_paging(conn, query, base_dn_arg(), attributes=attributes, size_limit=limit, page_size=page_size)
    for user in flat:
        if user.get("userAccountControl"):
            user["userAccountControlFields"] = uac_fields(first_uac(user["userAccountControl"]))
    out({"users": flat, "count": len(flat)})


try:
    run()
except SystemExit:
    raise
except LDAPException as e:
    fail("LDAP error: " + str(e))
except Exception as e:
    fail(str(e))
finally:
    try:
        if CONN is not None:
            CONN.unbind()
    except Exception:
        pass