utils/session_observer.py

"""Pure Python executor for Sentinel Session Observer programs.

The Sentinel SDK receives two encrypted VM programs.  The collector installs
event handlers and initializes state; the snapshot program must then run in
the same process map.  ``build_session_observer_token`` mirrors that sequence
and returns the JSON value used by ``OpenAI-Sentinel-SO-Token``.

Malformed programs raise ``ValueError``.  Runtime JavaScript-like errors are
handled by the VM in the same places as the browser SDK and normally produce
an encoded diagnostic value instead of escaping to the caller.
"""

from __future__ import annotations

import base64
import json
import math
import random
import time
from collections import defaultdict
from collections.abc import Callable
from dataclasses import dataclass
from typing import Any


class _Undefined:
    def __repr__(self) -> str:
        return "undefined"


UNDEFINED = _Undefined()


def _js_string(value: Any) -> str:
    if value is UNDEFINED:
        return "undefined"
    if value is None:
        return "null"
    if value is True:
        return "true"
    if value is False:
        return "false"
    if isinstance(value, str):
        return value
    if isinstance(value, float):
        if math.isnan(value):
            return "NaN"
        if math.isinf(value):
            return "Infinity" if value > 0 else "-Infinity"
        if value.is_integer():
            return str(int(value))
    if isinstance(value, (list, tuple)):
        return ",".join("" if item is None or item is UNDEFINED else _js_string(item) for item in value)
    if isinstance(value, _JSObject):
        return "[object Object]"
    if callable(value):
        return "function () { [native code] }"
    return str(value)


def _js_number(value: Any) -> float:
    if value is None:
        return 0.0
    if value is True:
        return 1.0
    if value is False:
        return 0.0
    if isinstance(value, (int, float)):
        return float(value)
    try:
        text = _js_string(value).strip()
        return 0.0 if not text else float(text)
    except (TypeError, ValueError):
        return math.nan


def _latin1_btoa(value: Any) -> str:
    return base64.b64encode(_js_string(value).encode("latin1", errors="replace")).decode("ascii")


def _latin1_atob(value: Any) -> str:
    return base64.b64decode(_js_string(value), validate=True).decode("latin1")


def _xor_text(value: str, key: str) -> str:
    if not key:
        raise ValueError("requirements_token must not be empty")
    return "".join(chr(ord(char) ^ ord(key[index % len(key)])) for index, char in enumerate(value))


class _JSObject(dict[str, Any]):
    pass


class _JSMap:
    def __init__(self) -> None:
        self.values: dict[Any, Any] = {}


class _EventTarget(_JSObject):
    def __init__(self, values: dict[str, Any] | None = None) -> None:
        super().__init__(values or {})
        self.listeners: dict[str, list[Callable[..., Any]]] = defaultdict(list)

    def add_event_listener(self, event_type: Any, callback: Any, *_: Any) -> None:
        if callable(callback):
            self.listeners[_js_string(event_type)].append(callback)

    def remove_event_listener(self, event_type: Any, callback: Any, *_: Any) -> None:
        listeners = self.listeners.get(_js_string(event_type), [])
        self.listeners[_js_string(event_type)] = [item for item in listeners if item is not callback]

    def dispatch(self, event_type: str, **values: Any) -> None:
        event = _JSObject({"type": event_type, **values})
        for callback in tuple(self.listeners.get(event_type, ())):
            try:
                callback(event)
            except Exception:
                continue


@dataclass
class _Clock:
    start_epoch_ms: float
    elapsed_ms: float = 12.0

    def now(self) -> float:
        self.elapsed_ms += 0.37
        return self.elapsed_ms

    def date_now(self) -> int:
        return int(self.start_epoch_ms + self.now())


def _json_value(value: Any) -> Any:
    if value is UNDEFINED:
        return None
    if isinstance(value, _JSObject):
        return {key: _json_value(item) for key, item in value.items() if item is not UNDEFINED}
    if isinstance(value, _JSMap):
        return {}
    if isinstance(value, (list, tuple)):
        return [_json_value(item) for item in value]
    if isinstance(value, float) and (math.isnan(value) or math.isinf(value)):
        return None
    return value


