#!/usr/bin/env python3

import hashlib
import json
import os
import re
import sys
import traceback
from typing import Any, Dict, List, Optional, Set, Tuple, Union

BASE_DIR = "/opt/rbta/ad/mgmtportal/api/core"

if BASE_DIR not in sys.path:
    sys.path.insert(0, BASE_DIR)
os.environ["DJANGO_SETTINGS_MODULE"] = "project.settings"


import ldap
from aldpro_core.utils.logger import Logger
from django.conf import settings
from ldap.controls import SimplePagedResultsControl

logger = Logger("aldpro-check-in-restore")

APPLDAP_IPA_SERVER = settings.LDAP_SERVER
APPLDAP_IPA_USERNAME = settings.DISCOVER_USERNAME
APPLDAP_IPA_PASSWORD = settings.DISCOVER_PASSWORD

LDAP_SSL = settings.LDAP_SSL
LDAP_PROTO = "ldaps" if LDAP_SSL else "ldap"
LDAP_PORT = 636 if LDAP_SSL else 389
BASE_DN = settings.BASE_DN


MAIN_ADMINISTRATOR_PRIVILEGES = {
    "ALDPRO - AD Connections Administrators",
    "ALDPRO - Additional Automation Tasks Administrators",
    "ALDPRO - Automation Tasks Reader",
    "ALDPRO - Computers remote access",
    "ALDPRO - DHCP Servers Administrators",
    "ALDPRO - DNS Administrators",
    "ALDPRO - Delegation Administrators",
    "ALDPRO - Domain Controllers Administrators",
    "ALDPRO - Event Log Administrators",
    "ALDPRO - File Servers Administrators",
    "ALDPRO - Group Policies Administrator",
    "ALDPRO - Group Policy Additional Parameters Administrators",
    "ALDPRO - HBAC Administrator",
    "ALDPRO - Host Administrator",
    "ALDPRO - Host Group Administrator",
    "ALDPRO - IPA Configuration Administrators",
    "ALDPRO - Installer Server Computers Administrators",
    "ALDPRO - Installer Server Systems Administrators",
    "ALDPRO - Installer Servers Administrators",
    "ALDPRO - Kerberos Administrators",
    "ALDPRO - Manage Syncer",
    "ALDPRO - Manage Trusts",
    "ALDPRO - Manage Update Policies",
    "ALDPRO - Monitoring Administrators",
    "ALDPRO - NTP Servers Administrators",
    "ALDPRO - Organizational Units Administrator",
    "ALDPRO - Password Policy Administrators",
    "ALDPRO - Preserved User Manager",
    "ALDPRO - Print Servers Administrators",
    "ALDPRO - Profiles Administrators",
    "ALDPRO - Proxy Addresses Administrators",
    "ALDPRO - Proxy Addresses Reader",
    "ALDPRO - Replication Administrators",
    "ALDPRO - Repository Administrators",
    "ALDPRO - Repository Servers Administrators",
    "ALDPRO - Repository Versions Administrators",
    "ALDPRO - Service Administrators",
    "ALDPRO - Sites and services Administrators",
    "ALDPRO - Software Administrators",
    "ALDPRO - Software Catalogs Administrators",
    "ALDPRO - Software Configuration Templates Administrators",
    "ALDPRO - Software Packages Administrators",
    "ALDPRO - Software Parameters Administrators",
    "ALDPRO - Sudo Administrator",
    "ALDPRO - Sysaccounts Administrator",
    "ALDPRO - User Group Administrator",
    "ALDPRO - User Manager",
    "ALDPRO - User Reader",
    "Replication Administrators",
    "Users History - Read",
}


