riposte-marketplace/integrations/openldap/scripts/test_connection.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 (
        ALL_ATTRIBUTES, ALL_OPERATIONAL_ATTRIBUTES, AUTO_BIND_NO_TLS,
        AUTO_BIND_TLS_BEFORE_BIND, BASE, SUBTREE, Connection, Server, Tls,
    )
    from ldap3.core.exceptions import (
        LDAPBindError, LDAPException, LDAPInvalidDnError, LDAPInvalidPortError,
        LDAPSocketOpenError, LDAPSocketReceiveError, LDAPStartTLSError,
    )
    from ldap3.utils.dn import parse_dn
except ImportError as _e:
    print(json.dumps({
        "error": "The 'ldap3' Python library is required for the OpenLDAP integration. "
                 "Install it on the execution host (engine): pip install ldap3",
        "detail": str(_e),
    }))
    sys.exit(1)

MAX_PAGE_SIZE = 2000


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


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


def arg_to_num(v):
    if v is None or v == "":
        return None
    return int(v)


def list_arg_to_ldap_filter(arg, prefix):
    items = arg_to_list(arg)
    joined = "".join(f"({prefix}={item})" for item in items)
    if len(items) > 1:
        return f"(&{joined})"
    return joined if items else ""


def create_entries_search_filter(args):
    cn = list_arg_to_ldap_filter(args.get("cn", ""), "cn")
    description = list_arg_to_ldap_filter(args.get("description", ""), "description")
    object_class = list_arg_to_ldap_filter(args.get("object_class", ""), "objectClass")
    uid = list_arg_to_ldap_filter(args.get("uid", ""), "uid")
    search_filter = args.get("search_filter", "")
    if not any([cn, description, object_class, uid, search_filter]):
        return "(objectClass=*)"
    return f"(|{cn}{description}{object_class}{uid}{search_filter})"


def get_search_attributes(attributes):
    if attributes == "all":
        return [ALL_ATTRIBUTES, ALL_OPERATIONAL_ATTRIBUTES]
    return {
        "none": None,
        "all_user_attributes": ALL_ATTRIBUTES,
        "all_operational_attributes": ALL_OPERATIONAL_ATTRIBUTES,
    }.get(attributes, arg_to_list(attributes))


def entries_paged_search(connection, search_params, page, page_size):
    if page == 1:
        return connection.search(**search_params, paged_size=page_size)
    results_to_skip = page_size * (page - 1)
    connection.search(**search_params, paged_size=results_to_skip)
    cookie = connection.result["controls"]["1.2.840.113556.1.4.319"]["value"]["cookie"]
    return connection.search(**search_params, paged_size=page_size, paged_cookie=cookie)


