py-app-runner 0.5.49.dev0__py3-none-any.whl

py_app_runner/bridge/websocket.py

@@ -0,0 +1,313 @@
+import logging
+import os
+import uuid
+from asyncio import Future, get_event_loop
+from typing import Any
+from urllib.parse import urlparse
+
+from tornado.websocket import WebSocketHandler as TornadoWebSocketHandler
+
+from py_app_runner.bridge.encoders import JsonEncoder, MessageEncoder, MsgpackEncoder
+from py_app_runner.config import is_env_prod
+from py_app_runner.http_exception import HTTPException
+from py_app_runner.registry import AppRegistry
+from py_app_runner.request_handler.handlers import RequestHandlerApiKeys
+from py_app_runner.return_model import MessageModel, ReturnModel, StatusModel
+from py_app_runner.utils import json_encode
+from py_app_runner.wbcm.ws_interface import WebSocketHandlerInterface
+
+
+def allowed_ws_origins() -> set[str]:
+    """Origins accepted for WebSocket upgrades.
+
+    Read on demand rather than at import time: the value comes from the app config
+    (with an env fallback), and neither is populated when this module is imported.
+    An empty set means "accept any origin".
+    """
+    configured = AppRegistry.config().get("ws_allowed_origins")
+    if configured is None:
+        configured = os.environ.get("WS_ALLOWED_ORIGINS", "")
+
+    if isinstance(configured, str):
+        configured = configured.split(",")
+
+    return {str(o).strip().lower().rstrip("/") for o in configured if str(o).strip()}
+
+
+# * BaseWebSocketHandler - protocol-agnostic WebSocket handler
+class BaseWebSocketHandler(RequestHandlerApiKeys, TornadoWebSocketHandler, WebSocketHandlerInterface):
+    """Base WebSocket request handler with pluggable encoding."""
+
+    encoder: MessageEncoder
+
+    #########################
+    ### Class lifecycle #####
+    #########################
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        logger_name = f"{__name__}.{self.__class__.__name__}"
+        self.logger = logging.getLogger(logger_name)
+
+        self.loop = get_event_loop()
+
+        self.uid = None
+        self.msg_id = None
+        self.service = None
+        self.current_session_sid = None
+        self.device_id: str | None = None
+        self._api_key_valid: bool | None = None
+        self._rejected_api_key: str | None = None
+
+        super().__init__(*args, **kwargs)
+
+    ###############
+    ### Helpers ###
+    ###############
+
+    def error_message(self, msg: str | HTTPException, code: int = -1, **kwargs: Any):
+        error: dict[str, Any]
+        if isinstance(msg, HTTPException):
+            error = msg.to_dict()
+        else:
+            error = {"code": code, "msg": msg}
+
+        message: dict[str, Any] = {"error": error}
+        message.update(kwargs)
+        self.write_custom_message(message)
+
+    #################
+    ### Overrides ###
+    #################
+
+    def write_custom_message(
+        self,
+        message: bytes | str | dict[str, Any] | list[Any],
+        wrap_in_data: bool = True,
+        **kwargs: Any,
+    ) -> Future[None]:
+        if isinstance(message, (StatusModel, MessageModel, ReturnModel)):
+            message = message.to_dict()
+        if wrap_in_data:
+            message = {"data": message}
+
+        response_message: str | bytes
+        if isinstance(message, dict):
+            # Add service
+            if self.service and "service" not in message:
+                message["service"] = self.service
+
+            # Add msg_id
+            if self.msg_id and "msg_id" not in message:
+                message["msg_id"] = self.msg_id
+
+            # Add additional fields
+            message.update(kwargs)
+
+            # Encode message using the pluggable encoder
+            response_message = self.encoder.encode(message)
+
+        elif isinstance(message, (str, bytes)):
+            # Already-encoded payload, pass through untouched
+            response_message = message
+
+        else:
+            response_message = self.encoder.encode(message)  # type: ignore[arg-type]
+
+        if self.ws_connection is not None and self.ws_connection.is_closing() is False:
+            return super().write_message(response_message, binary=self.encoder.is_binary)  # type: ignore
+
+        # Resolved, not dangling: callers may await the result.
+        dropped: Future[None] = Future()
+        dropped.set_result(None)
+        return dropped
+
+    ########################
+    ### Request handling ###
+    ########################
+
+    async def reload_token(self, new_auth_token: str) -> None:
+        old_auth_token = self.auth_token
+        if old_auth_token != new_auth_token:
+            self.auth_token = new_auth_token
+            await self._ensure_current_user()
+
+    def check_origin(self, origin: str) -> bool:
+        allowed = allowed_ws_origins()
+        if not allowed:
+            if is_env_prod():
+                self.logger.warning(
+                    "ws_allowed_origins is not configured; accepting WebSocket upgrade from origin %s."
+                    " Set it to prevent cross-site WebSocket hijacking.",
+                    origin,
+                )
+            return True
+
+        normalized = origin.lower().rstrip("/")
+        if normalized in allowed:
+            return True
+
+        # Also check just the scheme + host (ignoring path)
+        parsed = urlparse(normalized)
+        origin_host = f"{parsed.scheme}://{parsed.netloc}" if parsed.scheme and parsed.netloc else normalized
+        return origin_host in allowed
+
+    async def open(self, *args: str, **kwargs: str) -> None:
+        # Cache the connection
+        self.uid = str(uuid.uuid4())
+        self.wb_connection_manager.user_connections.add_connection(self)
+
+        self.logger.debug(
+            "WS open uid=%s remote_ip=%s origin=%s",
+            self.uid,
+            self.request.remote_ip,
+            self.request.headers.get("Origin", "<none>"),
+        )
+
+        # Validate API key from query param or header at connection time
+        status = await self.has_valid_api_key()
+        if status is True:
+            self._api_key_valid = True
+        elif status == "API key is missing":
+            self._api_key_valid = None
+        else:
+            self._api_key_valid = False
+            self.logger.debug("WS closing uid=%s reason=api_key_failed status=%s", self.uid, status)
+            self.close(4401, f"API key validation failed: {status}")
+
+    async def on_message(self, message: str | bytes) -> None:
+        self.logger.debug(f"Received message ({self.encoder.name}): {message!r}")
+        try:
+            message_data = self.encoder.decode(message)
+        except Exception:
+            msg_str = "It was impossible to parse the given data"
+            self.logger.exception(f"Error: {msg_str}\nData: {message!r}")
+            self.error_message(msg_str, code=1000)
+            return
+
+        try:
+            # Assign message id
+            msg_id = message_data.get("msg_id", None)
+            service = message_data.get("service", None)
+
+            if not msg_id or not service:
+                raise HTTPException(
+                    f"Missing input data; Data received: {json_encode(message_data, pretty=True)}",
+                    code=1001,
+                    http_status=400,
+                )
+
+            self.msg_id = int(msg_id)
+            self.service = service
+            res = self.find_service(service_name=service)
+            if not res.result:
+                raise HTTPException(
+                    f"Could not find service by provided name: {self.service}",
+                    code=1002,
+                    http_status=400,
+                )
+
+            bridge_request, requires_api_key = res.result
+            if not bridge_request:
+                raise HTTPException(f"Service not found: {self.service}", code=1003, http_status=400)
+
+            # Check cached API key validation from connection handshake;
+            # fall back to in-message api_key (browsers cannot send custom headers on WS)
+            if requires_api_key is not False:
+                if self._api_key_valid is None:
+                    in_msg_api_key = message_data.get("api_key", None)
+                    # Remember only the last rejected key, so retrying the same bad key
+                    # costs nothing while a corrected one can still be accepted.
+                    if in_msg_api_key and in_msg_api_key != self._rejected_api_key:
+                        status = await self.has_valid_api_key(in_msg_api_key)
+                        if status is True:
+                            self._api_key_valid = True
+                        else:
+                            self._rejected_api_key = in_msg_api_key
+
+                if self._api_key_valid is not True:
+                    raise HTTPException(
+                        "Application is not authenticated: API key is missing or invalid",
+                        code=1401,
+                        http_status=401,
+                    )
+
+            # Set auth token
+            auth_token = message_data.get("auth_token", None)
+            if auth_token:
+                await self.reload_token(auth_token)
+
+            # Device session token auth
+            device_session_token = message_data.get("device_session_token", None)
+            if device_session_token and not self.device_id:
+                await self._authenticate_device(device_session_token)
+
+            # Timer
+            async with self.timer.aenter("bridge.websockets.on_message.runRequestHandler"):
+                req_data = message_data.get("data", {})
+                action = req_data.get("action", None)
+                input_data = req_data.get("data", {})
+
+                if not action:
+                    raise HTTPException("Missing action", code=1005, http_status=400)
+
+                # Run request handler
+                return_data = await bridge_request(action, input_data, self)
+                if return_data is not None:
+                    self.write_custom_message(return_data)
+
+        except HTTPException as e:
+            self.log_request(error=e)
+            self.error_message(str(e.message), e.code or -1)
+            return
+
+        except Exception as e:
+            self.log_request(error=e)
+            self.logger.exception(f"Error processing request with exception: {e}")
+            self.error_message(
+                "Found an error while processing your request. Please try again later."
+                " Send us a message if the problem persists."
+            )
+            return
+
+        finally:
+            self.msg_id = None
+            self.service = None
+
+            self.timer.print_timer_stats()
+            self.timer.reset_timers()
+
+    async def _authenticate_device(self, session_token: str) -> None:
+        """Validate a device session token and register the device connection.
+        Override in subclasses to implement project-specific device authentication."""
+        pass
+
+    def on_close(self):
+        self.logger.debug(
+            "WS close uid=%s code=%s reason=%s",
+            self.uid,
+            self.close_code,
+            self.close_reason or "<none>",
+        )
+        if self.uid:
+            self.wb_connection_manager.user_connections.remove_connection(self)
+            if self.device_id:
+                self.wb_connection_manager.device_connections.remove_connection(self)
+                self.device_id = None
+            self.uid = None
+        self.auth_token = None
+        self.current_user = None
+        self.current_session_sid = None
+
+
+# * WebSocketHandler - JSON WebSocket handler (backward compatible)
+class WebSocketHandler(BaseWebSocketHandler):
+    """JSON WebSocket handler at /v1/socket."""
+
+    encoder = JsonEncoder()
+
+
+# * MsgpackWebSocketHandler - MessagePack WebSocket handler
+class MsgpackWebSocketHandler(BaseWebSocketHandler):
+    """MessagePack WebSocket handler at /v1/socket/msgpack."""
+
+    encoder = MsgpackEncoder()

