# Copyright (c) Streamlit Inc. (2018-2022) Snowflake Inc. (2022-2026) # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Utility functions for the Starlette server implementation.""" from __future__ import annotations import binascii import os import time def parse_range_header(range_header: str, total_size: int) -> tuple[int, int]: """Parse the Range header and return the start and end byte positions. This is used for serving media files with range requests. Parameters ---------- range_header : str The value of the Range header (e.g. "bytes=0-1023"). total_size : int The total size of the resource in bytes. Returns ------- tuple[int, int] A tuple containing (start, end) byte positions. """ if total_size <= 0: raise ValueError("empty content") units, sep, range_spec = range_header.partition("=") if units.strip().lower() != "bytes" or sep == "" or "," in range_spec: raise ValueError("invalid range") range_spec = range_spec.strip() if range_spec.startswith("-"): try: suffix = int(range_spec[1:]) except ValueError: raise ValueError("invalid suffix range") from None if suffix <= 0: raise ValueError("invalid suffix range") if suffix >= total_size: return 0, total_size - 1 return total_size - suffix, total_size - 1 start_str, sep, end_str = range_spec.partition("-") if not start_str: raise ValueError("missing range start") start = int(start_str) if start < 0 or start >= total_size: raise ValueError("start out of range") if sep == "" or not end_str: end = total_size - 1 else: end = int(end_str) if end < start: raise ValueError("end before start") end = min(end, total_size - 1) return start, end def websocket_mask(mask: bytes, data: bytes) -> bytes: """Mask or unmask data for WebSocket transmission per RFC 6455. Each byte of data is XORed with mask[i % 4]. This operation is bidirectional - applying it twice with the same mask returns the original data. Parameters ---------- mask : bytes A 4-byte masking key. data : bytes The data to mask or unmask. Returns ------- bytes The masked/unmasked data. """ if len(mask) != 4: raise ValueError("mask must be 4 bytes") result = bytearray(len(data)) for i, byte in enumerate(data): result[i] = byte ^ mask[i % 4] return bytes(result) def create_signed_value( secret: str, name: str, value: str | bytes, ) -> bytes: """Create a signed cookie value using itsdangerous. Parameters ---------- secret The secret key used for signing. name The cookie name (used as salt for additional security). value The value to sign. Returns ------- bytes The signed value as bytes. """ from itsdangerous import URLSafeTimedSerializer serializer = URLSafeTimedSerializer(secret, salt=name) if isinstance(value, bytes): value = value.decode("utf-8") return serializer.dumps(value).encode("utf-8") def decode_signed_value( secret: str, name: str, value: str | bytes, max_age_days: float = 31, ) -> bytes | None: """Decode a signed cookie value using itsdangerous. Parameters ---------- secret The secret key used for signing. name The cookie name (used as salt for additional security). value The signed value to decode. max_age_days Maximum age of the cookie in days (default: 31). Returns ------- bytes | None The decoded value as bytes, or None if invalid/expired. """ from itsdangerous import BadSignature, SignatureExpired, URLSafeTimedSerializer if not value: return None try: if isinstance(value, bytes): value = value.decode("utf-8") serializer = URLSafeTimedSerializer(secret, salt=name) decoded = serializer.loads(value, max_age=int(max_age_days * 86400)) if isinstance(decoded, str): return decoded.encode("utf-8") if isinstance(decoded, bytes): return decoded # Unexpected type from deserializer — treat as invalid return None except (BadSignature, SignatureExpired, UnicodeDecodeError): return None def generate_xsrf_token_string( token_bytes: bytes | None = None, timestamp: int | None = None ) -> str: """Generate a version 2 XSRF token string. Format: 2|mask|masked_token|timestamp Parameters ---------- token_bytes The raw token bytes to encode. If None, generates 16 random bytes. timestamp The Unix timestamp to include in the token. If None, uses current time. Returns ------- str The encoded XSRF token string in version 2 format. """ if token_bytes is None: token_bytes = os.urandom(16) if timestamp is None: timestamp = int(time.time()) mask = os.urandom(4) masked_token = websocket_mask(mask, token_bytes) return "2|{}|{}|{}".format( binascii.b2a_hex(mask).decode("ascii"), binascii.b2a_hex(masked_token).decode("ascii"), timestamp, ) def decode_xsrf_token_string( cookie_value: str, ) -> tuple[bytes | None, int | None]: """Decode an XSRF token string. Supports version 2 (masked) and version 1 (unmasked) tokens. Parameters ---------- cookie_value The XSRF token cookie value to decode. Returns ------- tuple[bytes | None, int | None] A tuple of (token_bytes, timestamp). Both values are None if decoding fails. """ if not cookie_value: return None, None value = cookie_value.strip("\"'") if not value: return None, None try: # V2 tokens: if value.startswith("2|"): _, mask_hex, masked_hex, timestamp_str = value.split("|") mask = binascii.a2b_hex(mask_hex.encode("ascii")) masked = binascii.a2b_hex(masked_hex.encode("ascii")) token = websocket_mask(mask, masked) return token, int(timestamp_str) # V1 tokens: # TODO(lukasmasuch): This is likely unused in Streamlit since only V2 tokens # are used. We might be able to just remove this part. token = binascii.a2b_hex(value.encode("ascii")) if not token: return None, None # V1 tokens don't have an embedded timestamp, so we use current time # as a placeholder. This timestamp is informational only and not used # for token validation. return token, int(time.time()) except (binascii.Error, ValueError): return None, None def generate_random_hex_string(num_bytes: int = 32) -> str: """Generate a cryptographically secure random hex string. Parameters ---------- num_bytes Number of random bytes to generate (default: 32). The resulting hex string will be twice this length. Returns ------- str A hex-encoded random string. """ return binascii.b2a_hex(os.urandom(num_bytes)).decode("ascii") def validate_xsrf_token(supplied_token: str | None, xsrf_cookie: str | None) -> bool: """Validate the XSRF token from the WebSocket subprotocol against the cookie. This ensures the frontend can share XSRF logic between WebSocket handshake and HTTP uploads. """ if not supplied_token or not xsrf_cookie: return False # Decode the supplied token from the subprotocol supplied_token_bytes, _ = decode_xsrf_token_string(supplied_token) # Decode the expected token from the cookie expected_token_bytes, _ = decode_xsrf_token_string(xsrf_cookie) if not supplied_token_bytes or not expected_token_bytes: return False import hmac return hmac.compare_digest(supplied_token_bytes, expected_token_bytes)