def _json_stringify(value: Any) -> Any:
    if value is UNDEFINED or callable(value):
        return UNDEFINED
    return json.dumps(_json_value(value), ensure_ascii=False, separators=(",", ":"))


class _SessionObserverVM:
    def __init__(self, *, device_id: str, flow: str, seed: int | None = None) -> None:
        self.device_id = device_id
        self.flow = flow
        self.random = random.Random(seed)
        self.clock = _Clock(time.time() * 1000)
        self.window = self._make_window()
        self.slots: dict[Any, Any] = {}
        self.operation_count = 0
        self.result: str | None = None
        self._settled = False

    def _make_window(self) -> _EventTarget:
        window = _EventTarget()
        performance = _JSObject({"now": self.clock.now, "timeOrigin": self.clock.start_epoch_ms})
        page_path = {
            "authorize_continue": "/email-verification",
            "oauth_create_account": "/about-you",
            "password_verify": "/log-in/password",
            "username_password_create": "/create-account/password",
        }.get(self.flow, "/create-account/password")
        location = _JSObject(
            {
                "href": f"https://auth.openai.com{page_path}",
                "origin": "https://auth.openai.com",
                "pathname": page_path,
                "search": "",
            }
        )
        navigator = _JSObject(
            {
                "userAgent": (
                    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 "
                    "(KHTML, like Gecko) Chrome/145.0.0.0 Safari/537.36"
                ),
                "language": "en-US",
                "languages": ["en-US", "en"],
                "platform": "Win32",
                "hardwareConcurrency": 8,
                "deviceMemory": 8,
                "maxTouchPoints": 0,
                "webdriver": False,
                "cookieEnabled": True,
                "vendor": "Google Inc.",
            }
        )
        document = _EventTarget(
            {
                "location": location,
                "visibilityState": "visible",
                "hidden": False,
                "readyState": "complete",
                "referrer": "https://auth.openai.com/",
                "cookie": f"oai-did={self.device_id}",
                "scripts": [],
                "documentElement": _JSObject(),
                "body": _JSObject(),
            }
        )
        local_storage = _JSObject({"oai-did": self.device_id})
        math_object = _JSObject(
            {
                "abs": abs,
                "hypot": math.hypot,
                "sqrt": math.sqrt,
                "round": round,
                "floor": math.floor,
                "random": self.random.random,
            }
        )
        object_object = _JSObject(
            {
                "create": lambda *_: _JSObject(),
                "keys": lambda value: list(value.keys()) if isinstance(value, (dict, _JSObject)) else [],
            }
        )
        reflect = _JSObject({"set": self._reflect_set})
        date_object = _JSObject({"now": self.clock.date_now})
        json_object = _JSObject({"parse": lambda value: json.loads(_js_string(value)), "stringify": _json_stringify})
        window.update(
            {
                "window": window,
                "self": window,
                "top": window,
                "parent": window,
                "globalThis": window,
                "document": document,
                "navigator": navigator,
                "location": location,
                "performance": performance,
                "screen": _JSObject({"width": 1920, "height": 1080, "availWidth": 1920, "availHeight": 1040}),
                "localStorage": local_storage,
                "sessionStorage": _JSObject(),
                "Math": math_object,
                "Object": object_object,
                "Reflect": reflect,
                "Date": date_object,
                "JSON": json_object,
                "btoa": _latin1_btoa,
                "atob": _latin1_atob,
                "setTimeout": lambda *_: 1,
                "clearTimeout": lambda *_: None,
                "scrollX": 0,
                "scrollY": 0,
                "innerWidth": 1920,
                "innerHeight": 937,
            }
        )
        window["addEventListener"] = window.add_event_listener
        window["removeEventListener"] = window.remove_event_listener
        document["addEventListener"] = document.add_event_listener
        document["removeEventListener"] = document.remove_event_listener
        return window

    @staticmethod
    def _reflect_set(target: Any, key: Any, value: Any) -> bool:
        _js_set(target, key, value)
        return True

    def _reset(self, requirements_token: str) -> None:
        self.slots.clear()
        self.operation_count = 0
        self.result = None
        self._settled = False
        self.slots.update(
            {
                0: self._run_nested_dx,
                1: self._op_xor,
                2: self._op_literal,
                3: self._op_resolve,
                4: self._op_reject,
                5: self._op_add,
                6: self._op_get_property,
                7: self._op_call,
                8: self._op_copy,
                9: [],
                10: self.window,
                11: self._op_find_script,
                12: self._op_process_map,
                13: self._op_call_raw,
                14: self._op_json_parse,
                15: self._op_json_stringify,
                16: requirements_token,
                17: self._op_call_assign,
                18: self._op_atob,
                19: self._op_btoa,
                20: self._op_if_equal,
                21: self._op_if_distance,
                22: self._op_nested_queue,
                23: self._op_if_defined,
                24: self._op_bind,
                25: self._op_noop,
                26: self._op_noop,
                27: self._op_subtract,
                28: self._op_noop,
                29: self._op_less_than,
                30: self._op_make_function,
                33: self._op_multiply,
                34: self._op_resolve_value,
                35: self._op_divide,
            }
        )

    def run(self, dx: str, *, requirements_token: str | None = None) -> str:
        if requirements_token is not None:
            self._reset(requirements_token)
        if 16 not in self.slots:
            raise ValueError("collector must run before snapshot")
        try:
            decoded = base64.b64decode(dx, validate=True).decode("latin1")
            instructions = json.loads(_xor_text(decoded, _js_string(self.slots[16])))
        except (ValueError, TypeError, json.JSONDecodeError) as exc:
            raise ValueError("invalid Session Observer dx program") from exc
        if not isinstance(instructions, list):
            raise ValueError("Session Observer program must be an instruction array")
        self.result = None
        self._settled = False
        self.slots[9] = list(instructions)
        try:
            self._execute_queue()
        except Exception as exc:
            if not self._settled:
                self.result = _latin1_btoa(f"{self.operation_count}: {exc}")
                self._settled = True
        return self.result or ""

    def emit_synthetic_activity(self) -> None:
        x = 420 + self.random.randint(-20, 20)
        y = 280 + self.random.randint(-20, 20)
        self.window.dispatch("pointermove", clientX=x, clientY=y)
        self.window.dispatch("pointermove", clientX=x + 37, clientY=y + 19)
        self.window.dispatch("click", clientX=x + 37, clientY=y + 19)
        self.window.dispatch("keydown", key="a", altKey=False, ctrlKey=False, metaKey=False)
        self.window["scrollY"] = 184
        self.window.dispatch("scroll")
        self.window.dispatch("wheel")

    def _execute_queue(self) -> None:
        while isinstance(self.slots.get(9), list) and self.slots[9]:
            queue = self.slots[9]
            row = queue.pop(0)
            if not isinstance(row, list) or not row:
                self.operation_count += 1
                continue
            operation = self.slots.get(row[0], UNDEFINED)
            if callable(operation):
                operation(*row[1:])
            self.operation_count += 1

    def _run_nested_dx(self, dx: Any) -> str:
        return self.run(_js_string(dx))

    def _op_xor(self, target: Any, other: Any) -> None:
        self.slots[target] = _xor_text(_js_string(self.slots.get(target, UNDEFINED)), _js_string(self.slots.get(other, UNDEFINED)))

    def _op_literal(self, target: Any, value: Any) -> None:
        self.slots[target] = value

    def _op_resolve(self, value: Any) -> None:
        if not self._settled:
            self.result = _latin1_btoa(value)
            self._settled = True

    def _op_reject(self, value: Any) -> None:
        if not self._settled:
            self.result = _latin1_btoa(value)
            self._settled = True

    def _op_add(self, target: Any, other: Any) -> None:
        current = self.slots.get(target, UNDEFINED)
        incoming = self.slots.get(other, UNDEFINED)
        if isinstance(current, list):
            current.append(incoming)
        elif isinstance(current, str) or isinstance(incoming, str):
            self.slots[target] = _js_string(current) + _js_string(incoming)
        else:
            self.slots[target] = _js_number(current) + _js_number(incoming)

    def _op_get_property(self, target: Any, owner: Any, key: Any) -> None:
        self.slots[target] = _js_get(self.slots.get(owner, UNDEFINED), self.slots.get(key, UNDEFINED))

    def _op_call(self, function_slot: Any, *argument_slots: Any) -> Any:
        function = self.slots.get(function_slot, UNDEFINED)
        if not callable(function):
            raise TypeError(f"{_js_string(function)} is not a function")
        return function(*(self.slots.get(slot, UNDEFINED) for slot in argument_slots))

    def _op_copy(self, target: Any, source: Any) -> None:
        self.slots[target] = self.slots.get(source, UNDEFINED)

    def _op_find_script(self, target: Any, pattern_slot: Any) -> None:
        self.slots[target] = None

    def _op_process_map(self, target: Any) -> None:
        self.slots[target] = self.slots

    def _op_call_raw(self, error_slot: Any, function_slot: Any, *raw_arguments: Any) -> None:
        try:
            function = self.slots.get(function_slot, UNDEFINED)
            if not callable(function):
                raise TypeError(f"{_js_string(function)} is not a function")
            function(*raw_arguments)
        except Exception as exc:
            self.slots[error_slot] = f"{type(exc).__name__}: {exc}"

    def _op_json_parse(self, target: Any, source: Any) -> None:
        self.slots[target] = json.loads(_js_string(self.slots.get(source, UNDEFINED)))

    def _op_json_stringify(self, target: Any, source: Any) -> None:
        self.slots[target] = _json_stringify(self.slots.get(source, UNDEFINED))

    def _op_call_assign(self, target: Any, function_slot: Any, *argument_slots: Any) -> None:
        try:
            function = self.slots.get(function_slot, UNDEFINED)
            if not callable(function):
                raise TypeError(f"{_js_string(function)} is not a function")
            self.slots[target] = function(*(self.slots.get(slot, UNDEFINED) for slot in argument_slots))
        except Exception as exc:
            self.slots[target] = f"{type(exc).__name__}: {exc}"

    def _op_atob(self, target: Any) -> None:
        self.slots[target] = _latin1_atob(self.slots.get(target, UNDEFINED))

    def _op_btoa(self, target: Any) -> None:
        self.slots[target] = _latin1_btoa(self.slots.get(target, UNDEFINED))

    def _op_if_equal(self, left: Any, right: Any, function_slot: Any, *raw_arguments: Any) -> None:
        if _js_strict_equal(self.slots.get(left, UNDEFINED), self.slots.get(right, UNDEFINED)):
            function = self.slots.get(function_slot, UNDEFINED)
            if callable(function):
                function(*raw_arguments)

    def _op_if_distance(self, left: Any, right: Any, distance: Any, function_slot: Any, *raw_arguments: Any) -> None:
        if abs(_js_number(self.slots.get(left, UNDEFINED)) - _js_number(self.slots.get(right, UNDEFINED))) > _js_number(
            self.slots.get(distance, UNDEFINED)
        ):
            function = self.slots.get(function_slot, UNDEFINED)
            if callable(function):
                function(*raw_arguments)

    def _op_nested_queue(self, error_slot: Any, instructions: Any) -> None:
        old_queue = self.slots.get(9, [])
        try:
            self.slots[9] = list(instructions)
            self._execute_queue()
        except Exception as exc:
            self.slots[error_slot] = str(exc)
        finally:
            self.slots[9] = old_queue

    def _op_if_defined(self, value_slot: Any, function_slot: Any, *raw_arguments: Any) -> None:
        if self.slots.get(value_slot, UNDEFINED) is not UNDEFINED:
            function = self.slots.get(function_slot, UNDEFINED)
            if callable(function):
                function(*raw_arguments)

    def _op_bind(self, target: Any, owner: Any, key: Any) -> None:
        self.slots[target] = _js_get(self.slots.get(owner, UNDEFINED), self.slots.get(key, UNDEFINED))

    @staticmethod
    def _op_noop(*_: Any) -> None:
        return None

    def _op_subtract(self, target: Any, other: Any) -> None:
        current = self.slots.get(target, UNDEFINED)
        incoming = self.slots.get(other, UNDEFINED)
        if isinstance(current, list):
            try:
                current.remove(incoming)
            except ValueError:
                pass
        else:
            self.slots[target] = _js_number(current) - _js_number(incoming)

    def _op_less_than(self, target: Any, left: Any, right: Any) -> None:
        self.slots[target] = _js_number(self.slots.get(left, UNDEFINED)) < _js_number(self.slots.get(right, UNDEFINED))

    def _op_make_function(self, target: Any, result_slot: Any, parameter_slots: Any, body: Any = UNDEFINED) -> None:
        has_parameters = isinstance(body, list)
        parameters = list(parameter_slots) if has_parameters and isinstance(parameter_slots, list) else []
        instructions = list(body if has_parameters else (parameter_slots or []))

        def generated(*arguments: Any) -> Any:
            old_queue = self.slots.get(9, [])
            for index, slot in enumerate(parameters):
                self.slots[slot] = arguments[index] if index < len(arguments) else UNDEFINED
            try:
                self.slots[9] = list(instructions)
                self._execute_queue()
                return self.slots.get(result_slot, UNDEFINED)
            except Exception as exc:
                return str(exc)
            finally:
                self.slots[9] = old_queue

        self.slots[target] = generated

    def _op_multiply(self, target: Any, left: Any, right: Any) -> None:
        self.slots[target] = _js_number(self.slots.get(left, UNDEFINED)) * _js_number(self.slots.get(right, UNDEFINED))

    def _op_resolve_value(self, target: Any, source: Any) -> None:
        self.slots[target] = self.slots.get(source, UNDEFINED)

    def _op_divide(self, target: Any, left: Any, right: Any) -> None:
        numerator = _js_number(self.slots.get(left, UNDEFINED))
        denominator = _js_number(self.slots.get(right, UNDEFINED))
        self.slots[target] = 0 if denominator == 0 else numerator / denominator


