misc_scripts/ldap_extra/getAttribs.py

#!/usr/bin/env python3

import argparse
import csv
import sys

import ldap
from ldap.filter import escape_filter_chars

# ============================================================================
# Configuration
# ============================================================================

LDAP_URI = "ldaps://ldap.example.com"
BIND_DN = "uid=admin,ou=People,dc=example,dc=com"
BIND_PW = "password"
BASE_DN = "dc=example,dc=com"

DEFAULT_KEY_ATTRIBUTE = "customStaffNo"
DEFAULT_DELIMITER = "|"
DEFAULT_OUTPUT_FILE = "/tmp/ldap_output.csv"

INFO = "[i]"
START = "[*]"
SUCCESS = "[+]"
WARN = "[!]"
ERROR = "[x]"

# ============================================================================


def read_input_file(filename):
    values = []

    with open(filename, "r") as f:
        for line in f:
            line = line.strip()
            if line:
                values.append(line)

    return values


def ldap_connect():
    conn = ldap.initialize(LDAP_URI)
    conn.protocol_version = ldap.VERSION3
    conn.simple_bind_s(BIND_DN, BIND_PW)
    return conn


def decode(value):
    if isinstance(value, bytes):
        return value.decode("utf-8", errors="replace")
    return str(value)


def update_progress(current, total, matched, not_found):
    """Update the progress display in-place."""

    sys.stdout.write("\033[3A")

    sys.stdout.write(f"\r{INFO} Progress            : {current:<6} / {total}\n")
    sys.stdout.write(f"\r{INFO} Entries matched     : {matched:<6}\n")
    sys.stdout.write(f"\r{INFO} Not found           : {not_found:<6}\n")

    sys.stdout.flush()


def main():
    parser = argparse.ArgumentParser()

    group = parser.add_mutually_exclusive_group(required=True)

    group.add_argument(
        "-F",
        "--file",
        help="Input file containing one value per line",
    )

    group.add_argument(
        "-v",
        "--values",
        help="Comma-separated list of search values",
    )

    parser.add_argument(
        "-k",
        "--key",
        default=DEFAULT_KEY_ATTRIBUTE,
        help=f"LDAP attribute used for searching (default: {DEFAULT_KEY_ATTRIBUTE})",
    )

    parser.add_argument(
        "-a",
        "--attributes",
        required=True,
        help="Comma-separated LDAP attributes to retrieve",
    )

    parser.add_argument(
        "-o",
        "--output",
        default=DEFAULT_OUTPUT_FILE,
        help=f"Output CSV file (default: {DEFAULT_OUTPUT_FILE})",
    )

    parser.add_argument(
        "-d",
        "--delimiter",
        default=DEFAULT_DELIMITER,
        help=f"Output CSV delimiter (default: {DEFAULT_DELIMITER})",
    )

    args = parser.parse_args()

    attributes = [x.strip() for x in args.attributes.split(",") if x.strip()]

    if args.file:
        search_values = read_input_file(args.file)
    else:
        search_values = [
            x.strip()
            for x in args.values.split(",")
            if x.strip()
        ]

    total = len(search_values)

    matched = 0
    not_found = []
    incomplete = []

    print(f"{START} Fetching LDAP attributes...")
    print(f"{INFO} Search key          : {args.key}")
    print(f"{INFO} Requested attributes: {', '.join(attributes)}")
    print(f"{INFO} Input values        : {total}")

    print(f"{INFO} Progress            : 0 / {total}")
    print(f"{INFO} Entries matched     : 0")
    print(f"{INFO} Not found           : 0")

    conn = ldap_connect()

    with open(args.output, "w", newline="") as csvfile:
        writer = csv.writer(csvfile, delimiter=args.delimiter)

        writer.writerow(attributes)

        for current, value in enumerate(search_values, start=1):
            flt = f"({args.key}={escape_filter_chars(value)})"

            result = conn.search_s(
                BASE_DN,
                ldap.SCOPE_SUBTREE,
                flt,
                attributes,
            )

            if not result:
                not_found.append(value)
                update_progress(current, total, matched, len(not_found))
                continue

            matched += 1

            _, entry = result[0]

            row = []
            missing_attribute = False

            for attr in attributes:
                vals = entry.get(attr)

                if vals:
                    row.append(";".join(decode(v) for v in vals))
                else:
                    row.append("")
                    missing_attribute = True

            if missing_attribute:
                incomplete.append(value)

            writer.writerow(row)

            update_progress(current, total, matched, len(not_found))

    conn.unbind_s()

    print()

    if not_found:
        print(f"{WARN} Missing entries")
        print(f"    {args.key}: {', '.join(not_found)}")
        print()

    if incomplete:
        print(f"{WARN} Incomplete entries")
        print("    Missing one or more requested attributes:")
        print(f"    {args.key}: {', '.join(incomplete)}")
        print()

    print(f"{SUCCESS} Output written to {args.output}")


if __name__ == "__main__":
    main()