offwork 0.4.0__py3-none-any.whl

offwork/core/pairing.py

@@ -0,0 +1,389 @@
+"""PIN-based pairing protocol for client-worker key exchange.
+
+Pairing allows a client and worker to establish a shared secret by
+entering the same short PIN on both sides.  The protocol is inspired by
+SPAKE2 / SAS-based verification and works as follows:
+
+1. One side (the *initiator*) generates a random 6-digit PIN and shows it
+   to the user.
+2. The user enters the same PIN on the other side (the *responder*).
+3. Both sides independently derive an *intermediate secret* from the PIN
+   using a key-derivation step (HMAC-SHA256 with a fixed salt).
+4. The initiator generates a random *challenge nonce* and sends it to the
+   responder over the backend channel.
+5. The responder computes ``HMAC(intermediate_secret, nonce)`` and sends
+   the result back.
+6. The initiator verifies the response.  If it matches, both sides derive
+   the final shared secret from
+   ``HMAC(intermediate_secret, nonce ‖ "confirmed")``.
+
+The protocol prevents replay attacks (nonce is random per session) and
+ensures that a passive eavesdropper who observes the nonce+response
+cannot recover the PIN or the shared secret.
+
+All primitives are stdlib-only.
+"""
+
+import os
+import hmac
+import json
+import time
+import asyncio
+import hashlib
+import logging
+import secrets
+from typing import Any
+from pathlib import Path
+from dataclasses import field, dataclass
+
+from offwork.core.errors import PairingError
+
+logger = logging.getLogger(__name__)
+
+# Fixed salt used to bind the PIN derivation to offwork
+_PIN_SALT = b"offwork-pairing-v1"
+
+# Where key material is persisted
+_DEFAULT_KEY_DIR = Path.home() / ".offwork"
+_CLIENT_KEY_FILE = "client.key"
+_WORKER_KEY_FILE = "worker.key"
+
+# PIN length (digits)
+PIN_LENGTH = 6
+
+
+def generate_pin() -> str:
+    """Generate a random numeric PIN of :data:`PIN_LENGTH` digits."""
+    # secrets.randbelow is cryptographically secure
+    num = secrets.randbelow(10 ** PIN_LENGTH)
+    return str(num).zfill(PIN_LENGTH)
+
+
+def _derive_intermediate(pin: str) -> bytes:
+    """Derive a 32-byte intermediate key from a PIN string."""
+    return hmac.new(
+        _PIN_SALT, pin.encode("utf-8"), hashlib.sha256
+    ).digest()
+
+
+def generate_challenge() -> bytes:
+    """Generate a 32-byte random challenge nonce."""
+    return os.urandom(32)
+
+
+def compute_response(intermediate: bytes, challenge: bytes) -> bytes:
+    """Compute the challenge-response: ``HMAC(intermediate, challenge)``."""
+    return hmac.new(intermediate, challenge, hashlib.sha256).digest()
+
+
+def verify_response(
+    intermediate: bytes, challenge: bytes, response: bytes
+) -> bool:
+    """Verify a challenge-response in constant time."""
+    expected = compute_response(intermediate, challenge)
+    return hmac.compare_digest(expected, response)
+
+
+def derive_shared_secret(intermediate: bytes, challenge: bytes) -> bytes:
+    """Derive the final 32-byte shared secret after successful verification."""
+    material = challenge + b"confirmed"
+    return hmac.new(intermediate, material, hashlib.sha256).digest()
+
+
+# -- High-level pairing state -----------------------------------------------
+
+
+@dataclass
+class PairingResult:
+    """Outcome of a successful pairing exchange."""
+
+    shared_key: bytes
+    peer_role: str  # "client" or "worker"
+    paired_at: float = field(default_factory=time.time)
+
+
+# -- Pairing messages (JSON-serializable) -----------------------------------
+
+
+def _encode_bytes(b: bytes) -> str:
+    """Encode bytes as hex for JSON transport."""
+    return b.hex()
+
+
+def _decode_bytes(s: str) -> bytes:
+    """Decode hex string back to bytes."""
+    return bytes.fromhex(s)
+
+
+def make_challenge_message(challenge: bytes) -> str:
+    """Build the JSON pairing-challenge message."""
+    return json.dumps({
+        "type": "pairing_challenge",
+        "challenge": _encode_bytes(challenge),
+    })
+
+
+def parse_challenge_message(msg: str) -> bytes:
+    """Extract the challenge nonce from a pairing-challenge message.
+
+    Raises
+    ------
+    ValueError
+        If the message is not a valid pairing challenge.
+    """
+    data = json.loads(msg)
+    if data.get("type") != "pairing_challenge":
+        raise ValueError(f"Expected pairing_challenge, got {data.get('type')!r}")
+    return _decode_bytes(data["challenge"])
+
+
+def make_response_message(response: bytes) -> str:
+    """Build the JSON pairing-response message."""
+    return json.dumps({
+        "type": "pairing_response",
+        "response": _encode_bytes(response),
+    })
+
+
+def parse_response_message(msg: str) -> bytes:
+    """Extract the response from a pairing-response message.
+
+    Raises
+    ------
+    ValueError
+        If the message is not a valid pairing response.
+    """
+    data = json.loads(msg)
+    if data.get("type") != "pairing_response":
+        raise ValueError(f"Expected pairing_response, got {data.get('type')!r}")
+    return _decode_bytes(data["response"])
+
+
+def make_confirm_message() -> str:
+    """Build the JSON pairing-confirmed message."""
+    return json.dumps({"type": "pairing_confirmed"})
+
+
+def parse_confirm_message(msg: str) -> None:
+    """Validate a pairing-confirmed message.
+
+    Raises
+    ------
+    ValueError
+        If the message is not a valid pairing confirmation.
+    """
+    data = json.loads(msg)
+    if data.get("type") != "pairing_confirmed":
+        raise ValueError(f"Expected pairing_confirmed, got {data.get('type')!r}")
+
+
+# -- Key persistence --------------------------------------------------------
+
+
+def _ensure_key_dir(key_dir: Path | None = None) -> Path:
+    """Return the key directory, creating it if necessary."""
+    d = key_dir or _DEFAULT_KEY_DIR
+    d.mkdir(parents=True, exist_ok=True)
+    return d
+
+
+def save_shared_key(
+    shared_key: bytes,
+    role: str,
+    key_dir: Path | None = None,
+) -> Path:
+    """Persist *shared_key* to disk.
+
+    Parameters
+    ----------
+    shared_key
+        The 32-byte shared secret.
+    role
+        ``"client"`` or ``"worker"`` — determines the filename.
+    key_dir
+        Override the default ``~/.offwork`` directory.
+
+    Returns
+    -------
+    Path
+        The file that was written.
+    """
+    d = _ensure_key_dir(key_dir)
+    filename = _CLIENT_KEY_FILE if role == "client" else _WORKER_KEY_FILE
+    path = d / filename
+    path.write_bytes(shared_key)
+    # Restrict permissions: owner-only
+    path.chmod(0o600)
+    logger.info("Saved shared key to %s", path)
+    return path
+
+
+def load_shared_key(
+    role: str,
+    key_dir: Path | None = None,
+) -> bytes | None:
+    """Load a previously saved shared key, or return *None*.
+
+    Parameters
+    ----------
+    role
+        ``"client"`` or ``"worker"``.
+    key_dir
+        Override the default ``~/.offwork`` directory.
+    """
+    d = _ensure_key_dir(key_dir)
+    filename = _CLIENT_KEY_FILE if role == "client" else _WORKER_KEY_FILE
+    path = d / filename
+    if not path.exists():
+        return None
+    key = path.read_bytes()
+    if len(key) != 32:
+        logger.warning("Invalid key file %s (expected 32 bytes, got %d)", path, len(key))
+        return None
+    return key
+
+
+def clear_shared_key(
+    role: str,
+    key_dir: Path | None = None,
+) -> bool:
+    """Delete a saved shared key.  Returns ``True`` if a file was removed."""
+    d = _ensure_key_dir(key_dir)
+    filename = _CLIENT_KEY_FILE if role == "client" else _WORKER_KEY_FILE
+    path = d / filename
+    if path.exists():
+        path.unlink()
+        logger.info("Removed shared key %s", path)
+        return True
+    return False
+
+
+# -- Backend channel helpers ------------------------------------------------
+
+_PAIRING_CHANNEL = "offwork:pairing"
+
+
+async def initiate_pairing(
+    backend: Any,
+    pin: str,
+    timeout: float = 30.0,
+) -> PairingResult:
+    """Run the *initiator* side of the pairing protocol.
+
+    The initiator is typically the **worker**: it generates a challenge,
+    publishes it on the pairing channel, and waits for the client's
+    response.
+
+    Parameters
+    ----------
+    backend
+        A offwork :class:`~offwork.worker.backends.base.Backend` instance
+        with ``send_progress`` / ``get_progress`` used as a simple KV
+        channel, or any object that exposes the pairing channel methods.
+    pin
+        The PIN entered by the user.
+    timeout
+        Seconds to wait for the peer to respond.
+
+    Raises
+    ------
+    PairingError
+        On timeout or verification failure.
+    """
+    intermediate = _derive_intermediate(pin)
+    challenge = generate_challenge()
+
+    # Publish challenge
+    challenge_msg = make_challenge_message(challenge)
+    await backend.send_progress(_PAIRING_CHANNEL, challenge_msg)
+    logger.debug("Pairing: sent challenge")
+
+    # Wait for response
+    deadline = time.monotonic() + timeout
+    while time.monotonic() < deadline:
+        raw = await backend.get_progress(_PAIRING_CHANNEL + ":response")
+        if raw is not None:
+            try:
+                response = parse_response_message(raw)
+            except (ValueError, json.JSONDecodeError):
+                await asyncio.sleep(0.5)
+                continue
+
+            if not verify_response(intermediate, challenge, response):
+                raise PairingError("PIN mismatch — pairing failed")
+
+            # Derive shared secret and confirm
+            shared = derive_shared_secret(intermediate, challenge)
+            confirm_msg = make_confirm_message()
+            await backend.send_progress(_PAIRING_CHANNEL + ":confirm", confirm_msg)
+            logger.info("Pairing successful (initiator)")
+            return PairingResult(shared_key=shared, peer_role="client")
+
+        await asyncio.sleep(0.5)
+
+    raise PairingError(f"Pairing timed out after {timeout}s — no response from peer")
+
+
+async def respond_to_pairing(
+    backend: Any,
+    pin: str,
+    timeout: float = 30.0,
+) -> PairingResult:
+    """Run the *responder* side of the pairing protocol.
+
+    The responder is typically the **client**: it waits for the worker's
+    challenge, computes a response, and waits for confirmation.
+
+    Parameters
+    ----------
+    backend
+        A offwork backend instance.
+    pin
+        The PIN entered by the user.
+    timeout
+        Seconds to wait for the challenge and confirmation.
+
+    Raises
+    ------
+    PairingError
+        On timeout or if the initiator rejects the response.
+    """
+    intermediate = _derive_intermediate(pin)
+
+    # Wait for challenge
+    deadline = time.monotonic() + timeout
+    challenge: bytes | None = None
+    while time.monotonic() < deadline:
+        raw = await backend.get_progress(_PAIRING_CHANNEL)
+        if raw is not None:
+            try:
+                challenge = parse_challenge_message(raw)
+                break
+            except (ValueError, json.JSONDecodeError):
+                pass
+        await asyncio.sleep(0.5)
+
+    if challenge is None:
+        raise PairingError(f"Pairing timed out after {timeout}s — no challenge from peer")
+
+    # Compute and send response
+    response = compute_response(intermediate, challenge)
+    response_msg = make_response_message(response)
+    await backend.send_progress(_PAIRING_CHANNEL + ":response", response_msg)
+    logger.debug("Pairing: sent response")
+
+    # Wait for confirmation
+    while time.monotonic() < deadline:
+        raw = await backend.get_progress(_PAIRING_CHANNEL + ":confirm")
+        if raw is not None:
+            try:
+                parse_confirm_message(raw)
+                shared = derive_shared_secret(intermediate, challenge)
+                logger.info("Pairing successful (responder)")
+                return PairingResult(shared_key=shared, peer_role="worker")
+            except (ValueError, json.JSONDecodeError):
+                pass
+        await asyncio.sleep(0.5)
+
+    raise PairingError("Pairing failed — initiator did not confirm")
+

