import json
import functools
from typing import Callable, Optional, Union
from q2_sdk.core.exceptions import TectonError
from q2_sdk.core.http_handlers.tecton_base_handler import (
Q2TectonBaseRequestHandler,
MFAMessage,
EdvMfaType,
_detect_mfa_exit,
_mfa_handling_supported,
_build_mfa_response_payload,
_run_audit_check,
)
from q2_sdk.models.tecton import InternalServerError, Success, BadRequest
def _handle_mfa_exit(handler, error_handler):
detected, result = _detect_mfa_exit(handler, error_handler)
if not detected:
return None
return result if result is not None else Success({})
async def _handle_audit_and_response(
handler,
func,
error_handler: Optional[Callable],
audit_action_name: str,
audit_details: Optional[dict | str | Callable],
soft_failures: bool = False,
):
"""
Handle audit record creation and processing the response.
"""
return await _run_audit_check(
handler,
func,
error_handler,
audit_action_name,
audit_details,
format_challenge=lambda audit_response: Success(
_build_mfa_response_payload(audit_response)
),
format_denial=lambda err_msg, audit_response: InternalServerError(
"Action Denied",
{
"detail": err_msg,
"errorReturnCode": audit_response.get("return_code"),
},
),
soft_failures=soft_failures,
)
def inner_mfa_function(
func,
error_handler: Optional[Callable] = None,
audit_action_name: Optional[str] = None,
audit_details: Optional[Union[dict, str, Callable]] = None,
mfa_type: Optional[EdvMfaType] = None,
single_use: bool = False,
):
"""
Decorator for checking if session has been validated by the MFA workflow
:param mfa_type: Optional EdvMfaType enum value (defaults to PatrolEDV if not specified)
:param single_use: If True, deletes the cached token after the handler runs, forcing re-validation on the next request
"""
@functools.wraps(func)
async def wrapper(handler):
exit_response = _handle_mfa_exit(handler, error_handler)
if exit_response is not None:
return exit_response
validated = handler.validate_mfa_token()
if validated:
handler.logger.debug(MFAMessage.VALIDATED.value)
result = await func(handler)
if single_use:
handler.consume_mfa_token()
return result
else:
handler.logger.debug(MFAMessage.NOT_FOUND.value)
if _mfa_handling_supported(handler):
if audit_action_name:
return await _handle_audit_and_response(
handler, func, error_handler, audit_action_name, audit_details
)
else:
mfa_type_response = await handler._get_mfa_type(mfa_type=mfa_type)
response = _build_mfa_response_payload(mfa_type_response)
return Success(response)
if error_handler:
return error_handler(handler, MFAMessage.NOT_FOUND.value)
else:
handler.logger.debug(MFAMessage.NO_ERROR_HANDLER.value)
return InternalServerError(MFAMessage.NOT_FOUND.value)
return wrapper
def mfa_validation_required(
error_handler: Optional[Callable] = None,
audit_action_name: Optional[str] = None,
audit_details: Optional[Union[dict, str, Callable]] = None,
mfa_type: Optional[EdvMfaType] = None,
single_use: bool = False,
):
"""
Decorator for checking if session has been validated by the MFA workflow
:param error_handler: Optional callable to handle errors
:param audit_action_name: Optional audit action name for EDV workflow
:param audit_details: Optional audit details for EDV workflow
:param mfa_type: Optional EdvMfaType enum value (HQ defaults to PatrolEDV if not specified)
:param single_use: If True, deletes the cached token after the handler runs, forcing re-validation on the next request
"""
return functools.partial(
inner_mfa_function,
error_handler=error_handler,
audit_action_name=audit_action_name,
audit_details=audit_details,
mfa_type=mfa_type,
single_use=single_use,
)
def inner_audit_function(
func,
audit_action_name: str,
audit_details: Optional[Union[dict, str, Callable]] = None,
error_handler: Optional[Callable] = None,
csr_only: bool = False,
soft_failures: bool = False,
):
"""
Decorator for creating an audit record and handling edv response as needed
"""
@functools.wraps(func)
async def wrapper(handler):
if csr_only and handler.online_session.is_csr_assist is False:
return await func(handler)
exit_response = _handle_mfa_exit(handler, error_handler)
if exit_response is not None:
return exit_response
if _mfa_handling_supported(handler):
return await _handle_audit_and_response(
handler,
func,
error_handler,
audit_action_name,
audit_details,
soft_failures,
)
if error_handler:
return error_handler(handler, MFAMessage.NOT_FOUND.value)
else:
handler.logger.debug(MFAMessage.NO_ERROR_HANDLER.value)
return BadRequest(MFAMessage.NOT_FOUND.value)
return wrapper
def create_audit(
audit_action_name: str,
audit_details: Optional[Union[dict, str, Callable]] = None,
error_handler: Optional[Callable] = None,
csr_only: bool = False,
soft_failures: bool = False,
):
return functools.partial(
inner_audit_function,
error_handler=error_handler,
audit_action_name=audit_action_name,
audit_details=audit_details,
csr_only=csr_only,
soft_failures=soft_failures,
)
[docs]
class Q2TectonClientRequestHandler(Q2TectonBaseRequestHandler):
"""
RequestHandler meant to be used for requests incoming from Online using the Tecton Client Side rendered content
"""
OPTIONAL_CONFIGURATIONS = {"FEATURE": None}
def __init__(self, application, request, **kwargs):
super().__init__(application, request, **kwargs)
self.return_as_html = False
[docs]
async def q2_post(self, *args, **kwargs):
try:
self.form_fields = self._parse_tecton_payload()
route_response = await self.route_request()
except TectonError as exc:
self.logger.error(exc)
error_args = exc.args
error_message = error_args[0] if error_args else str(exc)
response_type = (
BadRequest
if len(error_args) > 2 and error_args[2] == 400
else InternalServerError
)
return json.dumps(
response_type(
error_message, error_args[1] if len(error_args) > 1 else None
).to_json()
)
if isinstance(route_response, dict):
json_dump = json.dumps(Success(route_response).to_json())
elif isinstance(route_response, str):
json_dump = route_response
else:
json_dump = json.dumps(vars(route_response))
self.logger.debug("Tecton response: %s", json_dump)
return json_dump if json_dump else ""
async def mfa_validate(self):
return json.loads(await super().mfa_validate())["data"]