py_app_runner/colors.py

@@ -0,0 +1,73 @@
+class Colors:
+    # Reset
+    ColorOff = "\033[0m"  # Text Reset
+
+    # Regular Colors
+    Black = "\033[0;30m"  # Black
+    Red = "\033[0;31m"  # Red
+    Green = "\033[0;32m"  # Green
+    Yellow = "\033[0;33m"  # Yellow
+    Blue = "\033[0;34m"  # Blue
+    Purple = "\033[0;35m"  # Purple
+    Cyan = "\033[0;36m"  # Cyan
+    White = "\033[0;37m"  # White
+
+    # Bold
+    BBlack = "\033[1;30m"  # Black
+    BRed = "\033[1;31m"  # Red
+    BGreen = "\033[1;32m"  # Green
+    BYellow = "\033[1;33m"  # Yellow
+    BBlue = "\033[1;34m"  # Blue
+    BPurple = "\033[1;35m"  # Purple
+    BCyan = "\033[1;36m"  # Cyan
+    BWhite = "\033[1;37m"  # White
+
+    # Underline
+    UBlack = "\033[4;30m"  # Black
+    URed = "\033[4;31m"  # Red
+    UGreen = "\033[4;32m"  # Green
+    UYellow = "\033[4;33m"  # Yellow
+    UBlue = "\033[4;34m"  # Blue
+    UPurple = "\033[4;35m"  # Purple
+    UCyan = "\033[4;36m"  # Cyan
+    UWhite = "\033[4;37m"  # White
+
+    # Background
+    BgBlack = "\033[40m"  # Black
+    BgRed = "\033[41m"  # Red
+    BgGreen = "\033[42m"  # Green
+    BgYellow = "\033[43m"  # Yellow
+    BgBlue = "\033[44m"  # Blue
+    BgPurple = "\033[45m"  # Purple
+    BgCyan = "\033[46m"  # Cyan
+    BgWhite = "\033[47m"  # White
+
+    # High Intensty
+    IBlack = "\033[0;90m"  # Black
+    IRed = "\033[0;91m"  # Red
+    IGreen = "\033[0;92m"  # Green
+    IYellow = "\033[0;93m"  # Yellow
+    IBlue = "\033[0;94m"  # Blue
+    IPurple = "\033[0;95m"  # Purple
+    ICyan = "\033[0;96m"  # Cyan
+    IWhite = "\033[0;97m"  # White
+
+    # Bold High Intensty
+    BIBlack = "\033[1;90m"  # Black
+    BIRed = "\033[1;91m"  # Red
+    BIGreen = "\033[1;92m"  # Green
+    BIYellow = "\033[1;93m"  # Yellow
+    BIBlue = "\033[1;94m"  # Blue
+    BIPurple = "\033[1;95m"  # Purple
+    BICyan = "\033[1;96m"  # Cyan
+    BIWhite = "\033[1;97m"  # White
+
+    # High Intensty backgrounds
+    BgIBlack = "\033[0;100m"  # Black
+    BgIRed = "\033[0;101m"  # Red
+    BgIGreen = "\033[0;102m"  # Green
+    BgIYellow = "\033[0;103m"  # Yellow
+    BgIBlue = "\033[0;104m"  # Blue
+    BgIPurple = "\033[10;95m"  # Purple
+    BgICyan = "\033[0;106m"  # Cyan
+    BgIWhite = "\033[0;107m"  # White

