""" AUTHOR: Khushal P Soonderji DATE: Friday, 2nd Aug., 2024 OBJECTIVE: To provide an easy way to perform repetitive tasks in quart. REFERENCES: N/A DOWNLOADS: N/A NOTES: 01. PLEASE USE "ResponseModel" AS THE RETURNED VALUE OF THE API ENDPOINT IF YOU ARE USING ANY OF THESE DECORATORS. """ import copy # ***************************************************************************************************************** # ***** **** # *** IMPORT *** # ***** **** # ***************************************************************************************************************** # To work with Quart: from quart import request, current_app, g, make_response # To make decorators: from functools import wraps # My utils: from utils_v2.string import json from utils_v2.date_time import date_time from utils_v2.security import sanitizers from utils_v2.api.codes import StatusCodes, HttpCodes from utils_v2.api.log import APILogModel from utils_v2.api.response import ResponseModel from utils_v2.logging.context import AsyncLoggerContext, AsyncMongoLogger # System-level activities: import io import os # To work with datatypes: from types import NoneType import pandas as pd # For Pydantic data-behaviour_models: import pydantic from typing import List # For hashing and shortening the hash: import hashlib import base64 # To make things human-readable: import humanize # For debugging: import traceback import random import string # For asynchronous activities: import asyncio # For timekeeping: import time import datetime # To make copies: import copy # ***************************************************************************************************************** # ***** **** # *** MACROS / ONE-TIME INIT *** # ***** **** # ***************************************************************************************************************** # Chars to choose from for random strings: ALPHANUMERIC_CHARS = string.ascii_letters + string.digits # To capture system information: PROCESS_ID = os.getppid() PARENT_PROCESS_ID = os.getppid() # ***************************************************************************************************************** # ***** **** # *** VARIABLES *** # ***** **** # ***************************************************************************************************************** # --- Nothing Yet # ***************************************************************************************************************** # ***** **** # *** FUNCTIONS *** # ***** **** # ***************************************************************************************************************** # --- Nothing Yet # ***************************************************************************************************************** # ***** **** # *** EXCEPTIONS *** # ***** **** # ***************************************************************************************************************** class AuthDetailsIncompleteException(Exception): def __str__(self): return "incomplete auth details" # ***************************************************************************************************************** # ***** **** # *** FUNCTIONS *** # ***** **** # ***************************************************************************************************************** async def make_ordered_json( json_data, no_space = True, http_code = 200 ): """ Quart sorts the fields of a dict when converting to a JSON response. Here we are manually making the response when the sequence of the fields is sensitive. :param json_data: The data (dict, list, etc.) to be converted to a JSON string. :param no_space: Set this to True to remove all excess white spaces from the JSON string. Saves bandwidth. :param http_code: The HTTP status code you want to send with the response. :return: The JSON-ified response such that the sequence of the fields is maintained. """ response = await make_response(json.to_string(json_data, no_space = no_space), http_code) response.headers["Content-Type"] = "application/json" return response # --------------------------------------------------------------------------------------------------------------------- async def data_from_request(inbound_request): """ Adaptively extract the params from incoming request in whichever way it was provided. :param inbound_request: The request that came in. :return: A dict (could be empty) of the data/params that came in with the request. """ inbound_data = {} # Extract data from the params in the URL: from_args = inbound_request.args.to_dict() if isinstance(from_args, dict): for k, v in from_args.items(): inbound_data[k] = v # Extract data from the raw JSON data: from_json = await request.get_json() if isinstance(from_json, dict): for k, v in from_json.items(): inbound_data[k] = v # Extract inputs from the form body: from_form = await inbound_request.form from_form = from_form.to_dict() if isinstance(from_form, dict): for k, v in from_form.items(): inbound_data[k] = v return inbound_data # --------------------------------------------------------------------------------------------------------------------- async def file_from_request(inbound_request, file_key): """ Extracts ONE file from the incoming request's form-data. WARNING: NOT RECOMMENDED FOR LARGE FILES. STRICTLY USE FOR SMALL FILES THAT WON'T CRASH THE SCRIPT. :param inbound_request: The incoming request. :param file_key: The key of the file that you want to extract. :return: A tuple of the file's name and data. """ file_name = None file_data = None files = await inbound_request.files if file_key in files: file_name = files[file_key].filename file_data = io.BytesIO(files[file_key].read()) return file_name, file_data # --------------------------------------------------------------------------------------------------------------------- async def files_from_request(inbound_request: request): """ Extracts ALL files from the incoming request's form-data. WARNING: NOT RECOMMENDED FOR LARGE FILES. STRICTLY USE FOR SMALL FILES THAT WON'T CRASH THE SCRIPT. :param inbound_request: The incoming request. :return: A dict describing the file's name, data, and size. """ files = await inbound_request.files inbound_files = {} for file_key in files: file_data = io.BytesIO(files[file_key].read()) file_size = file_data.seek(0, 2) file_data.seek(0) inbound_files[file_key] = { "name": files[file_key].filename, "data": file_data, "size": file_size, "type": files[file_key].content_type } return inbound_files # --------------------------------------------------------------------------------------------------------------------- async def headers_from_request( inbound_request, mandatory_keys: list = None ): """ Extract custom headers and some extra info. from the incoming request. Raises an exception if any mandatory key is missing. IMPORTANT: CUSTOMIZE THIS FOR THE NEEDS OF YOUR PROJECT. :param inbound_request: The request that came in. :param mandatory_keys: The keys that you need to have in the auth. :return: The extracted auth details. """ # Start by extracting whatever complies with the format of "X-{Header-Name}": head_json = {k: v for k, v in request.headers.items() if k.startswith("X-")} # Now note down things that you want to keep from incoming requests: head_json["Remote-IP"] = inbound_request.remote_addr head_json["Host"] = inbound_request.headers.get("Host") head_json["Origin"] = inbound_request.headers.get("Origin") head_json["User-Agent"] = inbound_request.headers.get("User-Agent") head_json["Content-Type"] = inbound_request.headers.get("Content-Type") head_json["Content-Length"] = inbound_request.headers.get("Content-Length") # Raise an exception if any of the mandatory auth details were missing: if mandatory_keys is not None: available_keys = head_json.keys() for mandatory_key in mandatory_keys: if mandatory_key not in available_keys: raise AuthDetailsIncompleteException # Done here: return head_json # --------------------------------------------------------------------------------------------------------------------- def cause_exception(): """ Call this from any function when you want to raise an exception. Example use case would be when receiving data from an API call and that field is not supposed to be null. :return: None. """ return 100/0 # --------------------------------------------------------------------------------------------------------------------- def describe_exception(exc): """ Describes the exception in detail. It extracts the type of exception, a brief message, and even the entire traceback. Useful for debugging in details without the terminal. You could either log the resultant dict or send it to the dev team over some service like WhatsApp/Telegram. :param exc: The exception that occurred. :return: The dict that explains the exception. """ exc_desc = { "type": type(exc).__name__, "msg": str(exc), "tb": [str(exc_tb) for exc_tb in traceback.format_exception(exc, value = exc, tb = exc.__traceback__)] } return exc_desc # --------------------------------------------------------------------------------------------------------------------- def messages_from_pydantic_exception(exception, as_str = True, sep = ", "): """ Creates a list of readable error messages from Pydantic's validation failure. :param exception: Pydantic's ValidationError :param as_str: Set to True to receive all messages as one string, False to receive an array of strings. :param sep: The separator to use when joining multiple messages as one string. :return: A list of messages of all the things that went wrong. """ # Make a variable to hold all individual messages: messages = [] # Interpret all the problems: for error in exception.errors(): loc = " --> ".join([str(item) for item in error["loc"]]) msg = error.get("msg") if error["type"] == "missing": messages.append(f"missing input: '{loc}'") elif error["type"] == "model_type": messages.append(f"invalid input: '{loc}'") elif error["type"] == "bool_parsing": messages.append(f"invalid bool: '{loc}'") elif error["type"] == "string_type": messages.append(f"invalid string: '{loc}'") elif error["type"] == "float_parsing": messages.append(f"invalid float: '{loc}'") elif error["type"] == "int_parsing": messages.append(f"invalid integer: '{loc}'") elif error["type"] == "extra_forbidden": messages.append(f"extra input: '{loc}'") elif error["type"] == "value_error": messages.append(f"validation failed: '{loc}' ({msg})") else: messages.append(f"invalid value: '{loc}'") # Return a response as per the preference of the user: if as_str: return sep.join(messages) else: return messages # --------------------------------------------------------------------------------------------------------------------- def summarize_variable( value, str_limit = 1024, expand: bool | int = False, sensitive_keys: list[str] = None ): """ To summarize an input value to capture the essence without hoarding to much data. :param value: Anything that you want to summarize. :param str_limit: The max. no. of chars of a string to retain. :param expand: Set to True for full expansion, False for no expansion, and an integer for a specific level of expansion. Applicable on iterables and dicts. The smaller this number, the more concise the summary will be, and vice versa. :param sensitive_keys: The list of keys (of a dict) to obscure when summarizing. :return: The summarized version of the input. """ # Basic prep: if value is None: return None if sensitive_keys is None: sensitive_keys = [] value_copy = copy.deepcopy(value) # If the value is a Pydantic model, we convert to dict: if isinstance(value_copy, pydantic.BaseModel): value_copy = value_copy.model_dump() # Handle datatypes that you don't want to modify: if isinstance(value_copy, (int, float, bool, NoneType)): pass # When the value is a list or similar iterable: elif isinstance(value_copy, (list, tuple, set)): if expand: if not isinstance(expand, bool): expand -= 1 value_copy = [AsyncLoggerContext.summarize( v, expand = expand, sensitive_keys = sensitive_keys ) for v in value_copy] else: value_copy = f"array of {len(value_copy)} item(s)" # If the value is a dict: elif isinstance(value_copy, dict): if expand: if not isinstance(expand, bool): expand -= 1 value_copy = { k: AsyncLoggerContext.summarize( v, expand = expand, sensitive_keys = sensitive_keys ) if k not in sensitive_keys else f"{len(str(v))} sensitive char(s)" for k, v in value_copy.items() } else: value_copy = f"object of {len(value_copy.keys())} field(s) [{', '.join(value_copy.keys())}]" # When a dataframe is passed: elif isinstance(value_copy, pd.DataFrame): cols = value_copy.columns.to_list() value_copy = f"table with {len(cols)} col(s) [{', '.join(cols)}] and {len(value_copy)} row(s)" str_limit = 1024 # If the input is some form of non-standard object: else: value_copy = str(value_copy) # Handle strings: if isinstance(value_copy, str): if len(value_copy) > str_limit: value_copy = value_copy[:str_limit] + "...(trunc'd)" # Done here: return value_copy # --------------------------------------------------------------------------------------------------------------------- def set_api_version(api_version): """ Use this decorator to automatically note down the API version no. and propagate it throughout the downstream decorators. Use this as the entry point if possible. :param api_version: The version code to assign to the API. :return: The decorator factory. """ def decorator(func): @wraps(func) async def wrapper(*args, **kwargs): # set the version information in the variable. # This makes it available to the downstream decorators too! kwargs["api_version"] = api_version # we are ready to call the function that we are wrapping: kwargs["decorator_count"] = kwargs.get("decorator_count", 0) + 1 response = await func(*args, **kwargs) kwargs["decorator_count"] -= 1 # Add the API version to the response: if isinstance(response, ResponseModel): response.api_version = api_version # Done here: if ( kwargs["decorator_count"] == 0 and isinstance(response, ResponseModel) ): response = response.for_quart() return response return wrapper return decorator # --------------------------------------------------------------------------------------------------------------------- def read_input( read_headers = True, read_data = True, read_files = True, sanitize_headers = True, sanitize_data = True ): """ Use this decorator to read the inputs from the incoming request and sanitize them. Sanitization makes the inputs safe against certain threats like injections attacks. If you expect to take in inputs that you want to use to run database commands, you could disable them manually. PLEASE USE "ResponseModel" AS THE RETURNED VALUE OF THE API ENDPOINT IF YOU ARE USING THIS DECORATOR. :param read_headers: Whether, or not, to read the headers of the request. :param read_data: Whether, or not, to read the data of the request. :param read_files: Whether, or not, to read the files of the request. :param sanitize_headers: Whether, or not, you would like to sanitize the params coming in through the headers. :param sanitize_data: Whether, or not, you would like to sanitize the params coming in through the body or query. :return: The decorator factory. """ def decorator(func): @wraps(func) async def wrapper(*args, **kwargs): # Get the headers: if read_headers: kwargs["inbound_headers"] = await headers_from_request(request) if sanitize_headers: kwargs["inbound_headers"] = sanitizers.for_mongo(kwargs["inbound_headers"]) # else: kwargs["inbound_headers"] = None # Get the data: if read_data: kwargs["inbound_data"] = await data_from_request(request) if sanitize_data: kwargs["inbound_data"] = sanitizers.for_mongo(kwargs["inbound_data"]) # else: kwargs["inbound_data"] = None # Get small files from the request: if read_files: kwargs["inbound_files"] = await files_from_request(request) # else: kwargs["inbound_files"] = None # # We also make a provision for capturing an identifier # # for the logs that we make through a sister decorator: # kwargs["log_id"] = "".join(random.choice(ALPHANUMERIC_CHARS) for _ in range(8)) # kwargs["log_chain"] = kwargs.get("inbound_headers", {}).get("X-Log-Chain") or kwargs["log_id"] # Now that we have unpacked the incoming data, # we are ready to run the function that we are wrapping: kwargs["decorator_count"] = kwargs.get("decorator_count", 0) + 1 response = await func(*args, **kwargs) kwargs["decorator_count"] -= 1 # Done here: if ( kwargs["decorator_count"] == 0 and isinstance(response, ResponseModel) ): response = response.for_quart() return response return wrapper return decorator # --------------------------------------------------------------------------------------------------------------------- def get_session_info( key: str, model: str = None, session_func: str = None, session_coro: str = None, get: str = None, mandatory: bool = False, sensitive_keys: List[str] = None ): """ To get the information about the user from the session token or some similar identifier. Use this after 'read_input', and note that the callable func/coro should only take on parameter - the thing to identify the session by, and it should return a dictionary with all the needed details. PLEASE USE "ResponseModel" AS THE RETURNED VALUE OF THE API ENDPOINT IF YOU ARE USING THIS DECORATOR. :param key: The key in either the 'inbound_data' or the 'inbound_headers' from which the session info will be available. Examples: 'sessionToken' or 'X-Session-Token'. :param model: The object from which the func/coro should be called. This could be the name of the variable holding an instance of a class. Should be available in the scope of 'current_app'. :param session_func: The synchronous func of the model to call. This can either be an independent function or a class's method. Should be available in the scope of 'current_app'. :param session_coro: The asynchronous func of the model to call. This can either be an independent function or a class's method. Ignored if 'func' was provided. Should be available in the scope of 'current_app'. :param get: The field inside the dict to get. Specify this in dot notation. The whole response will be returned as is if this is null. :param mandatory: If set to True, the API call will enforce session checking; and if a session is not found, the client will receive an unauthorized failure message. :param sensitive_keys: Any keys to obscure when retrieving details about the session. :return: The decorator factory. """ def decorator(func): @wraps(func) async def wrapper(*args, **kwargs): # Start by assuming failure: session_info = None # Get the identifier of the session from the values extracted in the 'read_input' decorator: session_id = kwargs["inbound_data"].get(key, kwargs["inbound_headers"].get(key)) # If the key is present, # we call the function that will get the session's info for us: if session_id is not None: try: # For synchronous functions: if session_func: if model: session_info = getattr(getattr(current_app, model), session_func)(session_id) else: session_info = getattr(current_app, session_func)(session_id) # For asynchronous coroutines: elif session_coro: if model: session_info = await getattr(getattr(current_app, model), session_coro)(session_id) else: session_info = await getattr(current_app, session_coro)(session_id) # If something goes wrong: except Exception as exception: if hasattr(current_app, "printer"): getattr(current_app, "printer")(exception) # We note down whatever we got: if session_info and isinstance(get, str): for subkey in get.split("."): session_info = session_info.get(subkey, {}) if isinstance(session_info, dict) else {} session_info = summarize_variable(session_info, expand = True, sensitive_keys = sensitive_keys) kwargs["session_info"] = session_info # If no session info was found, but it was mandatory: if mandatory and not session_info: response = ResponseModel( status_code = StatusCodes.FAILED, http_code = HttpCodes.UNAUTHORIZED, message = f"invalid session" ) # Now we are ready to run the function that we are wrapping: else: kwargs["decorator_count"] = kwargs.get("decorator_count", 0) + 1 response = await func(*args, **kwargs) kwargs["decorator_count"] -= 1 # Done here: if ( kwargs["decorator_count"] == 0 and isinstance(response, ResponseModel) ): response = response.for_quart() return response return wrapper return decorator # --------------------------------------------------------------------------------------------------------------------- def validate_input( mandatory_header_keys = None, mandatory_data_keys = None, mandatory_file_keys = None, header_validator = None, data_validator = None ): """ USE THIS ONLY AFTER YOU HAVE USED 'read_input'. This decorator will help you run validation on the inputs that were extracted from the request. The mandatory keys will be checked first and the validations will be run after that. If your validator already checks for keys, you may skip mentioning mandatory keys. DO NOTE THAT YOUR VALIDATOR FUNCTIONS MUST RAISE AN EXCEPTION FOR THIS DECORATOR TO WORK. :param mandatory_header_keys: The keys in 'inbound_headers' that are absolutely necessary. :param mandatory_data_keys: The keys in 'inbound_data' that are absolutely necessary. :param mandatory_file_keys: The keys in 'inbound_files' that are absolutely necessary. :param header_validator: The function to use to validate the 'inbound_headers'. :param data_validator: The function to use to validate 'inbound_data'. :return: The decorator factory. """ def decorator(func): @wraps(func) async def wrapper(*args, **kwargs): # We first validate the mandatory header keys. # Having a null value in this case is NOT allowed: if mandatory_header_keys is not None: for mandatory_key in mandatory_header_keys: if kwargs["inbound_headers"].get(mandatory_key) is None: return ResponseModel( status_code = StatusCodes.HEADERS_INCOMPLETE, message = f"missing: '{mandatory_key}'" ) # Return with failure if any of the mandatory JSON details are missing. # Having a null value is allowed, but is should be sent by the user on intention. if mandatory_data_keys is not None: for mandatory_key in mandatory_data_keys: try: kwargs["inbound_data"][mandatory_key] except: return ResponseModel( status_code = StatusCodes.DATA_INCOMPLETE, message = f"missing: '{mandatory_key}'" ) # Return with failure if any of the mandatory file-keys details are missing: if mandatory_file_keys is not None: provided_file_keys = kwargs["inbound_files"].keys() for mandatory_key in mandatory_file_keys: if mandatory_key not in provided_file_keys: return ResponseModel( status_code = StatusCodes.FILE_MISSING, message = f"missing: '{mandatory_key}'" ) # Next we validate the headers: if header_validator is not None: # Try validate the data: try: kwargs["inbound_headers"] = header_validator(kwargs["inbound_headers"]) # In case some needed field is missing: except KeyError as exception: return ResponseModel( status_code = StatusCodes.DATA_VALIDATION_FAILURE, message = "missing: " + str(exception), http_code = HttpCodes.BAD_REQUEST, ) # In case some pydantic data model fails validation: except pydantic.ValidationError as exception: return ResponseModel( status_code = StatusCodes.DATA_VALIDATION_FAILURE, message = messages_from_pydantic_exception(exception), http_code = HttpCodes.BAD_REQUEST ) # In case some other exception was raised: except Exception as exception: return ResponseModel( status_code = StatusCodes.DATA_VALIDATION_FAILURE, message = str(exception), http_code = HttpCodes.BAD_REQUEST ) # Finally, we validate the incoming data: if data_validator is not None: # Try validate the data: try: kwargs["inbound_data"] = data_validator(kwargs["inbound_data"]) # In case some needed field is missing: except KeyError as exception: return ResponseModel( status_code = StatusCodes.DATA_VALIDATION_FAILURE, message = "missing: " + str(exception), http_code = HttpCodes.BAD_REQUEST ) # In case some pydantic data model fails validation: except pydantic.ValidationError as exception: return ResponseModel( status_code = StatusCodes.DATA_VALIDATION_FAILURE, message = messages_from_pydantic_exception(exception), http_code = HttpCodes.BAD_REQUEST ) # In case some other exception was raised: except Exception as exception: return ResponseModel( status_code = StatusCodes.DATA_VALIDATION_FAILURE, message = str(exception), http_code = HttpCodes.BAD_REQUEST ) # Now that we have unpacked the incoming data, # we are ready to run the function that we are wrapping: kwargs["decorator_count"] = kwargs.get("decorator_count", 0) + 1 response = await func(*args, **kwargs) kwargs["decorator_count"] -= 1 # Done here: if ( kwargs["decorator_count"] == 0 and isinstance(response, ResponseModel) ): response = response.for_quart() return response return wrapper return decorator # --------------------------------------------------------------------------------------------------------------------- def log_request_to_mongo( attr_name, collection: str = "logs", api_version: str = None, project: str = None, log_type: str = None, operation: str = None, log_input: bool | int = True, log_output: bool | int = True, sensitive_keys: list = None ): """ USE THIS ONLY AFTER YOU HAVE USED 'read_input'. This decorator will log the whole process of the API call to MongoDB. The variable that holds the instance of 'AsyncMongo' needs to be accessible in the scope of 'current_app'. PLEASE USE "ResponseModel" AS THE RETURNED VALUE OF THE API ENDPOINT IF YOU ARE USING THIS DECORATOR. :param attr_name: The name of the variable that holds the instance of 'AsyncMongo'. It should be accessible in the scope of 'current_app'. :param collection: The name of the collection to write the log into. :param api_version: The version code of the API endpoint that is being logged. :param project: The name of the project that the endpoint was built for. :param log_type: A hint to identify what the log was for. :param operation: A hint to identify what was action was being performed. :param log_input: Set to True to capture everything that went into the function, False to capture the least info, and set it to an integer to capture a certain depth of the input (applicable on iterables and dicts. :param log_output: The same as 'log_input', but applicable to the response from the function. :param sensitive_keys: The list of keys to not log. :return: The decorator factory. """ def decorator(func): @wraps(func) async def wrapper(*args, **kwargs): # Create ids for getting logs: kwargs["log_id"] = "".join(random.choice(ALPHANUMERIC_CHARS) for _ in range(8)) kwargs["log_chain"] = kwargs.get("inbound_headers", {}).get("X-Log-Chain") or kwargs["log_id"] # Set the api version as needed: kwargs["api_version"] = kwargs.get("api_version", api_version) # Make variables and extract available info.: exception = None response = None request_ts = date_time.get_current_utc_date_time() start_ts = time.perf_counter() cpu_start_ts = time.process_time() # Execute the function that is being wrapped: kwargs["decorator_count"] = kwargs.get("decorator_count", 0) + 1 try: response = await func(*args, **kwargs) except Exception as exc: exception = exc kwargs["decorator_count"] -= 1 # Ensure that the response is not null: response = response if response is not None else ResponseModel( status_code = StatusCodes.UNKNOWN_ERROR, message = "null response for request" ) # Add params to the response. # THIS IS ONLY APPLICABLE WHEN THE TYPE OF THE RESPONSE IS 'ResponseModel': if isinstance(response, ResponseModel): response.api_version = kwargs.get("api_version") response.log_id = kwargs.get("log_id") # Extract the response to log: response_to_log = "not logged" http_code_to_log = 200 if log_output: if isinstance(response, ResponseModel): response_to_log, http_code_to_log = response.for_quart() elif isinstance(response, tuple): response_to_log, http_code_to_log = response else: response_to_log, http_code_to_log = str(response), 200 response_to_log = summarize_variable( response_to_log, expand = log_output, sensitive_keys = sensitive_keys ) # Try to get the information about the request. # There will be no data in any of these if the decorator was used to catch start-up and shut-down events. request_method = None request_url = None request_route = None try: request_method = f"{request.method}" request_url = f"{request.url}" request_route = str(request.url_rule.rule) except: pass # Redact the sensitive keys: if sensitive_keys: for k in sensitive_keys: for var in ["inbound_headers", "inbound_data"]: try: kwargs[var][k] = f"{len(str(kwargs[var][k]))} sensitive char(s)" except: pass # Construct the log: api_log = APILogModel( pid = PROCESS_ID, ppid = PARENT_PROCESS_ID, project = project, log = log_type, operation = operation, apiVer = kwargs.get("api_version"), logId = kwargs.get("log_id"), logChain = kwargs.get("log_chain", kwargs.get("inbound_headers", {}).get("X-Log-Chain")), ts = request_ts, tat = time.perf_counter() - start_ts, cpuTime = time.process_time() - cpu_start_ts, sessionInfo = summarize_variable(kwargs.get("session_info"), expand = True), method = request_method, url = request_url, route = request_route, headers = kwargs.get("inbound_headers"), data = summarize_variable( kwargs.get("inbound_data"), expand = log_input, sensitive_keys = sensitive_keys ), files = { k: { "name": v["name"], "size": v["size"] } for k, v in kwargs.get("inbound_files", {}).items() }, exception = None if exception is None else describe_exception(exception), response = response_to_log, httpCode = http_code_to_log ) # Write the log: app_attr = getattr(current_app, attr_name) inserted_id = await app_attr.insert_one( collection = collection, document = api_log.model_dump() ) # Return the response from the wrapped function. if ( kwargs["decorator_count"] == 0 and isinstance(response, ResponseModel) ): response = response.for_quart() return response return wrapper return decorator # --------------------------------------------------------------------------------------------------------------------- def log_chain_to_mongo(attr_name): """ Use this to run everything in the context of the logging decorator. Everything that is decorated with the custom decorator will be logged if this is used. :param attr_name: The name of the db connection (an instance of 'AsyncMongo' or 'AsyncMongoStorage') :return: The decorator factory. """ def decorator(func): @wraps(func) async def wrapper(*args, **kwargs): # Let the next in-line decorator know that it has been wrapped: kwargs["decorator_count"] = kwargs.get("decorator_count", 0) + 1 # Fetch the attribute and get it to log the whole chain: async with AsyncLoggerContext.logging_context( logger = AsyncMongoLogger(getattr(current_app, attr_name)), log_chain = kwargs.get("log_chain") ): response = await func(*args, **kwargs) # Done here: kwargs["decorator_count"] -= 1 if ( kwargs["decorator_count"] == 0 and isinstance(response, ResponseModel) ): response = response.for_quart() return response return wrapper return decorator # --------------------------------------------------------------------------------------------------------------------- def should_not_be_under_maintenance(attr_name): """ Use this decorator to reject a request when the app is being marked as "under-maintenance". You will need to create a boolean variable within the scope of the 'current_app' for this to work. An alternate to this is to set the value in an environment variable named 'IS_UNDER_MAINTENANCE' to a string value of either 'True' or 'False' for multi-worker deployments. :param attr_name: The name of the boolean variable that will hold the information about the app being under maintenance. If its value is True at the time of checking, the incoming request will be rejected. :return: The decorator factory. """ def decorator(func): @wraps(func) async def wrapper(*args, **kwargs): # Get the attribute and check if it indicates that the app is under maintenance, # call the wrapped function if not under maintenance: app_attr = getattr(current_app, attr_name) env_attr = True if os.environ.get("IS_UNDER_MAINTENANCE", "False").lower() == "true" else False if app_attr or env_attr: response = ResponseModel(status_code = StatusCodes.DOWN_FOR_MAINTENANCE).for_quart() else: kwargs["decorator_count"] = kwargs.get("decorator_count", 0) + 1 response = await func(*args, **kwargs) kwargs["decorator_count"] -= 1 # Done here: if ( kwargs["decorator_count"] == 0 and isinstance(response, ResponseModel) ): response = response.for_quart() return response return wrapper return decorator # --------------------------------------------------------------------------------------------------------------------- def only_whitelisted_ips(attr_name): """ Use this decorator to reject any requests coming from unauthorized IPs. The list of IP addresses to allow must be in a list that is accessible in the context of 'current_app'. :param attr_name: The name of the boolean variable that will hold the information about the app being under maintenance. If its value is True at the time of checking, the incoming request will be rejected. :return: The decorator factory. """ def decorator(func): @wraps(func) async def wrapper(*args, **kwargs): # Get the attribute and check if the request's IP is in the permitted list: app_attr = getattr(current_app, attr_name) if request.remote_addr not in app_attr: return ResponseModel( status_code = StatusCodes.AUTHORIZATION_FAILED, message = "bad ip", http_code = HttpCodes.UNAUTHORIZED ).for_quart() # Now that we have checked that the IP is permitted, # we are ready to run the function that we are wrapping: kwargs["decorator_count"] = kwargs.get("decorator_count", 0) + 1 response = await func(*args, **kwargs) kwargs["decorator_count"] -= 1 # Done here: if ( kwargs["decorator_count"] == 0 and isinstance(response, ResponseModel) ): response = response.for_quart() return response return wrapper return decorator # --------------------------------------------------------------------------------------------------------------------- def limit_rate( attr_name, rate_limit: int = 5, seconds: float = 1.0, message = None, header_keys: list = None, data_keys: list = None, allow_if_exception = False, count_for_http_codes = None ): """ Use this decorator to apply rate-limiting to incoming requests. SUGGESTION: WHEN STACKING UP MANY RATE LIMITS, PUT THE SMALLEST TIME PERIOD ON TOP AND LARGEST TIME PERIOD AT THE BOTTOM. THIS ENSURES PROPER FUNCTIONALITY. WARNING: TIMING STARTS WHEN THE FIRST PERMITTED REQUEST GOES THROUGH. THIS MEANS THAT, IF YOU HAVE A PER-DAY LIMIT, AND YOU START MAKING REQUESTS AT 11:00 PM AND EXHAUST YOUR LIMIT AT 11:59 PM, YOUR LIMIT WILL BE REPLENISHED AT 11:00 PM OF THE NEXT DAY, NOT AT 12:00 AM. :param attr_name: The name of the variable that holds the instance of 'AsyncRedisCache'. Should be available in the context of 'current_app'. :param rate_limit: The number of requests per unit time. :param seconds: The time period in which the rate limit is to be applied. :param message: The custom message to respond with. :param header_keys: The keys in the header to consider when apply rate limits (like 'sessionToken'). :param data_keys: The keys in the header to consider when apply rate limits (like 'sessionToken'). :param allow_if_exception: In case Redis is unresponsive, would you prefer allowing the request to pass through or would you prefer the request getting blocked. :param count_for_http_codes: If this is provided, the counter will be incremented only if the response code was one of these values. If not provided, all requests will be counted. This can be used in cases when you want to count only when the request was successfully served. :return: The decorator factory. """ # Param-cleaning: rate_limit = max(1, rate_limit) if header_keys is None: header_keys = [] if data_keys is None: data_keys = [] def decorator(func): @wraps(func) async def wrapper(*args, **kwargs): # Note down the combination of values requested and the limits prescribed: params = {"route": str(request.url_rule), "header": {}, "data": {}, "limit": rate_limit, "seconds": seconds} for key in header_keys: params["header"][key] = kwargs["inbound_headers"].get(key) for key in data_keys: params["data"][key] = kwargs["inbound_data"].get(key) # Now make a unique key from this combination: params_json = json.to_string(params, no_space = True) sha256_hash = hashlib.sha256() sha256_hash.update(params_json.encode("utf-8")) hashed_key = sha256_hash.digest() base64_key = base64.b64encode(hashed_key).decode("utf-8") # We first get the value of the counter: app_attr = getattr(current_app, attr_name) counter_value = await app_attr.count(base64_key, value = 1, expiry = seconds) # If any exception occurred in getting the count, # and exceptions haven't been allowed: if counter_value is None and not allow_if_exception: response = ResponseModel( status_code = StatusCodes.RATE_LIMIT_EXCEEDED, message = "Please contact admin (E)" ) # If the rate-limit has already been crossed, # or when the counter was not fetched but exceptions are allowed: elif (counter_value or 0) > rate_limit: response = ResponseModel( status_code = StatusCodes.RATE_LIMIT_EXCEEDED, message = message or ", ".join([ f"rate limit: {rate_limit} in {humanize.naturaldelta(datetime.timedelta(seconds = seconds))}", f"this is your {humanize.ordinal(counter_value)} request in the given period" ]) ) # If the rate-limit hasn't been crossed: else: kwargs["decorator_count"] = kwargs.get("decorator_count", 0) + 1 response = await func(*args, **kwargs) kwargs["decorator_count"] -= 1 # If we have been told to count only for specific status codes, # and if the HTTP code of this response is not in the list of codes, we reduce the counter by one: if count_for_http_codes: http_code = 200 # ... Default assumption. if isinstance(response, (list, tuple, set)): http_code = response[1] elif isinstance(response, ResponseModel): http_code = response.http_code.value if http_code not in count_for_http_codes: await app_attr.count(base64_key, value = -1, expiry = seconds) # Return the response from the function call: if ( kwargs["decorator_count"] == 0 and isinstance(response, ResponseModel) ): response = response.for_quart() return response return wrapper return decorator # --------------------------------------------------------------------------------------------------------------------- def handle_cancelled_request(cleanup_func = None, cleanup_coro = None): """ Use this decorator to handle prematurely terminated requests. If your clean-up function needs access to variables, consider using 'g' to hold data in the scope of the request. :param cleanup_func: The function to call when the cancelled request needs graceful handling. :param cleanup_coro: The coroutine to call when the cancelled request needs graceful handling. :return: The decorator factory. """ def decorator(func): @wraps(func) async def wrapper(*args, **kwargs): try: # we are ready to run the function that we are wrapping: kwargs["decorator_count"] = kwargs.get("decorator_count", 0) + 1 response = await func(*args, **kwargs) kwargs["decorator_count"] -= 1 # Done here: if ( kwargs["decorator_count"] == 0 and isinstance(response, ResponseModel) ): response = response.for_quart() return response # In case the client closes the connection pre-maturely: except asyncio.CancelledError as exception: kwargs["decorator_count"] -= 1 if hasattr(current_app, "printer"): getattr(current_app, "printer")(exception) if cleanup_func is not None: cleanup_func() if cleanup_coro is not None: await cleanup_coro() return ResponseModel( api_version = kwargs.get("api_version"), status_code = StatusCodes.CLIENT_CLOSED_REQUEST ).for_quart() # We propagate any other kind of exception: except Exception as exception: kwargs["decorator_count"] -= 1 raise return wrapper return decorator # --------------------------------------------------------------------------------------------------------------------- def handle_failed_request(cleanup_func = None, cleanup_coro = None): """ Use this decorator to handle requests that fail due to exceptions. If your clean-up function needs access to variables, consider using 'g' to hold data in the scope of the request. THIS DECORATOR IS NOT MEANT TO SUPRESS EXCEPTIONS. IT IS, INSTEAD, MEANT TO PERFORM CLEAN-UP AND PROPAGATE THE EXCEPTION. :param cleanup_func: The function to call when the cancelled request needs graceful handling. :param cleanup_coro: The coroutine to call when the cancelled request needs graceful handling. :return: The decorator factory. """ def decorator(func): @wraps(func) async def wrapper(*args, **kwargs): try: # we are ready to run the function that we are wrapping: kwargs["decorator_count"] = kwargs.get("decorator_count", 0) + 1 response = await func(*args, **kwargs) kwargs["decorator_count"] -= 1 # Done here: if ( kwargs["decorator_count"] == 0 and isinstance(response, ResponseModel) ): response = response.for_quart() return response # In case the client closes the connection pre-maturely: except Exception as exception: kwargs["decorator_count"] -= 1 if hasattr(current_app, "printer"): getattr(current_app, "printer")(exception) if cleanup_func is not None: cleanup_func() if cleanup_coro is not None: await cleanup_coro() raise exception return wrapper return decorator # ***************************************************************************************************************** # ***** **** # *** MAIN PROGRAM *** # ***** **** # ***************************************************************************************************************** if __name__ == "__main__": pass