rsconnect-python 1.30.0__py3-none-any.whl

rsconnect/json_web_token.py

@@ -0,0 +1,178 @@
+"""
+Json Web Token (JWT) utilities
+"""
+
+from __future__ import annotations
+
+import base64
+import binascii
+import os
+from datetime import datetime, timedelta, timezone
+from typing import Any, Optional
+
+import jwt
+
+from .exception import RSConnectException
+from .http_support import HTTPResponse, JsonData
+from .models import BootstrapOutputDTO
+
+DEFAULT_ISSUER = "rsconnect-python"
+DEFAULT_AUDIENCE = "rsconnect"
+
+BOOTSTRAP_SCOPE = "bootstrap"
+BOOTSTRAP_EXP = timedelta(minutes=15)
+
+SECRET_KEY_ENV = "CONNECT_BOOTSTRAP_SECRETKEY"
+
+
+def read_secret_key(keypath: Optional[str]) -> bytes:
+    """
+    Reads a secret key as bytes given a path to a file containing a base64-encoded key.
+
+    The secret key can optionally be set with an environment variable.
+    """
+
+    env_raw_data = os.getenv(SECRET_KEY_ENV)
+
+    if keypath is not None and env_raw_data is not None:
+        raise RSConnectException("Cannot specify secret key using both a keyfile and environment variable.")
+
+    if keypath is None and env_raw_data is None:
+        raise RSConnectException("Must specify secret key using either a keyfile or environment variable.")
+
+    # check if secret key was specified using an env variable first
+    if env_raw_data is not None:
+        try:
+            return base64.b64decode(env_raw_data.encode("utf-8"))
+        except binascii.Error:
+            raise RSConnectException("Unable to decode base64 data from environment variable: " + SECRET_KEY_ENV)
+
+    if keypath is None:
+        raise RSConnectException("Keypath must not be None.")
+
+    if not os.path.exists(keypath):
+        raise RSConnectException("Keypath does not exist.")
+
+    with open(keypath, "r") as f:
+        raw_data = f.read()
+        if raw_data is None:
+            raise RSConnectException("Secret key cannot be 'None'")
+        try:
+            return base64.b64decode(raw_data)
+        except binascii.Error:
+            raise RSConnectException("Unable to decode base64 data from keyfile: " + keypath)
+
+
+# https://www.ibm.com/docs/vi/sva/9.0.6?topic=jwt-support
+def validate_hs256_secret_key(key: bytes):
+    if len(key) < 32:
+        raise RSConnectException("Secret key expected to be at least 32 bytes in length")
+
+
+def parse_client_response(response: BootstrapOutputDTO | HTTPResponse) -> tuple[int, BootstrapOutputDTO | JsonData]:
+    """
+    Helper to handle the response type from RSConnectClient, because
+    it can have different types depending on the response
+    """
+
+    if isinstance(response, dict):
+        return 200, response
+    elif isinstance(response, HTTPResponse):
+        # fail fast if a non-http exception occurred
+        if hasattr(response, "exception") and response.exception is not None:
+            raise RSConnectException(str(response.exception))
+
+        status = 500
+        if hasattr(response, "status"):
+            status = response.status
+
+        json_data: JsonData = {}
+        if hasattr(response, "json_data"):
+            json_data = response.json_data
+
+        return status, json_data
+
+    raise RSConnectException("Unrecognized response type: " + str(type(response)))
+
+
+def produce_bootstrap_output(status: int, json_data: BootstrapOutputDTO | JsonData) -> dict[str, int | str]:
+    """
+    Produces the expected programmatic output format from a request to the initial_admin endpoint
+    """
+
+    # Parse the returned API key if one is provided
+    api_key = ""
+    if isinstance(json_data, dict) and "api_key" in json_data:
+        api_key = json_data["api_key"]
+        if not isinstance(api_key, str):
+            raise RSConnectException("Connect returned a non-string value for api_key.")
+
+    # Catch unexpected response states and error early
+    if status == 200 and api_key == "":
+        raise RSConnectException("Connect returned a successful HTTP response but no API key.")
+
+    if status != 200 and api_key != "":
+        raise RSConnectException("Connect returned a non-successful HTTP response and an API key. ")
+
+    output = {"status": status, "api_key": api_key}
+
+    # Create a helpful error message
+    message = "Unexpected response status."
+    if status == 200:
+        message = "Success."
+    elif status == 401:
+        message = "JWT authorization failed."
+    elif status == 403:
+        message = "Unable to provision initial admin. Please check status of Connect database."
+    elif status == 404:
+        message = (
+            "Unable to find provisioning endpoint. Please check your 'rsconnect bootstrap --server' "
+            "parameter and your Connect configuration."
+        )
+
+    output["message"] = message
+
+    return output
+
+
+class JWTEncoder:
+    def __init__(self, issuer: str, audience: str, secret: str | bytes):
+        self.issuer = issuer
+        self.audience = audience
+        self.secret = secret
+
+    def generate_standard_claims(self, current_datetime: datetime, exp: timedelta):
+
+        if exp < timedelta(0):
+            raise RSConnectException("Unable to generate a token with a negative exp claim.")
+
+        return {
+            "exp": int((current_datetime + exp).timestamp()),
+            "iss": self.issuer,
+            "aud": self.audience,
+            "iat": int(current_datetime.timestamp()),
+        }
+
+    def new_token(self, custom_claims: dict[str, Any], exp: timedelta) -> str:
+
+        standard_claims = self.generate_standard_claims(datetime.now(tz=timezone.utc), exp)
+
+        claims: dict[str, Any] = {}
+        for c in [standard_claims, custom_claims]:
+            claims.update(c)
+
+        return jwt.encode(claims, self.secret, algorithm="HS256")  # pyright: ignore[reportUnknownMemberType]
+
+
+# Uses a generic encoder to create JWTs with specific custom scopes / expiration times
+class TokenGenerator:
+    """
+    Generates 'typed' JWTs with specific custom scopes / expiration times to serve a specific purpose.
+    """
+
+    def __init__(self, secret: str | bytes):
+        self.encoder = JWTEncoder(DEFAULT_ISSUER, DEFAULT_AUDIENCE, secret)
+
+    def bootstrap(self):
+        custom_claims = {"scope": BOOTSTRAP_SCOPE}
+        return self.encoder.new_token(custom_claims, BOOTSTRAP_EXP)