py_app_runner/config.py

@@ -0,0 +1,132 @@
+import logging
+import sys
+from copy import deepcopy
+from os import environ, getcwd, path
+from typing import Any
+
+from dotenv import load_dotenv
+
+
+def coerce_value(value: Any) -> Any:
+    """Turn a plain integer string into an int, leaving everything else alone.
+
+    The round-trip check keeps values where the text form carries meaning - a zero
+    padded id or a numeric secret must not silently become a different number.
+    """
+    if isinstance(value, list) or not isinstance(value, str):
+        return value
+
+    if value.isdigit() and str(int(value)) == value:
+        return int(value)
+
+    return value
+
+
+def replace_rec(
+    keys: list[str],
+    finalDict: dict[str, Any],
+    finalValue: Any,
+) -> dict[str, Any]:
+    key = keys.pop(0)
+    if len(keys) == 0:
+        finalDict[key] = coerce_value(finalValue)
+        return finalDict
+
+    if key not in finalDict:
+        finalDict[key] = {}
+
+    # An env var whose key path runs through an existing scalar cannot be merged
+    # (e.g. API_KEY_PEPPER against a config that already has api.key). Overwriting
+    # the scalar would silently destroy it, so refuse instead.
+    if not isinstance(finalDict[key], dict):
+        raise ValueError(f"Cannot descend into non-dict config key '{key}': it already holds a scalar value")
+
+    finalDict[key] = replace_rec(keys, finalDict[key], finalValue)
+    return finalDict
+
+
+def parse_splitted_values(value: str, separator: str = ",") -> list[str]:
+    valueList = value.split(separator)
+    formattedList: list[str] = []
+    for item in valueList:
+        item = item.strip(" \r\n\t\\")
+        if len(item) == 0 or item[0] == "#":
+            continue
+
+        formattedList.append(item)
+
+    return formattedList
+
+
+def is_env(env: str, config: dict[str, Any] | None = None) -> bool:
+    if config is None:
+        from py_app_runner.registry import AppRegistry
+
+        config = AppRegistry.config()
+    return env == config.get("environment", "")
+
+
+def is_env_dev(config: dict[str, Any] | None = None) -> bool:
+    return is_env("dev", config)
+
+
+def is_env_test(config: dict[str, Any] | None = None) -> bool:
+    return is_env("test", config)
+
+
+def is_env_prod(config: dict[str, Any] | None = None) -> bool:
+    return is_env("prod", config)
+
+
+def load_config(
+    defaults: dict[str, Any] | None = None,
+    split_value_keys: list[str] | None = None,
+) -> dict[str, Any]:
+    """
+    Load configuration from environment variables.
+
+    1. Calls load_dotenv() from the current working directory
+    2. Takes a `defaults` dict (the project's base config structure)
+    3. Applies replace_rec() to fill in from env vars
+    4. Handles split_value_keys (CSV splitting)
+    5. Returns the populated dict
+    """
+    currentDirectory = getcwd()
+    load_dotenv(path.join(currentDirectory, ".env"))
+
+    appEnv = environ.get("APP_ENV", None)
+    if not appEnv:
+        logging.getLogger().error("*\n*\n* It seems that `.env` file is not loaded\n*\n*")
+        sys.exit(-1)
+
+    if defaults is None:
+        defaults = {}
+
+    if split_value_keys is None:
+        split_value_keys = []
+
+    # Deep copy: replace_rec writes into nested dicts, which a shallow copy would
+    # share with the caller's (often module level) defaults.
+    config_dict: dict[str, Any] = deepcopy(defaults)
+
+    # Always set environment from APP_ENV
+    config_dict["environment"] = appEnv
+
+    for key in environ:
+        keys = key.split("_")
+        keys = [item.lower() for item in keys]
+        value = environ[key]
+
+        if len(keys) == 0 or config_dict.get(keys[0], None) is None:
+            continue
+
+        new_value: Any = parse_splitted_values(value) if key in split_value_keys else value
+
+        try:
+            config_dict = replace_rec(keys, config_dict, new_value)
+        except ValueError as e:
+            # One unmappable variable must not take the whole service down, but it
+            # must be loud: the value the operator set is not in effect.
+            logging.getLogger(__name__).error("Ignoring env var %s: %s", key, e)
+
+    return config_dict

