from dataclasses import dataclass
from ipaddress import ip_address
from pathlib import Path
from typing import Optional, List, Dict, Tuple, Union

import nassl

from sslyze.errors import TlsHandshakeFailed
from sslyze.plugins.certificate_info._cert_chain_analyzer import (
    CertificateDeploymentAnalyzer,
    CertificateDeploymentAnalysisResult,
)
from cryptography.x509 import DNSName, IPAddress
from sslyze.plugins.certificate_info._cli_connector import _CertificateInfoCliConnector
from sslyze.plugins.certificate_info._get_cert_chain import get_certificate_chain, ArgumentsToGetCertificateChain
from sslyze.plugins.certificate_info.trust_stores.trust_store import TrustStore
from sslyze.plugins.certificate_info.trust_stores.trust_store_repository import TrustStoresRepository
from sslyze.plugins.plugin_base import (
    ScanCommandImplementation,
    ScanJob,
    ScanCommandResult,
    ScanCommandExtraArgument,
    ScanJobResult,
)
from sslyze.server_connectivity import ServerConnectivityInfo, TlsVersionEnum


@dataclass(frozen=True)
class CertificateInfoExtraArgument(ScanCommandExtraArgument):
    """Additional configuration for running the certificate_info scan command.

    Attributes:
        custom_ca_file: The path to a custom trust store file to use for certificate validation. The file should contain
            PEM-formatted root certificates.
    """

    custom_ca_file: Path

    def __post_init__(self) -> None:
        if not self.custom_ca_file.is_file():
            raise ValueError(f'Could not open supplied CA file at "{self.custom_ca_file}"')


@dataclass(frozen=True)
class CertificateInfoScanResult(ScanCommandResult):
    """The result of retrieving and analyzing a server's certificates to verify their validity.

    Attributes:
        hostname_used_for_server_name_indication: The hostname sent by SSLyze as the Server Name Indication (SNI)
            extension.
        certificate_deployments: The list of certificate deployments detected by SSLyze with SNI enabled. Most servers
            only deploy one leaf certificate, but some websites (such as Facebook) return different leaf certificates
            depending on the client, as a way to maximize compatibility with older clients/devices.
        certificate_deployment_with_sni_disabled: The certificate deployment retrieved by SSLyze with SNI disabled.
            This allows discovery of default certificates that servers return to clients that don't support or provide
            SNI. Will be ``None`` if the non-SNI connection failed.
            Most clients always enable SNI when connecting to a server so this certificate deployment can be ignored.
    """

    hostname_used_for_server_name_indication: str
    certificate_deployments: List[CertificateDeploymentAnalysisResult]
    certificate_deployment_with_sni_disabled: Optional[CertificateDeploymentAnalysisResult]


_ListofPemCertificatesAndOptionalOcspResponse = Tuple[List[str], Optional[nassl._nassl.OCSP_RESPONSE]]


