Source code for pyprediktorutilities.shared

import requests
from pydantic import AnyUrl, validate_call
from typing import Literal
import logging
import os
import re
import ipaddress
from pathlib import Path
from urllib.parse import urlparse, urlunparse

logger = logging.getLogger(__name__)
logger.addHandler(logging.NullHandler())

[docs] @validate_call def validate_folder(folder: str): if not Path(folder).is_dir(): errormsg = f"Folder {folder} does not exist" logging.error(errormsg) raise FileNotFoundError(errormsg)
[docs] @validate_call def validate_file(file: str): if not Path(file).is_file(): errormsg = f"File {file} does not exist" logging.error(errormsg) raise FileNotFoundError(errormsg)
[docs] @validate_call def validate_safe_path(base_directory: str, filename: str) -> str: """Resolve `filename` relative to `base_directory` and ensure the result does not escape `base_directory`. Args: base_directory (str): The directory the file is expected to reside in filename (str): The (possibly untrusted) file name/path to validate Raises: ValueError: If the resolved path escapes base_directory Returns: str: The resolved, safe absolute path """ base = Path(base_directory).resolve() target = (base / filename).resolve() if target != base and base not in target.parents: errormsg = f"'{filename}' is not a valid file name" logging.error(errormsg) raise ValueError(errormsg) return str(target)
[docs] @validate_call def sanitize_file_path(file: str) -> str: """Sanitize a file path immediately before it is opened. Rejects NUL bytes and ".." path traversal segments, then returns a normalized, fully-resolved absolute path. Args: file (str): The file path to sanitize Raises: ValueError: If the path contains a NUL byte or a ".." traversal segment Returns: str: The normalized, absolute file path """ if "\x00" in file: errormsg = f"'{file}' contains a NUL byte and was rejected" logging.error(errormsg) raise ValueError(errormsg) normalized = os.path.normpath(file) parts = normalized.split(os.sep) if os.pardir in parts: errormsg = f"'{file}' contains a path traversal sequence and was rejected" logging.error(errormsg) raise ValueError(errormsg) return os.path.abspath(normalized)
_DISALLOWED_SSRF_HOSTNAMES = {"localhost", "metadata.google.internal"}
[docs] def validate_safe_url(url: str) -> str: """Guard against SSRF by rejecting obviously unsafe outbound request targets. Only allows the http/https schemes, rejects a small denylist of well-known internal hostnames, and rejects requests to loopback, private, link-local, reserved, multicast or unspecified IP addresses (this covers the common cloud metadata endpoint 169.254.169.254 when it is used directly as the host). But this performs no DNS resolution (to keep it fast and avoid depending on network access), so it does not protect against DNS rebinding attacks where a hostname resolves to an internal address only at connection time. Args: url (str): The fully-formed URL that a request is about to be made to Raises: ValueError: If the URL uses a disallowed scheme or targets a disallowed host Returns: str: The validated URL, unchanged """ parsed = urlparse(str(url)) if parsed.scheme not in ("http", "https"): errormsg = f"URL scheme '{parsed.scheme}' is not allowed" logging.error(errormsg) raise ValueError(errormsg) hostname = parsed.hostname if not hostname: errormsg = f"URL '{url}' does not have a valid hostname" logging.error(errormsg) raise ValueError(errormsg) if hostname.lower() in _DISALLOWED_SSRF_HOSTNAMES: errormsg = f"URL '{url}' targets a disallowed host ({hostname})" logging.error(errormsg) raise ValueError(errormsg) try: ip = ipaddress.ip_address(hostname) except ValueError: return str(url) if ( ip.is_private or ip.is_loopback or ip.is_link_local or ip.is_reserved or ip.is_multicast or ip.is_unspecified ): errormsg = f"URL '{url}' targets a disallowed address ({ip})" logging.error(errormsg) raise ValueError(errormsg) return str(url)
[docs] def build_validated_url(base_url: str, endpoint: str) -> str: """Build a validated URL by safely appending an endpoint to a base URL. Rejects endpoint values containing "../" (or its URL-encoded form) before the URL is ever assembled, then rebuilds the URL via `urlparse`/`urlunparse`. Args: base_url (str): The base URL (scheme, host, and optional base path) endpoint (str): The endpoint path to append Raises: ValueError: If the endpoint contains a path traversal sequence, or the URL could not be built Returns: str: The safely constructed URL """ try: if "../" in endpoint or re.search(r"%2e%2e/", endpoint, re.IGNORECASE): raise ValueError("Invalid path") parsed = urlparse(base_url) base_path = parsed.path.rstrip("/") endpoint_clean = endpoint.lstrip("/") new_path = f"{base_path}/{endpoint_clean}" if endpoint_clean else base_path parsed = parsed._replace(path=new_path) return urlunparse(parsed) except Exception: errormsg = f"Could not build a valid URL from base '{base_url}' and endpoint '{endpoint}'" logging.error(errormsg) raise ValueError(errormsg)
[docs] @validate_call def request_from_api( rest_url: AnyUrl, method: Literal["GET", "POST"], endpoint: str, data: str = None, params: dict = None, headers: dict = None, extended_timeout: bool = False, ) -> str: """Function to perform request against a REST API Args: rest_url (str): The URL with trailing shash method (str): "GET" or "POST" endpoint (str): The last part of the url (without the leading slash) data (str): defaults to None but can contain the data to send to the endpoint headers (str): default to None but can contain the headers og the request Raises: ValueError: If the resulting URL uses a disallowed scheme or targets a disallowed (private/loopback/link-local/etc.) address Returns: JSON: The result if successfull """ request_timeout = (3, 300 if extended_timeout else 27) combined_url = build_validated_url(str(rest_url), endpoint) combined_url = validate_safe_url(combined_url) if method == "GET": result = requests.get( combined_url, timeout=request_timeout, params=params, headers=headers, allow_redirects=False ) if method == "POST": result = requests.post( combined_url, data=data, headers=headers, timeout=request_timeout, params=params, allow_redirects=False ) result.raise_for_status() return result.json()