def _js_strict_equal(left: Any, right: Any) -> bool:
    if left is UNDEFINED or right is UNDEFINED:
        return left is right
    if type(left) is not type(right) and not (
        isinstance(left, (int, float)) and not isinstance(left, bool) and isinstance(right, (int, float)) and not isinstance(right, bool)
    ):
        return False
    return left == right


def _js_set(target: Any, key: Any, value: Any) -> None:
    key_text = _js_string(key)
    if isinstance(target, _JSMap):
        target.values[key] = value
    elif isinstance(target, dict):
        target[key_text] = value
    elif isinstance(target, list) and key_text.isdigit():
        index = int(key_text)
        while len(target) <= index:
            target.append(UNDEFINED)
        target[index] = value
    else:
        setattr(target, key_text, value)


def _js_get(owner: Any, key: Any) -> Any:
    key_text = _js_string(key)
    if owner is UNDEFINED or owner is None:
        raise TypeError(f"Cannot read properties of {_js_string(owner)}")
    if isinstance(owner, _JSMap):
        methods = {
            "get": lambda item: owner.values.get(item, UNDEFINED),
            "set": lambda item, value: owner.values.__setitem__(item, value) or owner,
            "has": lambda item: item in owner.values,
            "delete": lambda item: owner.values.pop(item, UNDEFINED) is not UNDEFINED,
            "clear": owner.values.clear,
            "keys": lambda: list(owner.values.keys()),
            "values": lambda: list(owner.values.values()),
            "size": len(owner.values),
        }
        return methods.get(key_text, UNDEFINED)
    if isinstance(owner, dict):
        return owner.get(key_text, UNDEFINED)
    if isinstance(owner, (list, tuple)):
        if key_text == "length":
            return len(owner)
        if key_text.lstrip("-").isdigit():
            index = int(key_text)
            return owner[index] if -len(owner) <= index < len(owner) else UNDEFINED
        methods = {
            "push": lambda *items: owner.extend(items) or len(owner),
            "indexOf": lambda item: owner.index(item) if item in owner else -1,
            "join": lambda separator=",": _js_string(separator).join(_js_string(item) for item in owner),
            "slice": lambda start=0, end=None: list(owner)[int(_js_number(start)) : None if end is None or end is UNDEFINED else int(_js_number(end))],
            "shift": lambda: owner.pop(0) if owner else UNDEFINED,
            "map": lambda function: [function(item, index, owner) for index, item in enumerate(owner)],
            "filter": lambda function: [item for index, item in enumerate(owner) if function(item, index, owner)],
        }
        return methods.get(key_text, UNDEFINED)
    if isinstance(owner, str):
        if key_text == "length":
            return len(owner)
        if key_text.lstrip("-").isdigit():
            index = int(key_text)
            return owner[index] if -len(owner) <= index < len(owner) else UNDEFINED
        methods = {
            "charCodeAt": lambda index=0: ord(owner[int(_js_number(index))]) if 0 <= int(_js_number(index)) < len(owner) else math.nan,
            "indexOf": lambda value: owner.find(_js_string(value)),
            "includes": lambda value: _js_string(value) in owner,
            "slice": lambda start=0, end=None: owner[int(_js_number(start)) : None if end is None or end is UNDEFINED else int(_js_number(end))],
            "substring": lambda start=0, end=None: owner[max(0, int(_js_number(start))) : None if end is None or end is UNDEFINED else max(0, int(_js_number(end)))],
            "toLowerCase": owner.lower,
            "toUpperCase": owner.upper,
            "split": lambda separator=None: list(owner) if separator == "" else owner.split(None if separator is None or separator is UNDEFINED else _js_string(separator)),
            "toString": lambda: owner,
        }
        return methods.get(key_text, UNDEFINED)
    if isinstance(owner, (int, float, bool)):
        if key_text in {"toString", "valueOf"}:
            return lambda: _js_string(owner) if key_text == "toString" else owner
        return UNDEFINED
    if callable(owner):
        if key_text == "bind":
            return lambda *_: owner
        if key_text == "call":
            return lambda _this=UNDEFINED, *args: owner(*args)
        if key_text == "apply":
            return lambda _this=UNDEFINED, args=(): owner(*args)
        if key_text == "toString":
            return lambda: _js_string(owner)
    return getattr(owner, key_text, UNDEFINED)