offwork/core/progress.py

@@ -0,0 +1,91 @@
+"""Progress reporting for remote task execution."""
+
+import json
+import contextvars
+from typing import Any, Self, overload
+from dataclasses import dataclass
+from collections.abc import Callable
+
+
+@dataclass(frozen=True)
+class ProgressInfo:
+    """Progress information reported by a running task.
+
+    Returned by :meth:`Result.progress`.
+    """
+
+    current: float = 0
+    total: float | None = None
+    message: str | None = None
+
+    @property
+    def percent(self) -> float | None:
+        """Return completion percentage, or ``None`` if *total* is unknown."""
+        if self.total is not None and self.total > 0:
+            return self.current / self.total * 100
+        return None
+
+    def to_json(self) -> str:
+        """Serialize to a JSON string."""
+        d: dict[str, Any] = {"current": self.current}
+        if self.total is not None:
+            d["total"] = self.total
+        if self.message is not None:
+            d["message"] = self.message
+        return json.dumps(d)
+
+    @classmethod
+    def from_json(cls, raw: str | bytes) -> Self:
+        """Deserialize from a JSON string or bytes."""
+        data = json.loads(raw)
+        return cls(
+            current=data["current"],
+            total=data.get("total"),
+            message=data.get("message"),
+        )
+
+
+_progress_callback: contextvars.ContextVar[
+    Callable[[float, float | None, str | None], None] | None
+] = contextvars.ContextVar("offwork_progress", default=None)
+
+
+@overload
+def progress(percent: float, /, *, message: str | None = None) -> None: ...
+
+
+@overload
+def progress(current: int, total: int, /, *, message: str | None = None) -> None: ...
+
+
+def progress(
+    _value: float,
+    _total: int | None = None,
+    /,
+    *,
+    message: str | None = None,
+) -> None:
+    """Report task progress from within a running function.
+
+    Call this inside a ``@trace``-decorated function to report progress
+    to the client.  When called outside a worker, this is a silent no-op.
+
+    Accepts either a percentage or a current/total pair::
+
+        progress(75.0)                          # 75 % complete
+        progress(75.0, message="loading model") # 75 % with message
+        progress(3, 10)                         # 3 of 10 (30 %)
+        progress(3, 10, message="step 3")       # 3 of 10 with message
+    """
+    cb = _progress_callback.get(None)
+    if cb is None:
+        return
+
+    if isinstance(_total, (int, float)):
+        if _total <= 0:
+            raise ValueError("Progress total must be a positive number.")
+        # current / total form
+        cb(_value, _total, message)
+    else:
+        # percent form — normalise to current/100
+        cb(_value, 100, message)