py_app_runner/crypto/__init__.py

@@ -0,0 +1,14 @@
+"""Field encryption, the shared client envelope, and password hashing."""
+
+from py_app_runner.crypto.errors import CryptoError
+from py_app_runner.crypto.fields import FieldCrypto, generate_key, is_encrypted, key_id_of
+from py_app_runner.crypto.passwords import PasswordHasher
+
+__all__ = [
+    "CryptoError",
+    "FieldCrypto",
+    "PasswordHasher",
+    "generate_key",
+    "is_encrypted",
+    "key_id_of",
+]

py_app_runner/crypto/_service.py

@@ -0,0 +1,75 @@
+import asyncio
+import logging
+from argparse import Namespace
+
+import psycopg
+
+from py_app_runner.crypto.commands import Out, cmd_key, cmd_rotate
+from py_app_runner.crypto.errors import CryptoError
+from py_app_runner.crypto.fields import FieldCrypto
+from py_app_runner.migrations._service import connect_kwargs
+from py_app_runner.pybridge import PyBridge
+from py_app_runner.registry import AppRegistry
+
+_DEFAULT_DB = "main"
+
+
+async def init_service(args: Namespace, _pybridge: PyBridge, logger: logging.Logger) -> None:
+    out: Out = print
+
+    # `key` runs before any config is read and any connection is opened. Generating the
+    # first key must not require a working install, because until it exists there isn't
+    # one - the environment variable it goes into is what the config refers to.
+    if args.step == "key":
+        raise SystemExit(cmd_key(out))
+
+    # Same reasoning as migrations: runner.py catches Exception around init_service and
+    # returns normally, which exits 0. A rotate that could not reach the database must not
+    # report success to a deploy script, so every non-SystemExit failure becomes a non-zero
+    # exit here rather than in runner.py.
+    code = 1
+    try:
+        config = AppRegistry.config()
+        crypto = FieldCrypto.from_config(config)
+
+        db_name = getattr(args, "db", None) or _DEFAULT_DB
+        db_config = config.get("db") or {}
+        if db_name not in db_config:
+            out(
+                f'error: no database {db_name!r} in config["db"]; '
+                f"configured are: {', '.join(sorted(db_config)) or 'none'}."
+            )
+            raise SystemExit(2)
+
+        if args.step == "rotate":
+            async with await psycopg.AsyncConnection.connect(**connect_kwargs(db_config[db_name])) as conn:
+                code = await cmd_rotate(
+                    conn,
+                    crypto,
+                    args.table,
+                    args.column,
+                    args.id,
+                    args.batch,
+                    getattr(args, "dry_run", False),
+                    out,
+                )
+        else:
+            out(f"error: unknown crypto command {args.step!r}")
+            code = 1
+
+    except CryptoError as e:
+        # Configuration and key material problems are the expected failure here, and their
+        # messages already say what to do. A stack trace would bury that.
+        out(f"error: {e}")
+        raise SystemExit(1) from None
+    except (KeyboardInterrupt, asyncio.CancelledError):
+        # Above `except Exception` because both derive from BaseException. An interrupted
+        # rotate has re-encrypted some rows and not others, which is safe to resume but must
+        # not be reported as done.
+        logger.error("crypto: interrupted; the column may be partially rotated")
+        raise SystemExit(1) from None
+    except Exception:
+        logger.exception("crypto: unhandled failure")
+        raise SystemExit(1) from None
+
+    raise SystemExit(code)

