# Copyright © 2024 Contrast Security, Inc.
# See https://www.contrastsecurity.com/enduser-terms-0317a for more details.
from typing import Optional

from contrast.agent import agent_state
from contrast.agent.middlewares.route_coverage import common

from contrast_vendor import structlog as logging

logger = logging.getLogger("contrast")


class RoutesMixin:
    """
    Class for route coverage work.
    Route coverage is assess only.
    """

    def handle_routes(self, context, request) -> None:
        """
        Method that should run for all middlewares at the end of request processing.
        """
        if context and request:
            self.update_route_information(context, request)
            self.append_route_to_findings(context)

    def update_route_information(self, context, request) -> None:
        """
        Store the information for this request's route on RequestContext.
        """
        logger.debug("updating observed route information")

        signature = self._get_signature(context, request)

        context.observed_route.signature = signature
        context.observed_route.url = context.request.get_normalized_uri()
        context.observed_route.verb = request.method
        logger.debug(
            "observed route",
            signature=context.observed_route.signature,
            url=context.observed_route.url,
            verb=context.observed_route.verb,
        )

    def _get_signature(self, context, request) -> Optional[str]:
        """
        The logic here is somewhat convoluted and depends on state being set in a few
        other locations. Much of this will change as we move route discovery to a
        background thread. There are a few different strategies we might use to obtain a
        route signature. In order of preference:

        1.  Obtain a signature at some point during the request; this string is set on
            `context.view_func_str`. For supported frameworks, this is accomplished with
            framework-specific patches. One exception to this is aiohttp, where it
            happens directly in the middleware. It is essential that for a particular
            route this signature exactly matches the one found during route discovery.
            In the future, all frameworks will take this approach.

        2.  Obtain a reference to the view function at some point during the request.
            This reference is set on `context.view_func`.
            2a. Use this reference to construct a key to the discovered_routes dict.
                Look up the previously discovered route, and use its signature. An
                unfortunate consequence of this approach is that route discovery must be
                completed first, so we can't do this anymore if we want to move route
                discovery to a background thread.
            2b. If we don't find the corresponding discovered route, build a reasonable
                signature using the view function we have. This is a "dynamic route".

        3.  If neither context.view_func_str nor context.view_func are set, we didn't
            get any framework-specific information about the view function. In this case
            we use the request's normalized URI as the signature. We expect to hit this
            case in pure WSGI / ASGI.

        Framework-specific route observation code should set either view_func_str, which
        takes priority, or view_func. It doesn't make sense to set both, since only
        view_func_str would be used in that case.
        """
        if context.view_func_str:
            return context.view_func_str

        # TODO: PYT-3317 remove this entire block
        if context.view_func:
            discovered_routes = agent_state.get_discovered_routes()
            key = common.build_key(str(id(context.view_func)), request.method)
            if route := discovered_routes.get(key):
                return route.signature
            return common.build_route(context.view_func.__name__, context.view_func)

        logger.debug(
            "Did not find a view function signature for the current request. "
            "Falling back on normalized URI."
        )
        return context.request.get_normalized_uri()

    def append_route_to_findings(self, context):
        """
        Append the observed route to any existing findings. We can't necessarily do this
        at finding creation time, because we might not have route info yet.
        """
        if not context.observed_route:
            logger.debug("No observed route to append to findings")
            return

        if not context.findings:
            logger.debug("No findings to append route to")
            return

        for finding in context.findings:
            if not finding.routes:
                logger.debug(
                    "Appending route %s:%s to %s",
                    context.observed_route.verb,
                    context.observed_route.url,
                    finding.rule_id,
                )
                finding.routes.append(context.observed_route)

            finding.set_version()
