## sasl.py
# SASL authentication support. Currently we only support PLAIN authentication.
import asyncio
import base64
from functools import partial

try:
    import puresasl
    import puresasl.client
except ImportError:
    puresasl = None

from . import cap

__all__ = ["SASLSupport"]


RESPONSE_LIMIT = 400
EMPTY_MESSAGE = "+"
ABORT_MESSAGE = "*"


class SASLSupport(cap.CapabilityNegotiationSupport):
    """SASL authentication support. Currently limited to the PLAIN mechanism."""

    SASL_TIMEOUT = 10

    ## Internal overrides.

    def __init__(
        self,
        *args,
        sasl_identity="",
        sasl_username=None,
        sasl_password=None,
        sasl_mechanism=None,
        **kwargs,
    ):
        super().__init__(*args, **kwargs)
        self.sasl_identity = sasl_identity
        self.sasl_username = sasl_username
        self.sasl_password = sasl_password
        self.sasl_mechanism = sasl_mechanism

    def _reset_attributes(self):
        super()._reset_attributes()
        self._sasl_client = None
        self._sasl_timer = None
        self._sasl_challenge = b""
        self._sasl_mechanisms = None

    ## SASL functionality.

    async def _sasl_start(self, mechanism):
        """Initiate SASL authentication."""
        # The rest will be handled in on_raw_authenticate()/_sasl_respond().
        await self.rawmsg("AUTHENTICATE", mechanism)
        # create a partial, required for our callback to get the kwarg
        _sasl_partial = partial(self._sasl_abort, timeout=True)
        self._sasl_timer = asyncio.get_event_loop().call_later(
            self.SASL_TIMEOUT, _sasl_partial
        )

    async def _sasl_abort(self, timeout=False):
        """Abort SASL authentication."""
        if timeout:
            self.logger.error("SASL authentication timed out: aborting.")
        else:
            self.logger.error("SASL authentication aborted.")

        if self._sasl_timer:
            self._sasl_timer.cancel()

            self._sasl_timer = None

        # We're done here.
        await self.rawmsg("AUTHENTICATE", ABORT_MESSAGE)
        await self._capability_negotiated("sasl")

    async def _sasl_end(self):
        """Finalize SASL authentication."""
        if self._sasl_timer:
            self._sasl_timer.cancel()
            self._sasl_timer = None
        await self._capability_negotiated("sasl")

    async def _sasl_respond(self):
        """Respond to SASL challenge with response."""
        # Formulate a response.
        if self._sasl_client:
            try:
                response = self._sasl_client.process(self._sasl_challenge)
            except puresasl.SASLError:
                response = None

            if response is None:
                self.logger.warning(
                    "SASL challenge processing failed: aborting SASL authentication."
                )
                await self._sasl_abort()
        else:
            response = b""

        response = base64.b64encode(response).decode(self.encoding)
        to_send = len(response)
        self._sasl_challenge = b""

        # Send response in chunks.
        while to_send > 0:
            await self.rawmsg("AUTHENTICATE", response[:RESPONSE_LIMIT])
            response = response[RESPONSE_LIMIT:]
            to_send -= RESPONSE_LIMIT

        # If our message fit exactly in SASL_RESPOSE_LIMIT-byte chunks, send an empty message to indicate we're done.
        if to_send == 0:
            await self.rawmsg("AUTHENTICATE", EMPTY_MESSAGE)

    ## Capability callbacks.

    async def on_capability_sasl_available(self, value):
        """Check whether or not SASL is available."""
        if value:
            self._sasl_mechanisms = value.upper().split(",")
        else:
            self._sasl_mechanisms = None

        if self.sasl_mechanism == "EXTERNAL" or (
            self.sasl_username and self.sasl_password
        ):
            if self.sasl_mechanism == "EXTERNAL" or puresasl:
                return True
            self.logger.warning(
                "SASL credentials set but puresasl module not found: not initiating SASL authentication."
            )
        return False

    async def on_capability_sasl_enabled(self):
        """Start SASL authentication."""
        if self.sasl_mechanism:
            if (
                self._sasl_mechanisms
                and self.sasl_mechanism not in self._sasl_mechanisms
            ):
                self.logger.warning(
                    "Requested SASL mechanism is not in server mechanism list: aborting SASL authentication."
                )
                return cap.failed
            mechanisms = [self.sasl_mechanism]
        else:
            mechanisms = self._sasl_mechanisms or ["PLAIN"]

        if mechanisms == ["EXTERNAL"]:
            mechanism = "EXTERNAL"
        else:
            self._sasl_client = puresasl.client.SASLClient(
                self.connection.hostname,
                "irc",
                username=self.sasl_username,
                password=self.sasl_password,
                identity=self.sasl_identity,
            )

            try:
                self._sasl_client.choose_mechanism(mechanisms, allow_anonymous=False)
            except puresasl.SASLError:
                self.logger.exception(
                    "SASL mechanism choice failed: aborting SASL authentication."
                )
                return cap.FAILED
            mechanism = self._sasl_client.mechanism.upper()

        # Initialize SASL.
        await self._sasl_start(mechanism)
        # Tell caller we need more time, and to not end capability negotiation just yet.
        return cap.NEGOTIATING

    ## Message handlers.

    async def on_raw_authenticate(self, message):
        """Received part of the authentication challenge."""
        # Cancel timeout timer.
        if self._sasl_timer:
            self._sasl_timer.cancel()
            self._sasl_timer = None

        # Add response data.
        response = " ".join(message.params)
        if response != EMPTY_MESSAGE:
            self._sasl_challenge += base64.b64decode(response)

        # If the response ain't exactly SASL_RESPONSE_LIMIT bytes long, it's the end. Process.
        if len(response) % RESPONSE_LIMIT > 0:
            await self._sasl_respond()
        else:
            # Response not done yet. Restart timer.
            self._sasl_timer = asyncio.get_event_loop().call_later(
                self.SASL_TIMEOUT, self._sasl_abort(timeout=True)
            )

    on_raw_900 = (
        cap.CapabilityNegotiationSupport._ignored
    )  # You are now logged in as...

    async def on_raw_903(self, message):
        """SASL authentication successful."""
        await self._sasl_end()

    async def on_raw_904(self, message):
        """Invalid mechanism or authentication failed. Abort SASL."""
        await self._sasl_abort()

    async def on_raw_905(self, message):
        """Authentication failed. Abort SASL."""
        await self._sasl_abort()

    on_raw_906 = (
        cap.CapabilityNegotiationSupport._ignored
    )  # Completed registration while authenticating/registration aborted.
    on_raw_907 = (
        cap.CapabilityNegotiationSupport._ignored
    )  # Already authenticated over SASL.