class CheckInRestore:
    def __ldap_disconnect(self, initialize_ldap: ldap.ldapobject.SimpleLDAPObject) -> None:
        """Закрыть соединение к ldap

        Args:
            initialize_ldap (ldap.ldapobject.SimpleLDAPObject): объект соединения
        """
        initialize_ldap.unbind()

    def __ldap_connect(self, domain_controller: str) -> ldap.ldapobject.SimpleLDAPObject:
        """Открыть соединенеие к ldap

        Args:
            domain_controller (str): Адрес контролера домена, к которому подсодениться

        Returns:
            ldap.ldapobject.SimpleLDAPObject: объект соединения
        """
        ldap_url = f"{LDAP_PROTO}://{domain_controller}:{LDAP_PORT}"
        ldap.set_option(ldap.OPT_X_TLS_REQUIRE_CERT, ldap.OPT_X_TLS_NEVER)
        initialize_ldap = ldap.initialize(ldap_url)
        initialize_ldap.set_option(ldap.OPT_REFERRALS, 0)
        initialize_ldap.protocol_version = ldap.VERSION3
        try:
            initialize_ldap.simple_bind_s(
                f"uid={APPLDAP_IPA_USERNAME},cn=sysaccounts,cn=etc,{BASE_DN}",
                APPLDAP_IPA_PASSWORD,
            )
        except ldap.LDAPError as e:
            self.__ldap_disconnect(initialize_ldap)
            sys.exit(f"Ошибка соединения с ldap {e}")

        return initialize_ldap

    def __fetch_ldap_result(
        self,
        search_base: str,
        filter_str: str = "(objectClass=*)",
        domain_controller: str = APPLDAP_IPA_SERVER,
        attrlist=None,
        scope=ldap.SCOPE_ONELEVEL,
    ) -> List[Tuple[str, Dict[str, List[bytes]]]]:
        """Выполняет поиск в контейнере

        Args:
            domain_controller (str): Адрес контролера домена, к которому подсодениться
            search_base (str): Контейнер для поиска
            filter_str (str): Фильтр для поиска, по умолчанию - (objectClass=*) - вернет все объекты в контейнере

        Returns:
            List[Tuple[str, Dict[str, List[bytes]]]]: Перечень объектов и их cn
        """
        initialize_ldap = self.__ldap_connect(domain_controller)
        PAGE_CONTROL: SimplePagedResultsControl = SimplePagedResultsControl(True, size=150, cookie="")
        if attrlist is None:
            attrlist = [
                "cn",
            ]
        try:
            message_id = initialize_ldap.search_ext(
                base=f"{search_base},{BASE_DN}",
                scope=scope,
                filterstr=filter_str,
                attrlist=attrlist,
                serverctrls=[PAGE_CONTROL],
            )
            result = []
            pages = 0
            while True:
                pages += 1
                _, rdata, _, serverctrls = initialize_ldap.result3(message_id)
                result.extend(rdata)
                controls = [
                    control for control in serverctrls if control.controlType == SimplePagedResultsControl.controlType
                ]
                if not controls:
                    print("The server ignores RFC 2696 control")
                    break
                if not controls[0].cookie:
                    break
                PAGE_CONTROL.cookie = controls[0].cookie
                message_id = initialize_ldap.search_ext(
                    base=f"{search_base},{BASE_DN}",
                    scope=scope,
                    filterstr=filter_str,
                    attrlist=attrlist,
                    serverctrls=[PAGE_CONTROL],
                )
        except ldap.LDAPError as e:
            sys.exit(f"Ошибка поиска сущностей {e}")
        finally:
            self.__ldap_disconnect(initialize_ldap)
        return result

    def run(self) -> None:
        roles = self.__fetch_ldap_result(
            search_base="cn=roles,cn=accounts",
            filter_str="(&(objectClass=aldproMetaRole)(aldproMetaRoleRenderStatus=done))",
            attrlist=["aldproMetaRoleConfiguration", "memberOf"],
            scope=ldap.SCOPE_SUBTREE,
        )

        role_main_admin = self.__fetch_ldap_result(
            search_base="cn=roles,cn=accounts",
            filter_str="(cn=ALDPRO - Main Administrator)",
            attrlist=["aldproMetaRoleConfiguration", "memberOf"],
            scope=ldap.SCOPE_SUBTREE,
        )

        role_main_admin_dn, data = role_main_admin[0]
        has_all_privileges, missing_privileges = self.__is_role_valid(data=data, attrs=MAIN_ADMINISTRATOR_PRIVILEGES)
        if not has_all_privileges:
            for privilege in missing_privileges:
                privilege_dn = f"cn={privilege},cn=privileges,cn=pbac,{settings.BASE_DN}"
                self.__add_attrs(dn=privilege_dn, field="member", attribute=role_main_admin_dn)
                logger.info(f"В привилегию {privilege_dn} добавлена роль {role_main_admin_dn}")

        for role_dn, data in roles:
            has_all_privileges, missing_privileges = self.__is_role_valid(data=data)
            if not has_all_privileges:
                for privilege in missing_privileges:
                    privilege_dn = f"cn={privilege},cn=privileges,cn=pbac,{settings.BASE_DN}"
                    self.__add_attrs(dn=privilege_dn, field="member", attribute=role_dn)
                    logger.info(f"В привилегию {privilege_dn} добавлена роль {role_dn}")

    def __is_role_valid(
        self, data: Dict[str, Any], attrs: Optional[Set[str]] = None
    ) -> Tuple[bool, Set[Union[str, Any]]]:
        """Проверяет наличие в memberof необходимого набора привилегий."""

        aldpro_meta_role_configuration = data.get("aldproMetaRoleConfiguration")[0].decode("utf-8")
        aldpro_meta_role_configuration = json.loads(aldpro_meta_role_configuration)
        members_of = data.get("memberOf")

        privileges_in_config = set()
        if attrs:
            privileges_in_config.update(attrs)

        for key, value in aldpro_meta_role_configuration.items():
            if self.__is_privilege_hash_required(key):
                for data in value:
                    ou = data.get("ou")
                    ou_nested = data.get("ou_nested")
                    site = data.get("site")
                    context = self.__generate_context(ou=ou, ou_nested=ou_nested, site=site)
                    new_cn = self.__generate_new_cn(key, context)
                    privileges_in_config.add(new_cn)
            else:
                privileges_in_config.add(key)

        privileges_in_memberof = set()
        if members_of:
            for memberof in members_of:
                memberof = memberof.decode("utf-8")
                if "privileges" in memberof:
                    privilege_cn = memberof.split(",")[0].replace("cn=", "")
                    privileges_in_memberof.add(privilege_cn)

        has_all_privileges = privileges_in_config <= privileges_in_memberof
        missing_privileges = privileges_in_config - privileges_in_memberof

        return has_all_privileges, missing_privileges

    def __is_privilege_hash_required(self, privileges_name: str) -> bool:
        privilege = self.__fetch_ldap_result(
            filter_str=f"(cn={privileges_name})",
            attrlist=["aldproMetaPrivilegesIsOUPattern", "aldproMetaPrivilegesIsSitePattern"],
            search_base="cn=privileges,cn=pbac",
            scope=ldap.SCOPE_SUBTREE,
        )

        privilege_dn, data = privilege[0]

        is_ou_pattern = data.get("aldproMetaPrivilegesIsOUPattern")[0].decode("utf-8")
        is_site_pattern = data.get("aldproMetaPrivilegesIsSitePattern")[0].decode("utf-8")

        is_privilege_hash_required = False
        if is_ou_pattern == "TRUE" or is_site_pattern == "TRUE":
            is_privilege_hash_required = True

        return is_privilege_hash_required

    def __generate_context(self, ou: str, ou_nested: Optional[bool], site: str) -> str:
        """Генерирует контекст."""

        result_context_list = []
        if ou:
            result_context_list.append(ou)

        if ou_nested:
            result_context_list.append("NESTED")

        if site:
            result_context_list.append(site)

        result_context = "|".join(result_context_list)
        return result_context

    def __generate_new_cn(self, cn: str, context: str) -> str:
        """Генерирует cn с хэшем на основе контекста."""

        hash_ = hashlib.md5(context.encode()).hexdigest()
        new_cn = f"{cn} {hash_}"
        return new_cn

    def __add_attrs(self, dn: str, field: str, attribute: str) -> None:
        """Добавляет атрибут."""

        initialize_ldap = self.__ldap_connect(APPLDAP_IPA_SERVER)
        initialize_ldap.modify_s(dn, [(ldap.MOD_ADD, field, bytes(attribute, encoding="UTF-8"))])

    def __del_attrs(self, dn: str, field: str, attribute: str) -> None:
        """Удаляет атрибут."""

        initialize_ldap = self.__ldap_connect(APPLDAP_IPA_SERVER)
        initialize_ldap.modify_s(dn, [(ldap.MOD_DELETE, field, bytes(attribute, encoding="UTF-8"))])

    def check_invalid_uniqueid(self) -> None:
        """Проверяет наличие у пермишена в атрибуте member nsuniqueid=... и удаляет его."""
        permissions = self.__fetch_ldap_result(
            filter_str=f"(member=*+nsuniqueid=*)",
            attrlist=["member"],
            search_base="cn=permissions,cn=pbac",
            scope=ldap.SCOPE_SUBTREE,
        )

        pattern = r"\+nsuniqueid=[^,]+"

        for permission in permissions:
            permission_dn, data = permission

            members = data.get("member")

            for member in members:
                member = member.decode("utf-8")
                if "+nsuniqueid=" in member:
                    cleaned_member = re.sub(pattern, "", member)
                    try:
                        self.__add_attrs(dn=permission_dn, field="member", attribute=cleaned_member)
                    except ldap.TYPE_OR_VALUE_EXISTS:
                        logger.warning(
                            f"У разрешения {permission_dn} уже существует атрибут member со значением {cleaned_member}"
                        )
                    self.__del_attrs(dn=permission_dn, field="member", attribute=member)


def main() -> None:
    try:
        check_in_restore = CheckInRestore()
        logger.info('Демон  "check in restore" запущен')
        check_in_restore.run()
        check_in_restore.check_invalid_uniqueid()
        logger.info("Проверка завершена")
    except Exception:
        error_message = traceback.format_exc()
        logger.error(error_message)



if __name__ == "__main__":
    main()
