import json
import logging
import time
from typing import Dict, Optional

from cacheout import CacheManager
from msal import PublicClientApplication
from prometheus_client import Counter
from . import util
from ._version import __version__
from .openvpn import ManagementInterface
from .util import errors
from .util.thread_pool import ThreadPoolExecutorStackTraced

openvpn_auth_azure_ad_events = Counter(
    "openvpn_auth_azure_ad_events", "track events", ["event"]
)
openvpn_auth_azure_ad_auth_total = Counter(
    "openvpn_auth_azure_ad_auth_total", "auth total", ["flow"]
)
openvpn_auth_azure_ad_auth_succeeded = Counter(
    "openvpn_auth_azure_ad_auth_succeeded", "auth succeeded", ["flow"]
)
openvpn_auth_azure_ad_auth_failures = Counter(
    "openvpn_auth_azure_ad_auth_failures", "auth failures", ["flow"]
)


log = logging.getLogger(__name__)


class AADAuthenticatorFlows:
    USER_PASSWORD = "username_password"
    DEVICE_TOKEN = "device_token"
    AUTH_TOKEN = "auth_token"


class AADAuthenticator(object):
    token_scopes = ["User.ReadBasic.All"]

    def __init__(
        self,
        app: PublicClientApplication,
        graph_endpoint: str,
        authenticators: str,
        verify_certificate_claim: bool,
        auth_token: bool,
        auth_token_lifetime: int,
        remember_user: bool,
        threads: int,
        host: str = None,
        port: int = None,
        socket: str = None,
        password: str = None,
    ):
        self._app = app
        self._graph_endpoint = graph_endpoint
        self._authenticators = [s.strip() for s in authenticators.split(",")]
        self._openvpn = ManagementInterface(host, port, socket, password)
        self._openvpn.connect()
        self._states = CacheManager(
            {
                "challenge": {"maxsize": 256, "ttl": 600},
                "authenticated": {"maxsize": 256, "ttl": 0},
                "auth_token": {"maxsize": 256, "ttl": 86400},
            }
        )

        self._verify_certificate_claim = verify_certificate_claim
        self._auth_token_enabled = auth_token
        self._auth_token_lifetime = auth_token_lifetime
        self._remember_user_enabled = remember_user
        self._thread_pool = ThreadPoolExecutorStackTraced(max_workers=threads)

    def run(self) -> None:
        log.info("Running openvpn-auth-azure-ad %s" % __version__)
        try:
            while True:
                message = self._openvpn.receive()
                if not message:
                    log.error("Connection to OpenVPN closed. Reconnecting...")
                    self._openvpn.connect(True)
                    continue

                if message.startswith("ERROR:"):
                    log.error(message)
                    continue

                if message.startswith(">CLIENT:DISCONNECT"):
                    self._thread_pool.submit(self.client_disconnect, message)

                elif message.startswith(">CLIENT:CONNECT"):
                    self._thread_pool.submit(self.client_connect, message)

                elif message.startswith(">CLIENT:REAUTH"):
                    self._thread_pool.submit(self.client_reauth, message)

                self._states["challenge"].delete_expired()
                self._states["auth_token"].delete_expired()
        except KeyboardInterrupt:
            pass
        except Exception as e:
            log.error(str(e), exc_info=e)

    def send_authentication_success(self, client: dict) -> None:
        self.log_info(client, "authentication succeeded")
        if client["auth_token"] is not None:
            self._openvpn.send_command(
                'client-auth %s %s\npush "auth-token %s"\nEND'
                % (client["cid"], client["kid"], client["auth_token"])
            )
        else:
            self._openvpn.send_command(
                "client-auth-nt %s %s" % (client["cid"], client["kid"])
            )

    def send_authentication_challenge(self, client: Dict, client_message: str) -> None:
        self.log_debug(client, "authentication challenge: %s" % client_message)

        client_challenge = util.format_client_challenge(client, client_message)
        self._openvpn.send_command(
            'client-deny %s %s "%s" "%s"'
            % (client["cid"], client["kid"], "client_challenge", client_challenge)
        )

    def send_authentication_error(
        self, client: dict, message: str, client_message: Optional[str]
    ) -> None:
        if client_message is None:
            self._openvpn.send_command(
                'client-deny %s %s "%s"' % (client["cid"], client["kid"], message)
            )
        else:
            client_challenge = util.format_client_challenge(client, client_message)
            self._openvpn.send_command(
                'client-deny %s %s "%s" "%s"'
                % (client["cid"], client["kid"], message, client_challenge)
            )

    def setup_auth_token(self, client) -> None:
        if self._auth_token_enabled:
            client["auth_token"] = util.generated_id()
            self._states["auth_token"].set(
                client["auth_token"], client["state_id"], self._auth_token_lifetime
            )

    def verify_client_certificate(self, client, result) -> bool:
        if not self._verify_certificate_claim:
            return True

        if "common_name" not in client["env"]:
            return False

        if "id_token_claims" not in result:
            self.log_warn(client, "Could not get id_token")
            return False

        if self._verify_certificate_claim not in result["id_token_claims"]:
            self.log_warn(
                client,
                "claim %s does not exist in id_token"
                % (self._verify_certificate_claim,),
            )
            return False

        return (
            client["env"]["common_name"]
            == result["id_token_claims"][self._verify_certificate_claim]
        )

    def handle_reauth(self, client: Dict) -> Optional[dict]:
        if not self._auth_token_enabled:
            return None

        openvpn_auth_azure_ad_auth_total.labels(AADAuthenticatorFlows.AUTH_TOKEN).inc()

        # Get Auth-Token from Request
        client["auth_token"] = util.get_auth_token(client)
        if client["auth_token"] is None:
            return None

        # Receive state_id by auth-token
        client["state_id"] = self._states["auth_token"].get(client["auth_token"])
        if client["state_id"] is None:
            return None

        # Get authenticated state
        authenticated_state = self._states["authenticated"].get(client["state_id"])
        if authenticated_state is None:
            return None

        # check is authenticated state has required objects
        if "result" not in authenticated_state or "client" not in authenticated_state:
            return None

        self.log_info(client, "Authenticate using auth-token flow")

        # recheck if connected client is the same as inside the state.
        if (
            client["auth_token"] != authenticated_state["client"]["auth_token"]
            or client["env"]["common_name"]
            != authenticated_state["client"]["env"]["common_name"]
            or client["env"]["username"]
            != authenticated_state["client"]["env"]["username"]
        ):
            return None

        # check if user is active inside azure ad by using a refresh token
        result = self._app.acquire_token_by_refresh_token(
            authenticated_state["result"]["refresh_token"], scopes=self.token_scopes
        )

        return result

    def handle_response_challenge(self, client: Dict) -> Optional[dict]:
        if not client["env"]["password"].startswith("CRV1::"):
            return None

        # get state id from client
        state_id = util.get_state_id(client)
        if state_id is None:
            return None

        # get state from challenge state and check if exists
        state = self._states["challenge"].get(state_id)
        if state is None:
            return None

        # delete used challenge state
        self._states["challenge"].delete(state_id)

        # Check if required properties exist in state
        if "flow" not in state:
            return None

        self.log_info(client, "Continue to authenticate using device token flow")
        result = self.device_token_finish(state)

        return result

    def client_connect(self, data: str) -> None:
        client = AADAuthenticator.parse_client_data(data)
        self.log_info(client, "Received client connect")
        openvpn_auth_azure_ad_events.labels("connect").inc()
        self.authenticate_client(client)

    def client_disconnect(self, data: str) -> None:
        client = AADAuthenticator.parse_client_data(data)
        self.log_info(client, "Received client disconnect event")
        openvpn_auth_azure_ad_events.labels("disconnect").inc()

    def client_reauth(self, data: str) -> None:
        client = AADAuthenticator.parse_client_data(data)
        self.log_info(client, "Received client re auth event")
        openvpn_auth_azure_ad_events.labels("reauth").inc()
        self.authenticate_client(client)

    def authenticate_client(self, client: dict) -> None:
        result = {}

        if client["reason"] == "reauth":
            # allow clients to bypass azure ad authentication by reauthenticate via auth-token
            result = self.handle_reauth(client)
            if result is not None:
                if util.is_authenticated(result):
                    if self.verify_client_certificate(client, result):
                        self.log_info(
                            client, "client certificate does not match Azure AD user."
                        )
                        self.send_authentication_error(
                            client, "client_certificate_not_matched", None
                        )
                        return None

                    openvpn_auth_azure_ad_auth_succeeded.labels(
                        AADAuthenticatorFlows.AUTH_TOKEN
                    ).inc()
                    self.log_info(client, "auth-token flow succeeded")
                    self.send_authentication_success(client)
                    return

        if AADAuthenticatorFlows.DEVICE_TOKEN in self._authenticators:
            result = self.handle_response_challenge(client)
            if result is not None:
                client["state_id"] = util.get_state_id(client)
                if util.is_authenticated(result):
                    if not self.verify_client_certificate(client, result):
                        self.log_info(
                            client, "client certificate does not match Azure AD user."
                        )
                        self.send_authentication_error(
                            client, "client_certificate_not_matched", None
                        )
                        return None

                    self.setup_auth_token(client)
                    self._states["authenticated"].set(
                        client["state_id"], {"client": client, "result": result}
                    )
                    openvpn_auth_azure_ad_auth_succeeded.labels(
                        AADAuthenticatorFlows.DEVICE_TOKEN
                    ).inc()
                    self.log_info(client, "device token flow succeeded")
                    self.send_authentication_success(client)
                    return
                else:
                    openvpn_auth_azure_ad_auth_failures.labels(
                        AADAuthenticatorFlows.DEVICE_TOKEN
                    ).inc()
                    if 70016 in result.get("error_codes", []):
                        error = "device token flow errored: no user action"
                        self.log_info(client, error)
                        self.send_authentication_error(client, error, None)
                    else:
                        self.log_info(
                            client,
                            "device token flow errored: %s "
                            % util.format_error(result),
                        )
                        self.send_authentication_error(
                            client, util.format_error(result), None
                        )
                    return

        client["state_id"] = util.generated_id()
        if self._remember_user_enabled:
            accounts = self._app.get_accounts()
            if "common_name" not in client["env"] and accounts:
                result = self._app.acquire_token_silent(
                    self.token_scopes, account=client["env"]["common_name"]
                )

            if util.is_authenticated(result):
                self.setup_auth_token(client)
                self._states["authenticated"].set(
                    client["state_id"], {"client": client, "result": result}
                )
                self.send_authentication_success(client)
                return

        if AADAuthenticatorFlows.USER_PASSWORD in self._authenticators:
            self.log_debug(client, "Authenticate using username/password flow")
            openvpn_auth_azure_ad_auth_total.labels(
                AADAuthenticatorFlows.USER_PASSWORD
            ).inc()
            result = self._app.acquire_token_by_username_password(
                client["env"]["username"],
                client["env"]["password"],
                scopes=self.token_scopes,
            )

            if util.is_authenticated(result):
                if not self.verify_client_certificate(client, result):
                    self.log_info(
                        client, "client certificate does not match Azure AD user."
                    )
                    self.send_authentication_error(
                        client, "client_certificate_not_matched", None
                    )
                    return None
            else:
                openvpn_auth_azure_ad_auth_failures.labels(
                    AADAuthenticatorFlows.USER_PASSWORD
                ).inc()
                if 65001 in result.get("error_codes", []):
                    # AAD requires user consent for U/P flow
                    error = (
                        "Get consent first: %s"
                        % self._app.get_authorization_request_url(self.token_scopes)
                    )
                    self.send_authentication_error(client, "consent_needed", error)
                else:
                    error = util.format_error(result)
                    self.send_authentication_error(client, error, None)

                self.log_info(client, "password flow errored: %s" % (error,))
                return

            openvpn_auth_azure_ad_auth_succeeded.labels(
                AADAuthenticatorFlows.USER_PASSWORD
            ).inc()
            self.log_info(client, "password flow succeeded")
            # Do not send successful auth command, because device token can be enabled.

        if AADAuthenticatorFlows.DEVICE_TOKEN in self._authenticators:
            self.log_info(client, "Start to authenticate using device token flow")
            flow = self.device_auth_start()
            message = flow["message"] + " Then press OK here. No input required here."

            self.log_debug(client, "Save state as %s" % client["state_id"])
            self._states["challenge"].set(client["state_id"], {"flow": flow})

            openvpn_auth_azure_ad_auth_total.labels(
                AADAuthenticatorFlows.DEVICE_TOKEN
            ).inc()
            self.send_authentication_challenge(client, message)
            return

        self.setup_auth_token(client)
        self._states["authenticated"].set(
            client["state_id"], {"client": client, "result": result}
        )
        self.send_authentication_success(client)
        return

    def device_auth_start(self) -> dict:
        flow = self._app.initiate_device_flow(scopes=self.token_scopes)

        if "user_code" not in flow:
            raise ValueError(
                "Fail to create device flow. Err: %s" % json.dumps(flow, indent=4)
            )

        return flow

    def device_token_finish(self, state: dict) -> dict:
        # Block authentication at least for 30 seconds
        state["flow"]["expires_in"] = 30
        state["flow"]["expires_at"] = time.time() + state["flow"]["expires_in"]

        return self._app.acquire_token_by_device_flow(state["flow"])

    @staticmethod
    def parse_client_data(data: str) -> dict:
        client = {
            "env": {},
            "reason": None,
            "cid": None,
            "kid": None,
            "state_id": None,
        }

        for line in data.splitlines():
            try:
                if line.startswith(">CLIENT:CONNECT") or line.startswith(
                    ">CLIENT:REAUTH"
                ):
                    client_info = line.split(",")
                    client["reason"] = client_info[0].replace(">CLIENT:", "").lower()
                    client["cid"] = client_info[1]
                    client["kid"] = client_info[2]
                elif line.startswith(">CLIENT:DISCONNECT"):
                    client_info = line.split(",")
                    client["reason"] = client_info[0].replace(">CLIENT:", "").lower()
                    client["cid"] = client_info[1]
                elif line.startswith(">CLIENT:ENV,"):
                    client_env = line.split(",")[1].split("=")
                    client["env"][client_env[0]] = (
                        client_env[1] if len(client_env) == 2 else ""
                    )
                else:
                    raise errors.ParseError("Can't parse line: %s" % (line,))
            except Exception:
                raise errors.ParseError("Can't parse line: %s" % (line,))

        return client

    def log_warn(self, client: dict, message: str) -> None:
        prefix = self.format_log_prefix(client)

        log.warning("[%s]: %s" % (prefix, message))

    def log_info(self, client: dict, message: str) -> None:
        prefix = self.format_log_prefix(client)

        log.info("[%s]: %s" % (prefix, message))

    def log_debug(self, client: dict, message: str) -> None:
        prefix = self.format_log_prefix(client)

        log.debug("[%s]: %s" % (prefix, message))

    def format_log_prefix(self, client: dict) -> str:
        prefix = "cid: %s" % (client["cid"],)
        if "common_name" in client["env"]:
            prefix += " | %s" % client["env"]["common_name"]

        return prefix
