apimrequests.py

"""Module for making requests to Azure API Management endpoints with consistent logging and output formatting.

Autoreload: reload_conditional
Method edits reload. Recreate clients after constructor changes; close existing sessions normally.
"""

import json
import time
import warnings
from typing import Any

import requests
import urllib3

# APIM Samples imports
from apimtypes import HTTP_VERB, SLEEP_TIME_BETWEEN_REQUESTS_MS, SUBSCRIPTION_KEY_PARAMETER_NAME, HttpStatusCode
from console import BOLD_G, BOLD_R, RESET, print_debug, print_error, print_info, print_message, print_ok, print_val
from json_utils import pretty_json
from logging_config import is_debug_enabled, should_unmask_sensitive_headers

_SENSITIVE_HEADER_NAMES = frozenset(k.lower() for k in (
    'api-key', 'Ocp-Apim-Subscription-Key', 'Authorization', 'x-api-key',
    'Proxy-Authorization', 'Cookie', 'Set-Cookie', 'x-auth-token',
    'x-access-token', 'x-ms-token-aad-access-token', 'x-ms-token-aad-refresh-token',
))


def _redact_headers(headers: dict | None) -> dict | None:
    """Copy headers for output, masking secrets unless explicitly opted out."""
    if not headers:
        return headers
    if should_unmask_sensitive_headers():
        return dict(headers)

    return {
        k: ('***REDACTED***' if k.lower() in _SENSITIVE_HEADER_NAMES else v)
        for k, v in headers.items()
    }


# ------------------------------
#    CLASSES
# ------------------------------