rsconnect/log.py

@@ -0,0 +1,253 @@
+"""
+Logging wrapper and shared instance
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import sys
+from functools import partial, wraps
+from typing import TYPE_CHECKING, Any, Callable, Literal, Optional, Protocol, TypeVar
+
+if sys.version_info >= (3, 10):
+    from typing import Concatenate, ParamSpec
+else:
+    from typing_extensions import Concatenate, ParamSpec
+
+
+if TYPE_CHECKING:
+    from collections.abc import MutableMapping
+
+import click
+
+T = TypeVar("T")
+P = ParamSpec("P")
+
+
+_DATE_FORMAT = "%Y-%m-%dT%H:%M:%S%z"
+
+VERBOSE = int((logging.INFO + logging.DEBUG) / 2)
+logging.addLevelName(VERBOSE, "VERBOSE")
+
+
+class LogOutputFormat(object):
+    TEXT = "text"
+    JSON = "json"
+    DEFAULT = TEXT
+    _all = [TEXT, JSON]
+    All = Literal["text", "json"]
+
+
+class JsonLogFormatter(logging.Formatter):
+    """
+    https://stackoverflow.com/a/70223539
+    Formatter that outputs JSON strings after parsing the LogRecord.
+
+    @param dict fmt_dict: Key: logging format attribute pairs.
+    @param str datefmt: Key: strftime format string
+    """
+
+    def __init__(self, fmt_dict: Optional[dict[str, str]] = None, datefmt: str = _DATE_FORMAT):
+        self.fmt_dict = (
+            fmt_dict if fmt_dict is not None else {"timestamp": "asctime", "level": "levelname", "message": "message"}
+        )
+        self.datefmt = datefmt
+
+    def usesTime(self):
+        """
+        Overwritten to look for the attribute in the format dict values instead of the fmt string.
+        """
+        return "asctime" in self.fmt_dict.values()
+
+    def formatMessage(self, record: logging.LogRecord):  # pyright: ignore[reportIncompatibleMethodOverride]
+        """
+        Overwritten to return a dictionary of the relevant LogRecord attributes instead of a string.
+        KeyError is raised if an unknown attribute is provided in the fmt_dict.
+        """
+        return {fmt_key: record.__dict__[fmt_val] for fmt_key, fmt_val in self.fmt_dict.items()}
+
+    def format(self, record: logging.LogRecord):
+        """
+        Mostly the same as the parent's class method, the difference being that a dict is manipulated and dumped as JSON
+        instead of a string.
+        """
+        record.message = record.getMessage()
+
+        if self.usesTime():
+            record.asctime = self.formatTime(record, self.datefmt)
+
+        message_dict = self.formatMessage(record)
+
+        if record.exc_info:
+            # Cache the traceback text to avoid converting it multiple times
+            # (it's constant anyway)
+            if not record.exc_text:
+                record.exc_text = self.formatException(record.exc_info)
+
+        if record.exc_text:
+            message_dict["exc_info"] = record.exc_text
+
+        if record.stack_info:
+            message_dict["stack_info"] = self.formatStack(record.stack_info)
+
+        return json.dumps(message_dict, default=str)
+
+
+# This is a workaround for LoggerAdapter not being generic in Python<=3.10.
+# See also:
+# https://github.com/python/typeshed/issues/7855#issuecomment-1128857842
+if sys.version_info >= (3, 11):
+    _LoggerAdapter = logging.LoggerAdapter[logging.Logger]
+else:
+    _LoggerAdapter = logging.LoggerAdapter
+
+
+class RSLogger(_LoggerAdapter):
+    def __init__(self):
+        super(RSLogger, self).__init__(logging.getLogger("rsconnect"), {})
+        self._in_feedback = False
+        self._have_feedback_output = False
+        self._log_format = LogOutputFormat.DEFAULT
+
+    def addHandler(self, handler: logging.Handler):
+        self.logger.addHandler(handler)
+
+    def set_in_feedback(self, value: bool):
+        self._in_feedback = value
+        self._have_feedback_output = False
+
+    def set_log_output_format(self, value: LogOutputFormat.All):
+        self._log_format = value
+        if self._log_format == LogOutputFormat.JSON:
+            for h in self.logger.handlers:
+                h.setFormatter(JsonLogFormatter())
+        else:
+            for h in self.logger.handlers:
+                h.setFormatter(logging.Formatter("[%(levelname)s] %(asctime)s %(message)s", datefmt=_DATE_FORMAT))
+
+    def process(self, msg: str, kwargs: MutableMapping[str, Any]):
+        msg, kwargs = super(RSLogger, self).process(msg, kwargs)
+        if self._in_feedback and self.is_debugging():
+            if not self._have_feedback_output:
+                print()
+                self._have_feedback_output = True
+            msg = click.style(" %s" % msg, fg="green")
+        return msg, kwargs
+
+    def is_debugging(self):
+        return self.isEnabledFor(logging.DEBUG)
+
+    def setLevel(self, level: int | str):
+        """
+        Set the specified level on the underlying logger.
+
+        **Note:** This is present in newer Python versions but since it's missing
+        from 2.7, we replicate it here.
+        """
+        self.logger.setLevel(level)
+
+
+logger = RSLogger()
+logger.addHandler(logging.StreamHandler())
+logger.set_log_output_format(LogOutputFormat.DEFAULT)
+
+
+class ConsoleFormatter(logging.Formatter):
+    green = "\x1b[32;20m"
+    yellow = "\x1b[33;20m"
+    red = "\x1b[31;20m"
+    msg_format = "%(message)s"
+    reset = "\x1b[0m"
+
+    FORMATS = {
+        logging.DEBUG: green + msg_format + reset,
+        logging.INFO: reset + msg_format + reset,
+        logging.WARNING: yellow + msg_format + reset,
+        logging.ERROR: red + msg_format + reset,
+        logging.CRITICAL: red + msg_format + reset,
+    }
+
+    def format(self, record: logging.LogRecord):
+        log_fmt = self.FORMATS.get(record.levelno)
+        formatter = logging.Formatter(log_fmt)
+        return formatter.format(record)
+
+
+console_logger = logging.getLogger("console")
+console_logger.setLevel(logging.DEBUG)
+
+# create console handler
+console_handler = logging.StreamHandler()
+console_handler.terminator = ""
+console_handler.setLevel(logging.DEBUG)
+console_handler.setFormatter(ConsoleFormatter())
+console_logger.addHandler(console_handler)
+
+
+def logged(logger: logging.Logger, label: str):
+    def decorator(f: Callable[P, T]) -> Callable[P, T]:
+        @wraps(f)
+        def wrapper(*args: P.args, **kw: P.kwargs):
+            logger.info(label)
+            result = None
+            try:
+                result = f(*args, **kw)
+            except Exception as exc:
+                logger.error(" \t[ERROR]: {}\n".format(str(exc)))
+                raise
+            logger.debug(" \t[OK]\n")
+            return result
+
+        return wrapper
+
+    return decorator
+
+
+class HasLoggerMethod(Protocol):
+    logger: logging.Logger | None
+
+
+HasLoggerMethodT = TypeVar("HasLoggerMethodT", bound=HasLoggerMethod)
+
+
+def cls_logged(label: str):  # uses logger provided by a class' self.logger
+    def decorator(
+        method: Callable[Concatenate[HasLoggerMethodT, P], T],
+    ) -> Callable[Concatenate[HasLoggerMethodT, P], T]:
+
+        @wraps(method)
+        def wrapper(self: HasLoggerMethodT, *args: P.args, **kw: P.kwargs):
+            logger = self.logger
+            if logger:
+                logger.info(label)
+            result = None
+            try:
+                result = method(self, *args, **kw)
+            except Exception as exc:
+                msg = " \t[ERROR]: {}\n"
+                if logger:
+                    logger.error(msg.format(str(exc)))
+                else:
+                    print(msg)
+                raise
+            if logger:
+                logger.debug(" \t[OK]\n")
+            return result
+
+        return wrapper
+
+    return decorator
+
+
+console_logged = partial(logged, console_logger)
+
+
+# generic logger
+connect_logger = logging.getLogger("connect_logger")
+connect_logger.setLevel(logging.DEBUG)
+connect_handler = logging.StreamHandler()
+connect_handler.terminator = "\n"
+connect_handler.setLevel(logging.DEBUG)
+connect_handler.setFormatter(ConsoleFormatter())
+connect_logger.addHandler(connect_handler)