py_app_runner/crypto/_service_args.py

@@ -0,0 +1,54 @@
+"""CLI subparsers for the built-in crypto service.
+
+    python3 src/app.py crypto key
+    python3 src/app.py crypto rotate --table NAME --column NAME [--id NAME]
+                                     [--batch N] [--dry-run] [--db NAME]
+
+`--dry-run` shadows a real top-level flag on runner.py's parser, so it needs
+`default=SUPPRESS` or argparse copies the subparser default back over the parent
+namespace - the same collision migrations' `apply` has.
+"""
+
+import logging
+from argparse import SUPPRESS, ArgumentParser, _SubParsersAction  # type: ignore
+
+from py_app_runner.pybridge import PyBridge
+
+
+def reg_subparsers(
+    subparsers: "_SubParsersAction[ArgumentParser]",
+    _pybridge: PyBridge,
+    _base_logger: logging.Logger,
+) -> None:
+    """Command line subparsers"""
+
+    parser = subparsers.add_parser(
+        "crypto",
+        description="Generate key material and re-encrypt stored columns",
+        help="Field encryption",
+    )
+    group = parser.add_subparsers(title="command", dest="step", required=True)
+
+    group.add_parser("key", help="Print fresh key material for an environment variable")
+
+    rotate_parser = group.add_parser("rotate", help="Re-encrypt a column under the current key")
+    rotate_parser.add_argument("--table", required=True, help="Table holding the column")
+    rotate_parser.add_argument("--column", required=True, help="Encrypted column")
+    rotate_parser.add_argument("--id", default="id", help="Primary key to page through (default: id)")
+    rotate_parser.add_argument(
+        "--batch",
+        type=int,
+        default=500,
+        help="Rows read per statement (default: 500)",
+    )
+    rotate_parser.add_argument(
+        "--dry-run",
+        action="store_true",
+        default=SUPPRESS,
+        help="Report what would change, change nothing",
+    )
+    rotate_parser.add_argument(
+        "--db",
+        default=None,
+        help='Entry of config["db"] to rotate in (default: main)',
+    )