offwork/core/signing.py

@@ -0,0 +1,91 @@
+"""Cryptographic task signing and verification.
+
+After a client and worker are paired (see :mod:`offwork.core.pairing`),
+they share a secret key.  The client uses it to produce an HMAC-SHA256
+signature over the serialized task payload; the worker verifies that
+signature before executing the task.
+
+All primitives are stdlib-only (``hashlib``, ``hmac``, ``json``).
+"""
+
+import hmac
+import json
+import hashlib
+import logging
+from typing import Any
+
+from offwork.core.errors import SignatureError
+
+logger = logging.getLogger(__name__)
+
+
+def compute_signature(payload: str, key: bytes) -> str:
+    """Return a hex-encoded HMAC-SHA256 signature of *payload*.
+
+    Parameters
+    ----------
+    payload
+        The string to sign (typically the JSON body of a task).
+    key
+        Shared secret derived from the pairing process.
+    """
+    return hmac.new(key, payload.encode("utf-8"), hashlib.sha256).hexdigest()
+
+
+def verify_signature(payload: str, signature: str, key: bytes) -> bool:
+    """Verify an HMAC-SHA256 *signature* over *payload*.
+
+    Uses :func:`hmac.compare_digest` for constant-time comparison.
+    """
+    expected = hmac.new(key, payload.encode("utf-8"), hashlib.sha256).hexdigest()
+    return hmac.compare_digest(expected, signature)
+
+
+# -- Signed JSON helpers ----------------------------------------------------
+
+
+def sign_json(data: dict[str, Any], key: bytes) -> str:
+    """Serialize *data* to JSON, attach an HMAC-SHA256 signature, and return the envelope.
+
+    The returned JSON has the shape::
+
+        {"payload": "<inner-json>", "signature": "<hex>"}
+    """
+    payload = json.dumps(data, separators=(",", ":"), sort_keys=True)
+    sig = compute_signature(payload, key)
+    return json.dumps({"payload": payload, "signature": sig})
+
+
+def verify_and_load_json(envelope_json: str, key: bytes) -> dict[str, Any]:
+    """Parse *envelope_json*, verify the signature, and return the inner data.
+
+    Raises
+    ------
+    SignatureError
+        If the signature is missing or invalid.
+    """
+    try:
+        envelope = json.loads(envelope_json)
+    except (json.JSONDecodeError, TypeError) as exc:
+        raise SignatureError(f"Invalid signed envelope: {exc}") from exc
+
+    payload = envelope.get("payload")
+    signature = envelope.get("signature")
+    if payload is None or signature is None:
+        raise SignatureError("Envelope missing 'payload' or 'signature' field")
+
+    if not verify_signature(payload, signature, key):
+        raise SignatureError("HMAC signature verification failed")
+
+    return json.loads(payload)  # type: ignore[no-any-return]
+
+
+# -- Key derivation ---------------------------------------------------------
+
+
+def derive_key(shared_secret: bytes, context: str = "offwork-task-signing") -> bytes:
+    """Derive a 32-byte HMAC signing key from a shared secret.
+
+    Uses HKDF-like expansion via ``HMAC-SHA256(secret, context)``.
+    """
+    return hmac.new(shared_secret, context.encode("utf-8"), hashlib.sha256).digest()