open-standards/wai/wai-quantum-py/python/wai_quantum/keys.py

"""The key a receipt is signed with.

- ``Key.demo()``: the published demo key, for receipts that repeat byte for
  byte everywhere. Anyone can sign with it, so it proves nothing about who
  signed.
- ``Key.generate()``: a new key from the operating system's random source.
- ``Key.from_file(path)``, ``Key.from_hex(text)``, ``Key(seed)``: a key you
  keep.
- ``Key.user()``: this user's key, the one the command line signs with:
  the file ``$WAI_QUANTUM_KEY`` names, or else
  ``$XDG_CONFIG_HOME/wai/quantum-key`` (``~/.config/wai/quantum-key``),
  made on first use and readable by its owner only.

A key file holds the 32-byte Ed25519 seed as 64 hex digits, as the command
line's does, so one key signs on both.
"""

from __future__ import annotations

import os
import secrets
from pathlib import Path
from typing import Optional, Union

from . import _engine

__all__ = ["Key", "DEMO_SEED", "key_of", "user_key_path"]

#: The published demo key's seed (the command line's ``--demo-key``).
DEMO_SEED = bytes([42] * 32)


def user_key_path() -> Optional[Path]:
    """This user's key file, as the command line finds it."""
    base = os.environ.get("XDG_CONFIG_HOME")
    if base:
        return Path(base) / "wai" / "quantum-key"
    home = os.environ.get("HOME") or str(Path.home())
    return Path(home) / ".config" / "wai" / "quantum-key" if home else None


def _parse(text: str) -> Optional[bytes]:
    t = text.strip()
    if len(t) != 64:
        return None
    try:
        return bytes.fromhex(t)
    except ValueError:
        return None


class Key:
    """An Ed25519 signing key, by its 32-byte seed."""

    __slots__ = ("_seed", "path")

    def __init__(self, seed: bytes, path: Optional[Path] = None) -> None:
        seed = bytes(seed)
        if len(seed) != 32:
            raise ValueError(f"a signing key's seed is 32 bytes, not {len(seed)}")
        self._seed = seed
        #: The file the key was read from or made in, if any.
        self.path = path

    @classmethod
    def demo(cls) -> "Key":
        """The published demo key: a receipt signed with it shows the result
        is unaltered, not who signed it."""
        return cls(DEMO_SEED)

    @classmethod
    def generate(cls) -> "Key":
        """A new key from the operating system's random source."""
        return cls(secrets.token_bytes(32))

    @classmethod
    def from_hex(cls, text: str) -> "Key":
        seed = _parse(text)
        if seed is None:
            raise ValueError("a key is 64 hex digits")
        return cls(seed)

    @classmethod
    def from_file(cls, path: Union[str, os.PathLike]) -> "Key":
        p = Path(path)
        seed = _parse(p.read_text())
        if seed is None:
            raise ValueError(f"{p}: a key file holds 64 hex digits")
        return cls(seed, p)

    @classmethod
    def user(cls) -> "Key":
        """This user's key, the command line's: the file ``$WAI_QUANTUM_KEY``
        names (it must exist), or the user's own, made on first use."""
        named = os.environ.get("WAI_QUANTUM_KEY")
        if named:
            return cls.from_file(named)
        p = user_key_path()
        if p is None:
            raise ValueError("no home directory to keep a key in: pass key=")
        try:
            return cls.from_file(p)
        except FileNotFoundError:
            pass
        k = cls.generate()
        p.parent.mkdir(parents=True, exist_ok=True)
        try:
            k.save(p)
        except FileExistsError:
            # Made by another process meanwhile: that one is the key.
            return cls.from_file(p)
        k.path = p
        return k

    def save(self, path: Union[str, os.PathLike]) -> Path:
        """Write the key to a new file readable by its owner only; an existing
        file is never overwritten."""
        p = Path(path)
        fd = os.open(p, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
        with os.fdopen(fd, "w") as f:
            f.write(self._seed.hex() + "\n")
        return p

    @property
    def seed(self) -> bytes:
        """The 32-byte seed: the secret itself."""
        return self._seed

    @property
    def public_key(self) -> str:
        """The public key as 64 hex digits, as a receipt names it."""
        return _engine.public_key(self._seed)

    @property
    def did(self) -> str:
        """The public key as a ``did:key``."""
        return _engine.did_key(self.public_key)

    @property
    def is_demo(self) -> bool:
        return self._seed == DEMO_SEED

    def __eq__(self, other: object) -> bool:
        return isinstance(other, Key) and other._seed == self._seed

    def __hash__(self) -> int:
        return hash(self.public_key)

    def __repr__(self) -> str:
        what = "the published demo key" if self.is_demo else (str(self.path) if self.path else "in memory")
        return f"Key({self.did}, {what})"


def key_of(key: Union["Key", bytes, str, None]) -> Key:
    """A key from what a caller gives: a ``Key``, a 32-byte seed, 64 hex
    digits, or ``None`` for this user's key."""
    if key is None:
        return Key.user()
    if isinstance(key, Key):
        return key
    if isinstance(key, (bytes, bytearray, memoryview)):
        return Key(bytes(key))
    if isinstance(key, str):
        return Key.from_hex(key)
    raise TypeError(f"a key is a wai_quantum.Key, a 32-byte seed or 64 hex digits, not {type(key).__name__}")