def build_session_observer_token(
    *,
    collector_dx: str,
    snapshot_dx: str,
    requirements_token: str,
    response_token: str,
    device_id: str,
    flow: str,
) -> str:
    """Build the JSON header value for a required Session Observer challenge.

    ``requirements_token`` is the ``p`` value sent to Sentinel ``/req``.  It is
    the XOR key for both programs; it is not the subsequent proof-of-work
    answer.  Invalid encrypted programs raise ``ValueError``.
    """

    if not all(isinstance(value, str) and value for value in (collector_dx, snapshot_dx, requirements_token)):
        raise ValueError("collector_dx, snapshot_dx and requirements_token are required")
    vm = _SessionObserverVM(device_id=device_id, flow=flow)
    vm.run(collector_dx, requirements_token=requirements_token)
    vm.emit_synthetic_activity()
    snapshot = vm.run(snapshot_dx)
    if not snapshot:
        return ""
    return json.dumps(
        {"so": snapshot, "c": response_token, "id": device_id, "flow": flow},
        ensure_ascii=False,
        separators=(",", ":"),
    )


def execute_sentinel_vm_token(
    *,
    dx: str,
    requirements_token: str,
    device_id: str = "",
    flow: str = "",
) -> str:
    """Execute one current Sentinel VM program and return its encoded result."""

    if not isinstance(dx, str) or not dx or not isinstance(requirements_token, str) or not requirements_token:
        raise ValueError("dx and requirements_token are required")
    vm = _SessionObserverVM(device_id=device_id, flow=flow)
    return vm.run(dx, requirements_token=requirements_token)


__all__ = ["build_session_observer_token", "execute_sentinel_vm_token"]