class LdapClient:
    """LDAP authentication client supporting OpenLDAP and Active Directory."""

    OPENLDAP = "OpenLDAP"
    ACTIVE_DIRECTORY = "Active Directory"
    AUTO = "Auto"
    GROUPS_TOKEN = "primaryGroupToken"
    GROUPS_MEMBER = "memberOf"
    GROUPS_PRIMARY_ID = "primaryGroupID"
    TIMEOUT = 120
    CIPHERS_STRING = (
        "@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"
    )
    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,
    }

    def __init__(self, cfg):
        self._host = cfg.get("host")
        self._port = int(cfg["port"]) if cfg.get("port") else None
        self._username = cfg.get("username", "")
        self._password = cfg.get("password", "")
        self._base_dn = (cfg.get("base_dn") or "").strip()
        self._connection_type = (cfg.get("connection_type") or "none").lower()
        self._ssl_version = cfg.get("ssl_version", "None")
        self._fetch_groups = arg_to_bool(cfg.get("fetch_groups", True))
        self._verify = not arg_to_bool(cfg.get("insecure", False))
        self._ldap_server = self._initialize_ldap_server()
        self._ldap_server_vendor = cfg.get("ldap_server_vendor", self.AUTO)
        if self._ldap_server_vendor == self.AUTO:
            self._determine_ldap_vendor_automatically()
        self._page_size = int(cfg.get("page_size") or 500)
        self._groups_filter_class = (cfg.get("group_filter_class") or "posixGroup").strip()
        self._group_identifier_attribute = (cfg.get("group_identifier_attribute") or "gidNumber").strip()
        self._member_identifier_attribute = (cfg.get("member_identifier_attribute") or "memberUid").strip()
        self._user_filter_class = cfg.get("user_filter_class") or "posixAccount"
        self._user_identifier_attribute = cfg.get("user_identifier_attribute") or "uid"
        self._custom_attributes = cfg.get("custom_attributes") or ""

    @property
    def GROUPS_OBJECT_CLASS(self):
        return self._groups_filter_class

    @property
    def GROUPS_IDENTIFIER_ATTRIBUTE(self):
        return self._group_identifier_attribute

    @property
    def GROUPS_MEMBERSHIP_IDENTIFIER_ATTRIBUTE(self):
        return self._member_identifier_attribute

    @property
    def USER_OBJECT_CLASS(self):
        return self._user_filter_class

    @property
    def USER_IDENTIFIER_ATTRIBUTE(self):
        return self._user_identifier_attribute

    @property
    def CUSTOM_ATTRIBUTE(self):
        return self._custom_attributes

    def _get_ssl_version(self):
        return self.SSL_VERSIONS.get(self._ssl_version)

    def _get_tls_object(self):
        if self._verify:
            return Tls(validate=ssl.CERT_REQUIRED, ca_certs_file=os.environ.get("SSL_CERT_FILE"), version=self._get_ssl_version())
        return Tls(validate=ssl.CERT_NONE, ca_certs_file=None, version=self._get_ssl_version(), ciphers=self.CIPHERS_STRING)

    def _initialize_ldap_server(self):
        if self._connection_type == "ssl":
            return Server(host=self._host, port=self._port, use_ssl=True, tls=self._get_tls_object(), connect_timeout=self.TIMEOUT)
        if self._connection_type == "start tls":
            return Server(host=self._host, port=self._port, use_ssl=False, tls=self._get_tls_object(), connect_timeout=self.TIMEOUT)
        return Server(host=self._host, port=self._port, connect_timeout=self.TIMEOUT)

    def _determine_ldap_vendor_automatically(self):
        try:
            with Connection(self._ldap_server) as conn:
                conn.search(search_base="", search_filter="(objectClass=*)", search_scope=BASE, attributes=[ALL_ATTRIBUTES])
                entry = conn.entries[0]
                if "objectClass" in entry and "OpenLDAProotDSE" in entry["objectClass"].value:
                    self._ldap_server_vendor = self.OPENLDAP
                else:
                    self._ldap_server_vendor = self.ACTIVE_DIRECTORY
        except Exception as e:
            raise Exception(f"Could not determine the LDAP vendor automatically. Select the vendor manually. Error: {e}")

    @staticmethod
    def _parse_ldap_group_entries(ldap_group_entries, groups_identifier_attribute):
        return [
            {
                "DN": ldap_group.get("dn"),
                "Attributes": [{
                    "Name": LdapClient.GROUPS_TOKEN,
                    "Values": [str(ldap_group.get("attributes", {}).get(groups_identifier_attribute))],
                }],
            }
            for ldap_group in ldap_group_entries
        ]

    @staticmethod
    def _parse_ldap_group_entries_and_referrals(ldap_group_entries):
        referrals = []
        entries = []
        for ldap_group in ldap_group_entries:
            group_type = ldap_group.get("type")
            if group_type == "searchResRef":
                referrals.extend(ldap_group.get("uri") or [])
            elif group_type == "searchResEntry":
                entries.append({
                    "DN": ldap_group.get("dn"),
                    "Attributes": [{
                        "Name": LdapClient.GROUPS_TOKEN,
                        "Values": [str(ldap_group.get("attributes", {}).get(LdapClient.GROUPS_TOKEN))],
                    }],
                })
        return referrals, entries

    def _parse_and_authenticate_ldap_group_entries_and_referrals(self, ldap_group_entries, password):
        referrals = []
        entries = []
        for entry in ldap_group_entries:
            entry_type = entry.get("type")
            if entry_type == "searchResRef":
                referrals.extend(entry.get("uri") or [])
            elif entry_type == "searchResEntry":
                entry_dn = entry.get("dn", "")
                entry_attributes = entry.get("attributes", {})
                relevant = []
                for attr in entry_attributes:
                    attr_value = entry_attributes.get(attr, [])
                    if attr_value:
                        if not isinstance(attr_value, list):
                            attr_value = [str(attr_value)]
                        relevant.append({"Name": attr, "Values": attr_value})
                entries.append({"DN": entry_dn, "Attributes": relevant})
                self.authenticate_ldap_user(entry_dn, password)
        return referrals, entries

    @staticmethod
    def _parse_ldap_users_groups_entries(ldap_group_entries):
        return [ldap_group.get("dn") for ldap_group in ldap_group_entries]

    @staticmethod
    def _build_entry_for_user(user_groups, user_data, mail_attribute, name_attribute, phone_attribute):
        attributes = [
            {"Name": LdapClient.GROUPS_MEMBER, "Values": user_groups},
            {"Name": LdapClient.GROUPS_PRIMARY_ID, "Values": user_data["gid_number"]},
        ]
        if "name" in user_data:
            attributes.append({"Name": name_attribute, "Values": [user_data["name"]]})
        if "email" in user_data:
            attributes.append({"Name": mail_attribute, "Values": [user_data["email"]]})
        if "mobile" in user_data:
            attributes.append({"Name": phone_attribute, "Values": [user_data["mobile"]]})
        return {"DN": user_data["dn"], "Attributes": attributes}

    @staticmethod
    def _is_valid_dn(dn, user_identifier_attribute):
        try:
            parsed_dn = parse_dn(dn, strip=False)
            for attribute_and_value in parsed_dn:
                if attribute_and_value[0].lower() == user_identifier_attribute.lower():
                    return True, attribute_and_value[1]
            raise Exception(f"The {user_identifier_attribute} attribute was not found in the user DN: {dn}")
        except LDAPInvalidDnError:
            return False, dn

    def _get_formatted_custom_attributes(self):
        if not self.CUSTOM_ATTRIBUTE:
            return ""
        formatted = ""
        for att in self.CUSTOM_ATTRIBUTE.split(","):
            if len(att.split("=")) != 2:
                raise Exception(f'User defined attributes must be of the form "attrA=valA,attrB=valB,...", but got: {self.CUSTOM_ATTRIBUTE}')
            formatted += f"({att})"
        return formatted

    def _get_ldap_groups_entries_and_referrals_ad(self, ldap_conn, search_filter):
        ldap_group_entries = ldap_conn.extend.standard.paged_search(
            search_base=self._base_dn, search_filter=search_filter,
            attributes=[LdapClient.GROUPS_TOKEN], paged_size=self._page_size, generator=False,
        )
        return LdapClient._parse_ldap_group_entries_and_referrals(ldap_group_entries)

    def _fetch_all_groups(self):
        auto_bind = self._get_auto_bind_value()
        with Connection(self._ldap_server, self._username, self._password, auto_bind=auto_bind) as ldap_conn:
            if self._ldap_server_vendor == self.ACTIVE_DIRECTORY:
                referrals, entries = self._get_ldap_groups_entries_and_referrals_ad(
                    ldap_conn=ldap_conn, search_filter="(&(objectClass=group)(objectCategory=group))")
                return {"Controls": None, "Referrals": referrals, "Entries": entries}
            ldap_group_entries = ldap_conn.extend.standard.paged_search(
                search_base=self._base_dn, search_filter=f"(objectClass={self.GROUPS_OBJECT_CLASS})",
                attributes=[self.GROUPS_IDENTIFIER_ATTRIBUTE], paged_size=self._page_size)
            return {
                "Controls": None,
                "Referrals": ldap_conn.result.get("referrals"),
                "Entries": LdapClient._parse_ldap_group_entries(ldap_group_entries, self.GROUPS_IDENTIFIER_ATTRIBUTE),
            }

    def _fetch_specific_groups(self, specific_groups):
        auto_bind = self._get_auto_bind_value()
        dn_list = [group.strip() for group in arg_to_list(specific_groups, separator="#")]
        with Connection(self._ldap_server, self._username, self._password, auto_bind=auto_bind) as ldap_conn:
            if self._ldap_server_vendor == self.ACTIVE_DIRECTORY:
                dns_filter = "".join(f"(distinguishedName={dn})" for dn in dn_list)
                search_filter = f"(&(objectClass=group)(objectCategory=group)(|{dns_filter}))"
                referrals, entries = self._get_ldap_groups_entries_and_referrals_ad(
                    ldap_conn=ldap_conn, search_filter=search_filter)
                return {"Controls": None, "Referrals": referrals, "Entries": entries}
            parsed = []
            for dn in dn_list:
                ldap_group_entries = ldap_conn.extend.standard.paged_search(
                    search_base=dn, search_filter=f"(objectClass={self.GROUPS_OBJECT_CLASS})",
                    attributes=[self.GROUPS_IDENTIFIER_ATTRIBUTE], paged_size=self._page_size, search_scope=BASE)
                parsed.append(self._parse_ldap_group_entries(ldap_group_entries, self.GROUPS_IDENTIFIER_ATTRIBUTE))
            return {"Controls": None, "Referrals": ldap_conn.result.get("referrals"), "Entries": parsed}

    @staticmethod
    def _get_ad_username(logon_name):
        if "\\" in logon_name:
            return logon_name.split("\\")[1]
        if "@" in logon_name:
            return logon_name.split("@")[0]
        return logon_name

    @staticmethod
    def _has_wildcards_in_user_logon(logon_name):
        for wildcard in ("*", "?"):
            if wildcard in logon_name:
                raise Exception(
                    f"Authentication failed - wildcards were detected in the user logon name "
                    f"(input username: '{logon_name}'). Wildcards are not permitted for user authentication.")

    def _get_auto_bind_value(self):
        if self._connection_type == "start tls":
            return AUTO_BIND_TLS_BEFORE_BIND
        return AUTO_BIND_NO_TLS

    def get_ldap_groups(self, specific_group=""):
        if not self._fetch_groups and not specific_group:
            return {"Controls": None, "Referrals": None, "Entries": []}
        return self._fetch_specific_groups(specific_group) if not self._fetch_groups else self._fetch_all_groups()

    def authenticate_ldap_user(self, username, password):
        auto_bind = self._get_auto_bind_value()
        ldap_conn = Connection(server=self._ldap_server, user=username, password=password, auto_bind=auto_bind)
        if ldap_conn.bound:
            ldap_conn.unbind()
            return "Done"
        raise Exception(f"Authentication connection failed (server type: {self._ldap_server_vendor})")

    def search_user_data(self, username, attributes, search_user_by_dn=False):
        auto_bind = self._get_auto_bind_value()
        with Connection(self._ldap_server, self._username, self._password, auto_bind=auto_bind) as ldap_conn:
            if search_user_by_dn:
                search_filter = f"(&(objectClass={self.USER_OBJECT_CLASS})" + self._get_formatted_custom_attributes() + ")"
                ldap_conn.search(search_base=username, search_filter=search_filter, size_limit=1,
                                 attributes=attributes, search_scope=BASE)
            else:
                custom_attributes = self._get_formatted_custom_attributes()
                search_filter = (f"(&(objectClass={self.USER_OBJECT_CLASS})"
                                 f"({self.USER_IDENTIFIER_ATTRIBUTE}={username}){custom_attributes})")
                ldap_conn.search(search_base=self._base_dn, search_filter=search_filter, size_limit=1, attributes=attributes)
            if not ldap_conn.entries:
                raise Exception("LDAP user not found")
            entry = ldap_conn.entries[0]
            referrals = ldap_conn.result.get("referrals")
            if self.GROUPS_IDENTIFIER_ATTRIBUTE not in entry or not entry[self.GROUPS_IDENTIFIER_ATTRIBUTE].value:
                raise Exception(f"User's {self.GROUPS_IDENTIFIER_ATTRIBUTE} not found")
            return entry, referrals

    def get_user_data(self, username, pull_name, pull_mail, pull_phone,
                      name_attribute, mail_attribute, phone_attribute, search_user_by_dn=False):
        attributes = [self.GROUPS_IDENTIFIER_ATTRIBUTE]
        if pull_name:
            attributes.append(name_attribute)
        if pull_mail:
            attributes.append(mail_attribute)
        if pull_phone:
            attributes.append(phone_attribute)
        entry, referrals = self.search_user_data(username, attributes, search_user_by_dn)
        user_data = {
            "dn": entry.entry_dn,
            "gid_number": [str(entry[self.GROUPS_IDENTIFIER_ATTRIBUTE].value)],
            "referrals": referrals,
        }
        if name_attribute in entry and entry[name_attribute].value:
            user_data["name"] = entry[name_attribute].value
        if mail_attribute in entry and entry[mail_attribute].value:
            user_data["email"] = entry[mail_attribute].value
        if phone_attribute in entry and entry[phone_attribute].value:
            user_data["mobile"] = entry[phone_attribute].value
        return user_data

    def get_user_groups(self, user_identifier):
        auto_bind = self._get_auto_bind_value()
        with Connection(self._ldap_server, self._username, self._password, auto_bind=auto_bind) as ldap_conn:
            search_filter = (f"(&(objectClass={self.GROUPS_OBJECT_CLASS})"
                             f"({self.GROUPS_MEMBERSHIP_IDENTIFIER_ATTRIBUTE}={user_identifier}))")
            ldap_group_entries = ldap_conn.extend.standard.paged_search(
                search_base=self._base_dn, search_filter=search_filter,
                attributes=[self.GROUPS_IDENTIFIER_ATTRIBUTE], paged_size=self._page_size)
            return LdapClient._parse_ldap_users_groups_entries(ldap_group_entries)

    def authenticate_and_roles_openldap(self, username, password, pull_name=True, pull_mail=True, pull_phone=False,
                                        mail_attribute="mail", name_attribute="name", phone_attribute="mobile"):
        search_user_by_dn, user_identifier = LdapClient._is_valid_dn(username, self.USER_IDENTIFIER_ATTRIBUTE)
        user_data = self.get_user_data(
            username=username, search_user_by_dn=search_user_by_dn, pull_name=pull_name, pull_mail=pull_mail,
            pull_phone=pull_phone, mail_attribute=mail_attribute, name_attribute=name_attribute, phone_attribute=phone_attribute)
        self.authenticate_ldap_user(user_data["dn"], password)
        user_groups = self.get_user_groups(user_identifier)
        return {
            "Controls": None,
            "Referrals": user_data["referrals"],
            "Entries": [LdapClient._build_entry_for_user(
                user_groups=user_groups, user_data=user_data,
                mail_attribute=mail_attribute, name_attribute=name_attribute, phone_attribute=phone_attribute)],
        }

    def authenticate_and_roles_active_directory(self, username, password, pull_name=True, pull_mail=True, pull_phone=False,
                                                mail_attribute="mail", name_attribute="name", phone_attribute="mobile"):
        ad_username = self._get_ad_username(username)
        auto_bind = self._get_auto_bind_value()
        with Connection(self._ldap_server, self._username, self._password, auto_bind=auto_bind) as ldap_conn:
            attributes = [self.GROUPS_MEMBER, self.GROUPS_PRIMARY_ID]
            if pull_name:
                attributes.append(name_attribute)
            if pull_mail:
                attributes.append(mail_attribute)
            if pull_phone:
                attributes.append(phone_attribute)
            search_filter = f"(|(sAMAccountName={ad_username})(userPrincipalName={username}))"
            ldap_conn_entries = ldap_conn.extend.standard.paged_search(
                search_base=self._base_dn, search_filter=search_filter,
                attributes=attributes, paged_size=self._page_size, generator=False)
            referrals, entries = self._parse_and_authenticate_ldap_group_entries_and_referrals(
                ldap_group_entries=ldap_conn_entries, password=password)
            if not entries:
                raise Exception("LDAP user not found")
        return {"Controls": [], "Referrals": referrals, "Entries": entries}

    def authenticate_and_roles(self, username, password, pull_name=True, pull_mail=True, pull_phone=False,
                               mail_attribute="mail", name_attribute="name", phone_attribute="mobile"):
        self._has_wildcards_in_user_logon(username)
        if self._ldap_server_vendor == self.ACTIVE_DIRECTORY:
            return self.authenticate_and_roles_active_directory(
                username=username, password=password, pull_name=pull_name, pull_mail=pull_mail, pull_phone=pull_phone,
                mail_attribute=mail_attribute, name_attribute=name_attribute, phone_attribute=phone_attribute)
        return self.authenticate_and_roles_openldap(
            username=username, password=password, pull_name=pull_name, pull_mail=pull_mail, pull_phone=pull_phone,
            mail_attribute=mail_attribute, name_attribute=name_attribute, phone_attribute=phone_attribute)

    def entries_search_command(self, args):
        search_params = {
            "search_base": args.get("search_base") or self._base_dn,
            "search_scope": args.get("search_scope") or SUBTREE,
            "search_filter": create_entries_search_filter(args),
            "attributes": get_search_attributes(args.get("attributes", "all")),
        }
        auto_bind = self._get_auto_bind_value()
        with Connection(self._ldap_server, self._username, self._password, auto_bind=auto_bind) as ldap_conn:
            page = arg_to_num(args.get("page"))
            if page:
                page_size = int(args.get("page_size", 50))
                if page_size > MAX_PAGE_SIZE:
                    raise Exception("The page size must be less than or equal to 2000")
            else:
                page = 1
                page_size = int(args.get("limit", 50))
            entries_paged_search(connection=ldap_conn, search_params=search_params, page=page, page_size=page_size)
            outputs = [
                {**json.loads(entry.entry_to_json()).get("attributes", {}), "dn": json.loads(entry.entry_to_json()).get("dn")}
                for entry in ldap_conn.entries
            ]
        return {"results": outputs, "count": len(outputs)}

    def ad_authenticate(self, username, password):
        self._has_wildcards_in_user_logon(username)
        if self._ldap_server_vendor == self.OPENLDAP:
            search_user_by_dn, _ = LdapClient._is_valid_dn(username, self.USER_IDENTIFIER_ATTRIBUTE)
            entry, _ = self.search_user_data(username, [self.GROUPS_IDENTIFIER_ATTRIBUTE], search_user_by_dn)
            username = entry.entry_dn
        return self.authenticate_ldap_user(username, password)

    def test_module(self):
        self._get_formatted_custom_attributes()
        if self._ldap_server_vendor == self.OPENLDAP:
            try:
                parse_dn(self._username)
            except LDAPInvalidDnError:
                raise Exception("Invalid credentials input. The bind user must be a full DN.")
        self.authenticate_ldap_user(username=self._username, password=self._password)
        return "ok"


def run():
    client = LdapClient(S)
    client.test_module()
    out({"ok": True, "vendor": client._ldap_server_vendor})


try:
    run()
except SystemExit:
    raise
except LDAPBindError as e:
    fail("Authentication connection failed. Additional details: " + str(e))
except (LDAPSocketOpenError, LDAPSocketReceiveError, LDAPStartTLSError) as e:
    msg = "Failed to connect to the LDAP server. Additional details: " + str(e)
    if not arg_to_bool(S.get("insecure", False)):
        msg += ' Try the "Trust any certificate" option.'
    fail(msg)
except LDAPInvalidPortError:
    fail("Invalid LDAP server input. The server must be of the form: ip or ldap://ip")
except LDAPException as e:
    fail("LDAP error: " + str(e))
except Exception as e:
    fail(str(e))