def otp_redis(self) -> Redis:
"""
Return the Redis manager to store OtpData.
:return: the Redis manager to store OtpData.
:raise: ImmuniException if the manager is not initialized.
"""
if self._otp_redis is None:
raise ImmuniException("Cannot use the Redis manager before initialising it.")
return self._otp_redis
def exposure_mongo(self) -> MongoClient:
"""
Return the MongoDB manager to handle TEKs.
:return: the MongoDB manager to handle TEKs.
:raise: ImmuniException if the manager is not initialized.
"""
if self._exposure_mongo is None:
raise ImmuniException(
"Cannot use the MongoDB manager before initializing it.")
return self._exposure_mongo
def celery_redis(self) -> Redis:
"""
Return the Celery Redis manager to store analytics.
:return: the Celery Redis manager to store analytics.
:raise: ImmuniException if the manager is not initialized.
"""
if self._celery_redis is None:
raise ImmuniException(
"Cannot use the Celery Redis manager before initialising it.")
return self._celery_redis
def app_configuration_mongo(self) -> MongoClient:
"""
Return the MongoDB manager to handle app configurations.
:return: the MongoDB manager to handle app configurations.
:raise: ImmuniException if the manager is not initialized.
"""
if self._app_configuration_mongo is None:
raise ImmuniException(
"Cannot use the MongoDB manager before initializing it.")
return self._app_configuration_mongo
def from_env_var(cls: Type[T], value: str) -> T:
"""
Parse the environment variable value and provide an informative error message on failure.
:param value: the environment variable value.
:return: the corresponding Enum entry.
"""
try:
return cls(value) # type: ignore
except ValueError:
allowed = ", ".join(e.value for e in cls) # type: ignore
raise ImmuniException(
f"Invalid environment: {value} (allowed: {allowed})")
def weighted_random(pairs: List[WeightedPayload[T]]) -> T:
"""
Returns one of the values in the WeightedPair list randomly based on the
weights defined in the given WeightedPair list.
:param pairs: The list of WeightedPair to pick the random value from.
"""
# Note: We allow 0 weights so that this function is testable and tests are not random.
if any(pair.weight < 0 for pair in pairs):
raise ImmuniException("Cannot perform a weighted random with negative weights.")
return random.choices(
population=tuple(p.payload for p in pairs), weights=tuple(p.weight for p in pairs), k=1,
)[0]
def from_env_var(cls, value: str) -> LogLevel:
"""
Parse the environment variable value and provide an informative error message on failure.
:param value: the environment variable value.
:return: the corresponding LogLevel entry.
:raises: ImmuniException if the given log level is deprecated in the logging library.
"""
environment = super().from_env_var(value)
if environment.name not in logging._nameToLevel: # pylint: disable=protected-access
raise ImmuniException(
f"Deprecated log level: {value}. Code update needed.")
return environment
async def _wrapper(*args: Any, **kwargs: Any) -> HTTPResponse:
response = await f(*args, **kwargs)
if not config.CACHE_ENABLED:
return response
if _CACHE_CONTROL in response.headers:
raise ImmuniException(
f"Attempt to redefine {_CACHE_CONTROL} headers.")
# NOTE: Check for both defined or both undefined has already been done outside.
if max_age:
response.headers[
_CACHE_CONTROL] = f"public, max-age={int(max_age.total_seconds())}"
if no_store:
response.headers[_CACHE_CONTROL] = "no-store"
return response
def cache(max_age: Optional[timedelta] = None,
no_store: Optional[bool] = None) -> Callable:
"""
Decorator to add cache headers to an endpoint response.
:param max_age: the timedelta, if any, after which the cache would expire.
:param no_store: True if the response shall not be cached.
:return: the decorator.
:raises:
ImmuniException: if none or all arguments are defined.
ImmuniException: if attempting to redefine the Cache-Control header.
"""
def _decorator(
f: Callable[..., Awaitable[HTTPResponse]]
) -> Callable[..., Awaitable[HTTPResponse]]:
@wraps(f)
async def _wrapper(*args: Any, **kwargs: Any) -> HTTPResponse:
response = await f(*args, **kwargs)
if not config.CACHE_ENABLED:
return response
if _CACHE_CONTROL in response.headers:
raise ImmuniException(
f"Attempt to redefine {_CACHE_CONTROL} headers.")
# NOTE: Check for both defined or both undefined has already been done outside.
if max_age:
response.headers[
_CACHE_CONTROL] = f"public, max-age={int(max_age.total_seconds())}"
if no_store:
response.headers[_CACHE_CONTROL] = "no-store"
return response
return _wrapper
if not bool(max_age) ^ bool(no_store):
raise ImmuniException(
f"{cache.__name__} decorator arguments are mutually exclusive, and at least one shall "
f"be defined (max_age: {max_age}, no_store: {no_store}).")
return _decorator
async def _wrapper(request: Request, *args: Any,
**kwargs: Any) -> HTTPResponse:
data = getattr(request, location.value, {})
if location == Location.HEADERS:
data = _remap_data_keys(data)
elif location == Location.QUERY:
_validate_query_args_length(data)
elif location == Location.JSON:
_validate_json_content_type(request)
schema = Schema.from_dict(fields)
try:
valid_data = schema().load(data) # pylint: disable=no-member
except ValidationError as exc:
raise SchemaValidationException(exc.messages) from exc
if intersection := set(kwargs.keys()).intersection(
set(valid_data.keys())):
raise ImmuniException(
f"Trying to validate some fields more than once: {list(intersection)}."
)
def _validate_crontab(value: str) -> str:
if not croniter.is_valid(value):
raise ImmuniException(
f"Invalid crontab string for {config_name}: {value}.")
return value