class ApimRequests:
    """
    Methods for making requests to the Azure API Management service.
    Provides single and multiple request helpers with consistent logging.

    Note: This class intentionally uses camelCase naming for methods and parameters
    to maintain consistency with API naming conventions and existing usage.
    """

    # ------------------------------
    #    CONSTRUCTOR
    # ------------------------------

    def __init__(
        self,
        url: str,
        apimSubscriptionKey: str | None = None,
        headers: dict[str, str] | None = None,
        allowInsecureTls: bool = False,
    ) -> None:
        """
        Initialize the ApimRequests object.

        Args:
            url: The base URL for the APIM endpoint.
            apimSubscriptionKey: Optional subscription key for APIM.
            headers: Optional additional headers to include in requests.
            allowInsecureTls: Whether to bypass TLS certificate verification.
        """

        self._url = url
        self._headers: dict[str, str] = headers.copy() if headers else {}
        self._session: requests.Session | None = None
        self.subscriptionKey = apimSubscriptionKey
        self.allowInsecureTls = allowInsecureTls

        self._headers['Accept'] = 'application/json'

    # ------------------------------
    #    PROPERTIES
    # ------------------------------

    # apimSubscriptionKey
    @property
    def subscriptionKey(self) -> str | None:
        """
        Gets the APIM subscription key, if defined.

        Returns:
            str | None: The APIM subscrption key, if defined; otherwise None.
        """
        return self._subscriptionKey

    @subscriptionKey.setter
    def subscriptionKey(self, value: str | None) -> None:
        """
        Sets the APIM subscription key for the request to use.

        Args:
            value: The APIM subscription key to use or None to not use any key for the request
        """

        self._subscriptionKey = value

        if self._subscriptionKey:
            self._headers[SUBSCRIPTION_KEY_PARAMETER_NAME] = self._subscriptionKey
        else:
            # Remove subscription key from headers if it exists
            self._headers.pop(SUBSCRIPTION_KEY_PARAMETER_NAME, None)

    # headers
    @property
    def headers(self) -> dict[str, str]:
        """
        Get the HTTP headers used for requests.

        Returns:
            dict[str, str]: The headers dictionary.
        """
        return self._headers

    @headers.setter
    def headers(self, value: dict[str, str]) -> None:
        """
        Set the HTTP headers used for requests.

        Args:
            value: The new headers dictionary.
        """
        self._headers = value

    # allowInsecureTls
    @property
    def allowInsecureTls(self) -> bool:
        """
        Get whether TLS certificate verification is disabled for requests.

        Returns:
            bool: True when TLS verification is bypassed; otherwise False.
        """
        return self._allowInsecureTls

    @allowInsecureTls.setter
    def allowInsecureTls(self, value: bool) -> None:
        """
        Set whether TLS certificate verification is disabled for requests.

        Args:
            value: True to bypass TLS certificate verification; otherwise False.
        """
        self._allowInsecureTls = value

    # ------------------------------
    #    PRIVATE METHODS
    # ------------------------------

    def _get_session(self) -> requests.Session:
        """Return the reusable HTTP session, creating it on first use."""

        if self._session is None:
            self._session = requests.Session()

        return self._session

    def _execute_request(self, request_callable, *args, **kwargs) -> requests.Response:
        """
        Execute a request with the configured TLS verification behavior.

        Args:
            request_callable: Request function to invoke.
            *args: Positional arguments for the request function.
            **kwargs: Keyword arguments for the request function.

        Returns:
            requests.Response: The HTTP response.
        """
        request_kwargs = kwargs.copy()
        request_kwargs['verify'] = not self.allowInsecureTls

        if is_debug_enabled():
            print_debug(f'Request headers:\n{pretty_json(_redact_headers(request_kwargs.get("headers")))}')
            if request_kwargs.get('json') is not None:
                print_debug(f'Request body:\n{pretty_json(request_kwargs["json"])}')

        if self.allowInsecureTls:
            with warnings.catch_warnings():
                warnings.simplefilter('ignore', urllib3.exceptions.InsecureRequestWarning)
                response = request_callable(*args, **request_kwargs)
        else:
            response = request_callable(*args, **request_kwargs)

        if is_debug_enabled():
            print_debug(f'Response status: {response.status_code} - {response.reason}')
            print_debug(f'Response headers:\n{pretty_json(_redact_headers(response.headers))}')
            print_debug(f'Response body:\n{pretty_json(response.text)}')

        return response

    def _request(
        self,
        method: HTTP_VERB,
        path: str,
        headers: list[any] = None,
        data: any = None,
        msg: str | None = None,
        printResponse: bool = True,
    ) -> str | None:
        """
        Make a request to the Azure API Management service.

        Args:
            method: The HTTP method to use (e.g., 'GET', 'POST').
            path: The path to append to the base URL for the request.
            headers: Additional headers to include in the request.
            data: Data to include in the request body.
            printResponse: Whether to print the returned output.

        Returns:
            str | None: The JSON response as a string, or None on error.
        """

        try:
            if msg:
                print_message(msg, blank_above=True)

            # Ensure path has a leading slash
            if not path.startswith('/'):
                path = '/' + path

            url = self._url + path
            print_info(f'{method.value} {url}')

            merged_headers = self.headers.copy()

            if headers:
                merged_headers.update(headers)

            print_info(_redact_headers(merged_headers))

            response = self._execute_request(self._get_session().request, method.value, url, headers=merged_headers, json=data, timeout=30)

            content_type = response.headers.get('Content-Type')

            responseBody = None

            if content_type and 'application/json' in content_type:
                responseBody = json.dumps(response.json(), indent=4)
            else:
                responseBody = response.text

            if printResponse:
                self._print_response(response)

            return responseBody

        except requests.exceptions.RequestException as e:
            print_error(f'Error making request: {e}')
            return None

    def _multiRequest(
        self,
        method: HTTP_VERB,
        path: str,
        runs: int,
        headers: list[any] = None,
        data: any = None,
        msg: str | None = None,
        printResponse: bool = True,
        sleepMs: int | None = None,
    ) -> list[dict[str, Any]]:
        """
        Make multiple requests to the Azure API Management service.

        Args:
            method: The HTTP method to use (e.g., 'GET', 'POST').
            path: The path to append to the base URL for the request.
            runs: The number of times to run the request.
            headers: Additional headers to include in the request.
            data: Data to include in the request body.
            printResponse: Whether to print the returned output.
            sleepMs: Optional sleep time between requests in milliseconds (0 to not sleep).

        Returns:
            List of response dicts for each run.
        """

        api_runs = []

        merged_headers = self.headers.copy()
        if headers:
            merged_headers.update(headers)

        if msg:
            print_message(msg, blank_above=True)

        # Ensure path has a leading slash
        if not path.startswith('/'):
            path = '/' + path

        url = self._url + path
        print_info(f'{method.value} {url}')

        for i in range(runs):
            print_info(f'▶️ Run {i + 1}/{runs}:')

            start_time = time.time()
            response = self._execute_request(self._get_session().request, method.value, url, headers=merged_headers, json=data)
            response_time = time.time() - start_time
            print_info(f'⌚ {response_time:.2f} seconds')

            if printResponse:
                self._print_response(response)
            else:
                self._print_response_code(response)

            content_type = response.headers.get('Content-Type')

            if content_type and 'application/json' in content_type:
                resp_data = json.dumps(response.json(), indent=4)
            else:
                resp_data = response.text

            api_runs.append(
                {
                    'run': i + 1,
                    'response': resp_data,
                    'status_code': response.status_code,
                    'response_time': response_time,
                    'headers': dict(response.headers),
                }
            )

            # Sleep only between requests (not after the final run)
            if i < runs - 1:
                if sleepMs is not None:
                    if sleepMs > 0:
                        time.sleep(sleepMs / 1000)
                else:
                    time.sleep(SLEEP_TIME_BETWEEN_REQUESTS_MS / 1000)  # default sleep time

        return api_runs

    def _print_response(self, response) -> None:
        """
        Print the response headers and body with appropriate formatting.
        """

        self._print_response_code(response)
        print_val('Response headers', _redact_headers(response.headers), True)

        if response.status_code == HttpStatusCode.OK:
            try:
                data = json.loads(response.text)
                print_val('Response body', json.dumps(data, indent=4), True)
            except Exception:
                print_val('Response body', response.text, True)
        else:
            print_val('Response body', response.text, True)

    def _print_response_code(self, response) -> None:
        """
        Print the response status code with color formatting.
        """

        if HttpStatusCode.OK <= response.status_code < HttpStatusCode.MULTIPLE_CHOICES:
            status_code_str = f'{BOLD_G}{response.status_code} - {response.reason}{RESET}'
        elif response.status_code >= HttpStatusCode.BAD_REQUEST:
            status_code_str = f'{BOLD_R}{response.status_code} - {response.reason}{RESET}'
        else:
            status_code_str = str(response.status_code)

        print_val('Response status', status_code_str)

    def _poll_async_operation(self, location_url: str, headers: dict = None, timeout: int = 60, poll_interval: int = 2) -> requests.Response | None:
        """
        Poll an async operation until completion.

        Args:
            location_url: The URL from the Location header
            headers: Headers to include in polling requests
            timeout: Maximum time to wait in seconds
            poll_interval: Time between polls in seconds

        Returns:
            The final response when operation completes or None on error
        """
        start_time = time.time()

        while time.time() - start_time < timeout:
            try:
                print_info(f'GET {location_url}', True)
                print_info(_redact_headers(headers))
                response = self._execute_request(self._get_session().get, location_url, headers=headers or {}, timeout=30)

                print_info(f'Polling operation - Status: {response.status_code}')

                if response.status_code == HttpStatusCode.OK:
                    print_ok('Async operation completed successfully!')
                    return response

                if response.status_code == HttpStatusCode.ACCEPTED:
                    print_info(f'Operation still in progress, waiting {poll_interval} seconds...')
                    time.sleep(poll_interval)
                else:
                    print_error(f'Unexpected status code during polling: {response.status_code}')
                    return response

            except requests.exceptions.RequestException as e:
                print_error(f'Error polling operation: {e}')
                return None

        print_error(f'Async operation timeout reached after {timeout} seconds')
        return None

    # ------------------------------
    #    PUBLIC METHODS
    # ------------------------------

    def close(self) -> None:
        """Close the reusable HTTP session."""

        if self._session is not None:
            self._session.close()
            self._session = None

    def __enter__(self) -> 'ApimRequests':
        return self

    def __exit__(self, exc_type, exc_value, traceback) -> None:
        self.close()

    def singleGet(self, path: str, headers=None, msg: str | None = None, printResponse: bool = True) -> Any:
        """
        Make a GET request to the Azure API Management service.

        Args:
            path: The path to append to the base URL for the request.
            headers: Additional headers to include in the request.
            printResponse: Whether to print the returned output.

        Returns:
            str | None: The JSON response as a string, or None on error.
        """

        return self._request(method=HTTP_VERB.GET, path=path, headers=headers, msg=msg, printResponse=printResponse)

    def singlePost(self, path: str, *, headers=None, data=None, msg: str | None = None, printResponse: bool = True) -> Any:
        """
        Make a POST request to the Azure API Management service.

        Args:
            path: The path to append to the base URL for the request.
            headers: Additional headers to include in the request.
            data: Data to include in the request body.
            printResponse: Whether to print the returned output.

        Returns:
            str | None: The JSON response as a string, or None on error.
        """

        return self._request(
            method=HTTP_VERB.POST,
            path=path,
            headers=headers,
            data=data,
            msg=msg,
            printResponse=printResponse,
        )

    def multiGet(
        self,
        path: str,
        runs: int,
        headers=None,
        data=None,
        msg: str | None = None,
        printResponse: bool = True,
        sleepMs: int | None = None,
    ) -> list[dict[str, Any]]:
        """
        Make multiple GET requests to the Azure API Management service.

        Args:
            path: The path to append to the base URL for the request.
            runs: The number of times to run the GET request.
            headers: Additional headers to include in the request.
            data: Data to include in the request body.
            printResponse: Whether to print the returned output.
            sleepMs: Optional sleep time between requests in milliseconds (0 to not sleep).

        Returns:
            List of response dicts for each run.
        """

        return self._multiRequest(
            method=HTTP_VERB.GET,
            path=path,
            runs=runs,
            headers=headers,
            data=data,
            msg=msg,
            printResponse=printResponse,
            sleepMs=sleepMs,
        )

    def multiPost(
        self,
        path: str,
        runs: int,
        headers=None,
        data=None,
        msg: str | None = None,
        printResponse: bool = True,
        sleepMs: int | None = None,
    ) -> list[dict[str, Any]]:
        """Make POST requests and return each response's body, status and timing."""
        return self._multiRequest(
            method=HTTP_VERB.POST,
            path=path,
            runs=runs,
            headers=headers,
            data=data,
            msg=msg,
            printResponse=printResponse,
            sleepMs=sleepMs,
        )

    def singlePostAsync(
        self,
        path: str,
        *,
        headers=None,
        data=None,
        msg: str | None = None,
        printResponse=True,
        timeout=60,
        poll_interval=2,
    ) -> Any:
        """
        Make an async POST request to the Azure API Management service and poll until completion.

        Args:
            path: The path to append to the base URL for the request.
            headers: Additional headers to include in the request.
            data: Data to include in the request body.
            msg: Optional message to display.
            printResponse: Whether to print the returned output.
            timeout: Maximum time to wait for completion in seconds.
            poll_interval: Time between polls in seconds.

        Returns:
            str | None: The JSON response as a string, or None on error.
        """

        try:
            if msg:
                print_message(msg, blank_above=True)

            # Ensure path has a leading slash
            if not path.startswith('/'):
                path = '/' + path

            url = self._url + path
            print_info(f'POST {url}')

            merged_headers = self.headers.copy()

            if headers:
                merged_headers.update(headers)

            print_info(_redact_headers(merged_headers))

            # Make the initial async request
            response = self._execute_request(
                self._get_session().request,
                HTTP_VERB.POST.value,
                url,
                headers=merged_headers,
                json=data,
                timeout=30,
            )

            print_info(f'Initial response status: {response.status_code}')

            if response.status_code == HttpStatusCode.ACCEPTED:  # Accepted - async operation started
                location_header = response.headers.get('Location')

                if location_header:
                    print_info(f'Found Location header: {location_header}')

                    # Poll the location URL until completion
                    final_response = self._poll_async_operation(location_header, timeout=timeout, poll_interval=poll_interval)

                    if final_response and final_response.status_code == HttpStatusCode.OK:
                        if printResponse:
                            self._print_response(final_response)

                        content_type = final_response.headers.get('Content-Type')
                        responseBody = None

                        if content_type and 'application/json' in content_type:
                            responseBody = json.dumps(final_response.json(), indent=4)
                        else:
                            responseBody = final_response.text

                        return responseBody

                    print_error('Async operation failed or timed out')
                    return None

                print_error('No Location header found in 202 response')
                if printResponse:
                    self._print_response(response)
                return None

            # Non-async response, handle normally
            if printResponse:
                self._print_response(response)

            content_type = response.headers.get('Content-Type')
            responseBody = None

            if content_type and 'application/json' in content_type:
                responseBody = json.dumps(response.json(), indent=4)
            else:
                responseBody = response.text

            return responseBody

        except requests.exceptions.RequestException as e:
            print_error(f'Error making request: {e}')
            return None