class CertificateInfoImplementation(ScanCommandImplementation[CertificateInfoScanResult, None]):
    """Retrieve and analyze a server's certificate(s) to verify its validity."""

    cli_connector_cls = _CertificateInfoCliConnector

    @classmethod
    def scan_jobs_for_scan_command(
        cls, server_info: ServerConnectivityInfo, extra_arguments: Optional[CertificateInfoExtraArgument] = None
    ) -> List[ScanJob]:
        custom_ca_file = extra_arguments.custom_ca_file if extra_arguments else None

        # Try to retrieve different certificates from the server by having SSLyze's TLS handshake look like different
        # kinds of clients
        call_arguments: List[ArgumentsToGetCertificateChain] = []
        if server_info.tls_probing_result.highest_tls_version_supported.value >= TlsVersionEnum.TLS_1_3.value:
            # Get the default certificate chain sent to clients using TLS 1.3
            call_arguments.append((server_info, custom_ca_file, TlsVersionEnum.TLS_1_3, None, True))

            # Get the other certificate chains sent to clients using TLS 1.2 that support
            #  or don't support RSA for authentication
            call_arguments.append((server_info, custom_ca_file, TlsVersionEnum.TLS_1_2, "aRSA", True))
            call_arguments.append((server_info, custom_ca_file, TlsVersionEnum.TLS_1_2, "ALL:-aRSA", True))
        else:
            # Get the certificate chains sent to clients that support or don't support RSA for authentication
            call_arguments.append((server_info, custom_ca_file, None, None, True))
            call_arguments.append((server_info, custom_ca_file, None, "aRSA", True))
            call_arguments.append((server_info, custom_ca_file, None, "ALL:-aRSA", True))

        # Additionally, get the default certificate chain sent to clients without SNI
        call_arguments.append((server_info, custom_ca_file, None, None, False))

        # The custom_ca_file is not needed by get_certificate_chain() but we have to pass it anyway so we can eventually
        # use it in result_for_completed_scan_jobs()
        scan_jobs = [
            ScanJob(function_to_call=get_certificate_chain, function_arguments=call_arg) for call_arg in call_arguments
        ]
        return scan_jobs

    _EXPECTED_SCAN_JOB_RESULTS_COUNT = 4

    @classmethod
    def result_for_completed_scan_jobs(
        cls, server_info: ServerConnectivityInfo, scan_job_results: List[ScanJobResult]
    ) -> CertificateInfoScanResult:
        if len(scan_job_results) != cls._EXPECTED_SCAN_JOB_RESULTS_COUNT:
            raise RuntimeError(f"Unexpected number of scan jobs received: {scan_job_results}")

        # Process the results
        # Leaf certificate => certificate chain, OCSP response, was_sni_used
        all_configured_certificate_chains: Dict[str, _ListofPemCertificatesAndOptionalOcspResponse] = {}
        all_handshake_failed_exceptions: List[Exception] = []
        custom_ca_file = None
        non_sni_chain_and_ocsp_response: Optional[_ListofPemCertificatesAndOptionalOcspResponse] = None
        server_did_not_send_a_certificate = False
        for completed_job in scan_job_results:
            try:
                received_chain_as_pem, ocsp_response, custom_ca_file, was_sni_used = completed_job.get_result()
                assert received_chain_as_pem, "Should never happen as we always retrieve a certificate chain"

            except TlsHandshakeFailed as exc:
                # Can happen when trying to connect with specific cipher suites (such as RSA or non-RSA)
                # or when connectivity is bad
                all_handshake_failed_exceptions.append(exc)

            except nassl._nassl.OpenSSLError as exc:
                if "unrecognized name" in exc.args[0]:
                    # Can happen when trying to connect without SNI
                    all_handshake_failed_exceptions.append(exc)
                else:
                    raise

            except ValueError as exc:
                if "peer did not present a certificate" in exc.args[0]:
                    # Some servers don't send back a certificate chain at all (ie. ANON cipher suites)
                    # https://github.com/nabla-c0d3/sslyze/issues/705
                    server_did_not_send_a_certificate = True
                    all_handshake_failed_exceptions.append(exc)
                else:
                    raise

            else:
                if not was_sni_used:
                    # This is the SNI-disabled deployment, we handle it separately
                    non_sni_chain_and_ocsp_response = received_chain_as_pem, ocsp_response
                else:
                    # Only keep certificate chains that are different
                    all_configured_certificate_chains[received_chain_as_pem[0]] = received_chain_as_pem, ocsp_response

        if len(all_handshake_failed_exceptions) == cls._EXPECTED_SCAN_JOB_RESULTS_COUNT:
            # All TLS handshakes failed; either:
            if server_did_not_send_a_certificate:
                # The server is only supporting ANON cipher suites (ie. no certificate)
                pass
            else:
                # Or we have bad connectivity to the server: re-raise one of the handshake exceptions
                raise all_handshake_failed_exceptions[0]

        if not all_configured_certificate_chains and not server_did_not_send_a_certificate:
            raise ValueError("Should never happen")

        # Then validate each certificate/chain deployment
        all_trust_stores = TrustStoresRepository.get_default().get_all_stores()
        if custom_ca_file:
            all_trust_stores.append(TrustStore(custom_ca_file, "Supplied CA file", "N/A"))

        # Figure out the server's subject (ie. hostname or IP) to use for certificate validation
        subject_to_use_for_validation: Union[IPAddress, DNSName]
        if server_info.network_configuration.tls_server_name_indication == server_info.server_location.ip_address:
            # We don't know the server's hostname as only an IP address was supplied to sslyze
            if not server_info.server_location.ip_address:
                raise ValueError("Should never happen as it means the tls_server_name_indication would be None")

            subject_to_use_for_validation = IPAddress(ip_address(server_info.server_location.ip_address))
        else:
            subject_to_use_for_validation = DNSName(server_info.network_configuration.tls_server_name_indication)

        # Check the SNI-disabled certificate deployment
        sni_disabled_analysis_result: Optional[CertificateDeploymentAnalysisResult]
        if not non_sni_chain_and_ocsp_response:
            sni_disabled_analysis_result = None
        else:
            received_chain_as_pem, ocsp_response = non_sni_chain_and_ocsp_response
            sni_disabled_deployment_analyzer = CertificateDeploymentAnalyzer(
                server_subject=subject_to_use_for_validation,
                server_certificate_chain_as_pem=received_chain_as_pem,
                server_ocsp_response=ocsp_response,
                trust_stores_for_validation=all_trust_stores,
            )
            sni_disabled_analysis_result = sni_disabled_deployment_analyzer.perform()

        # Check all SNI-enabled certificate deployments
        all_sni_enabled_deployments: List[CertificateDeploymentAnalysisResult] = []
        for received_chain_as_pem, ocsp_response in all_configured_certificate_chains.values():
            deployment_analyzer = CertificateDeploymentAnalyzer(
                server_subject=subject_to_use_for_validation,
                server_certificate_chain_as_pem=received_chain_as_pem,
                server_ocsp_response=ocsp_response,
                trust_stores_for_validation=all_trust_stores,
            )
            analysis_result = deployment_analyzer.perform()
            all_sni_enabled_deployments.append(analysis_result)

        # All done
        return CertificateInfoScanResult(
            hostname_used_for_server_name_indication=server_info.network_configuration.tls_server_name_indication,
            certificate_deployments=all_sni_enabled_deployments,
            certificate_deployment_with_sni_disabled=sni_disabled_analysis_result,
        )
