import requests
from http import HTTPStatus
from requests.auth import HTTPBasicAuth, HTTPDigestAuth
from glpwnme.exploits.logger import Log
from glpwnme.exploits.exceptions import BadCredentialsException

class AnyAuthSession(requests.Session):
    def __init__(self, username="", password=""):
        super().__init__()
        self.username = username
        self.password = password
        self._negotiated = False  # have we picked a scheme yet?

    def __assert_authentication_worked(self, res):
        """
        Check that server authentication worked
        """
        if(self._negotiated and self.auth is not None and res.status_code == HTTPStatus.UNAUTHORIZED):
            raise BadCredentialsException("Server authentication failed")

    def request(self, method, url, **kwargs):
        # If we haven't negotiated auth yet for this session, do it now
        if not self._negotiated and self.auth is None:
            self._negotiate(method, url, **kwargs)
        return super().request(method, url, **kwargs)

    def _negotiate(self, method, url, **kwargs):
        probe_kwargs = {k: v for k, v in kwargs.items() if k != "auth"}
        resp = super().request(method, url, **probe_kwargs)

        is_auth_allowed = "WWW-Authenticate" in resp.headers.keys()
        if resp.status_code != HTTPStatus.UNAUTHORIZED:
            # no auth needed, or already authenticated (e.g. via cookies)
            self._negotiated = True
            return

        scheme = self._detect_scheme(resp)
        self.auth = self._build_auth(scheme)
        self._negotiated = True
        retry_resp = super().request(method, url, **probe_kwargs)
        self.__assert_authentication_worked(retry_resp)

    def _detect_scheme(self, response):
        auth_header = response.headers.get("WWW-Authenticate", "")
        message = f"Server requires authentication [gold3]{auth_header}[/gold3]"
        if not self.username:
            message += ", use [cyan]-su[/cyan] and [cyan]-sp[/cyan] for server credentials"
        Log.log(message)
        schemes = [s.strip().split(" ")[0].lower() for s in auth_header.split(",")]
        for pref in ("negotiate", "digest", "ntlm", "basic"):
            if pref in schemes:
                return pref
        raise ValueError(f"No supported auth scheme in: {auth_header!r}")

    def _build_auth(self, scheme):
        if scheme == "basic":
            return HTTPBasicAuth(self.username, self.password)
        if scheme == "digest":
            return HTTPDigestAuth(self.username, self.password)
        if scheme == "ntlm":
            try:
                from requests_ntlm import HttpNtlmAuth
                return HttpNtlmAuth(self.username, self.password)
            except ImportError:
                raise ImportError("Missing dependency: requests-ntlm. "
                                  "Install with: pip install requests-ntlm")
        if scheme == "negotiate":
            try:
                from requests_kerberos import HTTPKerberosAuth
                return HTTPKerberosAuth()
            except ImportError:
                raise ImportError("Missing dependency: requests-kerberos. "
                                  "Install with: pip install requests-kerberos")
        raise ValueError(f"Unsupported scheme: {scheme}")
