h3cx/rex
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Harden and rewrite for release

Browse Source
This commit is contained in:
h3cx
2026-08-27 17:24:54 +01:00
parent 7e28074eff
commit d817b971e0
22 changed files with 719 additions and 1133 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+134 -10
View File
@@ -1,19 +1,143 @@
## Rex
# Rex
Run the public and management servers:
Rex is a small remote-execution relay. An API caller requests a named action; a connected device runs the matching, locally configured `argv` command. Commands never cross the network.
## What you need
- Python 3.13 or later and [uv](https://docs.astral.sh/uv/).
- A server reachable by the devices that need to connect.
- WSS/TLS for traffic beyond a trusted private network. Use the supplied Nginx example or configure TLS directly in Rex.
Install the project from its checkout:
```console
uv run rex server
uv sync
```
Open the terminal management interface in another shell:
Run each command below as `uv run rex ...` from the checkout, or install Rex into the appropriate environment and use `rex ...` directly.
## Quick start
1. Create the server configuration, allow the device name, and create two separate keys. The secret printed by `keys create` is shown only then.
```console
uv run rex server devices create desk
uv run rex server keys create desk-client --permission connect
uv run rex server keys create automation --permission execute
```
2. Copy [examples/server.toml](examples/server.toml) to `~/.config/rex/server.toml`; replace the placeholder secrets, adjust the listener addresses, and remove the example blacklist entries. The CLI already created this file, so edit it rather than overwriting it if you used step 1.
3. On the device, copy [examples/client.toml](examples/client.toml) to `~/.config/rex/client.toml`, set its `device_name` and the `desk-client` secret, then configure its local actions.
4. Start the server and client in separate environments:
```console
uv run rex server run
uv run rex client
```
5. Invoke a registered action using the execution key:
```console
curl --fail-with-body -X POST \
-H "X-API-Key: EXECUTION_SECRET" \
https://rex.example.net/device/desk/action/lock
```
The endpoint returns `{"queued":true}` only when the named device is connected and currently has that action registered. A disconnect immediately removes all of its actions.
## Configuration
Rex creates `~/.config/rex/server.toml` and `~/.config/rex/client.toml` on first use, with owner-only file permissions. Supply another path with `--config` immediately after `server` or `client`, for example `rex server --config /etc/rex/server.toml run`. If the selected configuration is absent, Rex prints a warning naming the missing role and suggesting the other one; this catches a common `rex server` versus `rex client` typo before it quietly creates the wrong template.
### Server (`server.toml`)
Copy [examples/server.toml](examples/server.toml). All fields are optional except fields within a `[[keys]]` record.
| Setting | Default | Meaning |
| --- | --- | --- |
| `public_host` / `public_port` | `127.0.0.1` / `8010` | Public HTTP API and device WebSocket listener. |
| `management_host` / `management_port` | `127.0.0.1` / `8011` | Internal status API listener. Keep it loopback-only. |
| `tls_certfile`, `tls_keyfile` | unset | Set both when Rex terminates TLS itself. Leave both unset behind Nginx. |
| `rate_limit_per_minute` | `120` | Requests allowed per direct peer in a rolling minute. |
| `blacklist` | `[]` | IP addresses or CIDR networks denied before authentication. |
| `devices` | `[]` | Device names permitted to open a WebSocket. |
| `[[keys]]` | none | A key’s name, secret, and permission list. |
Keys must have at least one permission:
| Permission | Allows |
| --- | --- |
| `connect` | Opening a device WebSocket. |
| `execute` | Calling `POST /device/{device}/action/{action}`. |
| `admin` | Reading the internal management status API; it also implies `connect` and `execute`. |
Use separate connect and execute keys. Keys are capability scoped, not device scoped: a connect key can connect as any name listed in `devices`, so issue one per client and revoke it when that client is retired.
### Client (`client.toml`)
Copy [examples/client.toml]. The client sends `device_name` and its action **names** during connection setup. Each `[[actions]]` entry has a `name` (letters, digits, `.`, `_`, or `-`; maximum 64 characters) and non-empty `argv` list. `argv[0]` is the executable. Rex uses `exec`, not a shell: shell syntax such as pipes, redirects, and `$VAR` expansion is deliberately unavailable.
`scheme` must be `ws` or `wss`. `port` is optional: omit it for a reverse proxy on the standard WSS port; set it for a direct listener such as `ws://192.0.2.10:8010`.
## TLS and Nginx
For an Nginx-terminated deployment, use [examples/nginx-rex.conf](examples/nginx-rex.conf), point the client at `wss://rex.example.net` with no port, and leave Rex TLS paths unset. The example exposes only the public listener on `127.0.0.1:8010`; do not proxy the management listener.
Rate limiting identifies Nginx as the direct peer, so all traffic through one proxy shares a bucket. Set a limit high enough for the expected aggregate traffic and use Nginx’s own per-client `limit_req` controls if per-client reverse-proxy limiting is needed.
For direct TLS, expose `public_host` deliberately (for example `0.0.0.0`), set both TLS paths, protect the private key, and configure clients with `scheme = "wss"` and `port = 8010`.
## Management CLI
The CLI changes the local server TOML; it does not call the network management API.
```console
uv run rex manage
rex server run
rex server keys list
rex server keys create laptop-client --permission connect
rex server keys create deploy-bot --permission execute
rex server keys create break-glass --permission admin
rex server keys delete laptop-client
rex server devices list
rex server devices create laptop
rex server devices delete laptop
rex client
rex client actions list
rex client actions create lock -- /usr/local/bin/lock-screen
rex client actions create wake-display -- /usr/bin/dpms force on
rex client actions delete wake-display
```
The manager connects to `http://127.0.0.1:8001` by default. Use
`uv run rex manage --url http://host:port` to connect elsewhere. The footer
shows all controls: `H/L` switch tabs, `h/j/k/l` move, `y` copies the selected
value, `n` creates a key or device, `a` adds an action, `d` deletes the
selection, `r` refreshes, and `q` quits.
`rex client actions` changes the local `client.toml` action allow-list; it never sends a command to the server. Use `--` before the executable so its arguments are unambiguously part of the action, especially if an argument begins with `-`. There is no TUI. To protect a secret, do not pass it as a command-line value; `keys create` generates it and prints it once.
## HTTP API
The public API has no browsable OpenAPI/docs endpoints in production.
| Request | Required key | Result |
| --- | --- | --- |
| `POST /device/{device}/action/{action}` | `execute` | Forwards the action name when it is currently registered. |
| `GET /keys` on the management listener | `admin` | Returns key names and permissions, never secrets. |
| `GET /devices` on the management listener | `admin` | Returns allowed devices and actions registered by connected clients. |
Supply `X-API-Key` for every authenticated request. A missing/incorrect key or insufficient permission returns `401`; a blacklisted or rate-limited peer returns `429`; an offline device or unavailable action returns `404`.
## Security and operations
- Use `wss` outside a fully trusted network. WebSocket/API keys are bearer credentials.
- Run each client under a dedicated, minimally privileged OS account. Anyone who can edit its TOML can configure a command for that device.
- Client commands are local allow-list entries. Remote callers cannot provide arguments or commands.
- Keep the management API on loopback. Its admin-key requirement is defense in depth, not a reason to expose it.
- Treat IP blacklisting as a coarse control, not authentication. Revoke compromised keys with `rex server keys delete NAME`.
- Review device action definitions before deployment. Rex intentionally runs the locally configured commands with the permissions of its client service account.
## Troubleshooting
| Symptom | Check |
| --- | --- |
| Client cannot connect | Device name is listed on the server, the key has `connect`, client URL/scheme/port match the listener, and the TLS certificate is valid for the endpoint. |
| Invocation returns 404 | The client is offline or did not register that exact action name. Check the client TOML and service logs. |
| Invocation returns 401 | Use an `execute` (or `admin`) key in `X-API-Key`; do not use the client’s `connect` key. |
| Invocation returns 429 | Check `blacklist` and the rolling limit. Behind Nginx, its aggregate proxy bucket may be full. |
examples/client.toml
+25
View File
@@ -0,0 +1,25 @@
# Copy to ~/.config/rex/client.toml on the device.
# This file contains the connect key and must be readable only by the client
# service account (chmod 600 ~/.config/rex/client.toml).
device_name = "desk"
api_key = "REPLACE_WITH_THE_DESK_CONNECT_KEY"
# WSS through an Nginx reverse proxy. No port means the normal HTTPS/WSS port.
endpoint = "rex.example.net"
scheme = "wss"
# For a direct server listener instead, use for example:
# endpoint = "192.0.2.10"
# scheme = "ws" # only on a trusted network
# port = 8010
# Only names are sent to the server. argv remains local and is executed without
# a shell, so each action must name an executable followed by its arguments.
[[actions]]
name = "lock"
argv = ["/usr/local/bin/lock-screen"]
[[actions]]
name = "wake-display"
argv = ["/usr/bin/dpms", "force", "on"]
examples/nginx-rex.conf
+20
View File
@@ -0,0 +1,20 @@
# /etc/nginx/sites-available/rex.conf
# Public API and WebSocket only. The management listener remains loopback-only.
server {
listen 443 ssl http2;
server_name rex.example.net;
ssl_certificate /etc/letsencrypt/live/rex.example.net/fullchain.pem;
ssl_certificate_key /etc/letsencrypt/live/rex.example.net/privkey.pem;
location / {
proxy_pass http://127.0.0.1:8010;
proxy_http_version 1.1;
proxy_set_header Host $host;
proxy_set_header X-Real-IP $remote_addr;
proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
proxy_set_header X-Forwarded-Proto $scheme;
proxy_set_header Upgrade $http_upgrade;
proxy_set_header Connection "upgrade";
}
}
examples/server.toml
+42
View File
@@ -0,0 +1,42 @@
# Copy to ~/.config/rex/server.toml on the server, then replace every secret.
# Rex writes this file with mode 0600 when it creates or updates it.
# Public API listener. Keep 127.0.0.1 when Nginx terminates TLS on this host.
public_host = "127.0.0.1"
public_port = 8010
# Internal status API. Do not expose this through a reverse proxy.
management_host = "127.0.0.1"
management_port = 8011
# Uncomment both lines only if Rex, rather than a reverse proxy, terminates TLS.
# tls_certfile = "/etc/letsencrypt/live/rex.example.net/fullchain.pem"
# tls_keyfile = "/etc/letsencrypt/live/rex.example.net/privkey.pem"
# Per direct peer address, over a rolling 60-second window.
rate_limit_per_minute = 120
# Individual IP addresses and CIDR networks are accepted.
blacklist = ["203.0.113.24", "198.51.100.0/24"]
# Only listed names may connect a device WebSocket.
devices = ["desk"]
# Use a distinct secret for each role. Generate production keys with:
# rex server keys create <name> --permission <connect|execute|admin>
[[keys]]
name = "desk-client"
secret = "REPLACE_WITH_A_CONNECT_KEY"
permissions = ["connect"]
[[keys]]
name = "automation"
secret = "REPLACE_WITH_AN_EXECUTE_KEY"
permissions = ["execute"]
# Optional: this key can access the loopback-only management status API and
# also has connect and execute rights. Do not use it in a client or automation.
[[keys]]
name = "local-admin"
secret = "REPLACE_WITH_AN_ADMIN_KEY"
permissions = ["admin"]
pyproject.toml
-1
View File
@@ -11,7 +11,6 @@ dependencies = [
"fastapi[standard]>=0.141.1",
"httpx>=0.28.1",
"pydantic>=2.13.4",
"textual>=6.1.0",
"uvicorn>=0.52.4",
"websockets>=17.0.1",
]
src/rex/__init__.py
+119 -22
View File
@@ -1,29 +1,126 @@
import argparse
import sys
from pathlib import Path
from rex.client.config import ClientAction, ClientConfigManager
from rex.server.config import ServerConfigManager
def _config_argument(parser: argparse.ArgumentParser, name: str) -> None:
parser.add_argument("--config", default=f"~/.config/rex/{name}.toml", help="config file (default: %(default)s)")
def _warn_missing_config(config_path: str, role: str) -> None:
path = Path(config_path).expanduser()
if not path.exists():
other_role = "client" if role == "server" else "server"
print(
f"warning: {path} is missing; Rex will create a {role}.toml template. "
f"If this was unintentional, did you mean `rex {other_role} ...`?",
file=sys.stderr,
)
def _server_key(args: argparse.Namespace) -> None:
manager = ServerConfigManager(args.config)
if args.key_command == "list":
for key in manager.config.keys:
print(f"{key.name}\t{','.join(sorted(key.permissions))}")
elif args.key_command == "create":
key = manager.create_key(args.name, set(args.permission))
if key is None:
raise SystemExit(f"key already exists: {args.name}")
print(f"{key.name}\t{key.secret}\t{','.join(sorted(key.permissions))}")
elif not manager.delete_key(args.name):
raise SystemExit(f"key not found: {args.name}")
def _server_device(args: argparse.Namespace) -> None:
manager = ServerConfigManager(args.config)
if args.device_command == "list":
print("\n".join(manager.config.devices))
elif args.device_command == "create":
if args.name in manager.config.devices:
raise SystemExit(f"device already exists: {args.name}")
manager.config.devices.append(args.name)
manager.save()
else:
try:
manager.config.devices.remove(args.name)
except ValueError:
raise SystemExit(f"device not found: {args.name}")
manager.save()
def _client_action(args: argparse.Namespace) -> None:
manager = ClientConfigManager(args.config)
if args.action_command == "list":
for action in manager.config.actions:
print(f"{action.name}\t{' '.join(action.argv)}")
elif args.action_command == "create":
if any(action.name == args.name for action in manager.config.actions):
raise SystemExit(f"action already exists: {args.name}")
argv = args.argv[1:] if args.argv[:1] == ["--"] else args.argv
try:
action = ClientAction(name=args.name, argv=argv)
except ValueError as error:
raise SystemExit(f"invalid action: {error}") from error
manager.config.actions.append(action)
manager.save()
else:
for action in manager.config.actions:
if action.name == args.name:
manager.config.actions.remove(action)
manager.save()
return
raise SystemExit(f"action not found: {args.name}")
def main() -> None:
parser = argparse.ArgumentParser(prog="rex")
subparsers = parser.add_subparsers(dest="command", required=True)
subparsers.add_parser("server", help="run the Rex server")
subparsers.add_parser("client", help="run the Rex client")
manage_parser = subparsers.add_parser("manage", help="manage a Rex server")
manage_parser.add_argument(
"--url",
default="http://127.0.0.1:8011",
help="management API URL (default: %(default)s)",
)
parser = argparse.ArgumentParser(prog="rex", description="Remote execution with locally configured client actions.")
root = parser.add_subparsers(dest="role", required=True)
server = root.add_parser("server", help="run or manage a Rex server")
_config_argument(server, "server")
server_commands = server.add_subparsers(dest="server_command")
server_commands.add_parser("run", help="run public and internal servers")
keys = server_commands.add_parser("keys", help="manage API keys")
key_commands = keys.add_subparsers(dest="key_command", required=True)
key_commands.add_parser("list")
create_key = key_commands.add_parser("create")
create_key.add_argument("name")
create_key.add_argument("--permission", action="append", choices=("admin", "connect", "execute"), required=True)
delete_key = key_commands.add_parser("delete")
delete_key.add_argument("name")
devices = server_commands.add_parser("devices", help="manage allowed device names")
device_commands = devices.add_subparsers(dest="device_command", required=True)
device_commands.add_parser("list")
create_device = device_commands.add_parser("create")
create_device.add_argument("name")
delete_device = device_commands.add_parser("delete")
delete_device.add_argument("name")
client = root.add_parser("client", help="run or manage a Rex client")
_config_argument(client, "client")
client_commands = client.add_subparsers(dest="client_command")
client_commands.add_parser("run", help="run the client")
actions = client_commands.add_parser("actions", help="manage locally executable actions")
action_commands = actions.add_subparsers(dest="action_command", required=True)
action_commands.add_parser("list")
create_action = action_commands.add_parser("create")
create_action.add_argument("name")
create_action.add_argument("argv", nargs=argparse.REMAINDER, help="command and arguments; prefix with -- when needed")
delete_action = action_commands.add_parser("delete")
delete_action.add_argument("name")
args = parser.parse_args()
if args.command == "server":
_warn_missing_config(args.config, args.role)
if args.role == "server" and args.server_command in (None, "run"):
from rex.server import main as server_main
server_main()
elif args.command == "client":
server_main(args.config)
elif args.role == "server" and args.server_command == "keys":
_server_key(args)
elif args.role == "server":
_server_device(args)
elif args.role == "client" and args.client_command == "actions":
_client_action(args)
else:
from rex.client import main as client_main
client_main()
elif args.command == "manage":
from rex.management import main as management_main
management_main(args.url)
client_main(args.config)
src/rex/client/__init__.py
+28 -74
View File
@@ -1,25 +1,17 @@
import asyncio
import json
import os
import signal
from pathlib import Path
import websockets
from rex.client.action_manager import ClientActionManager
from rex.client.config_manager import ClientConfigManager
from rex.server.device_manager import Action
from rex.client.config import ClientAction, ClientConfigManager
client_config = ClientConfigManager("~/.config/rex/client.json")
action_manager = ClientActionManager("~/.config/rex/actions.json")
async def run_command(command: str) -> None:
process = await asyncio.create_subprocess_shell(
command,
start_new_session=os.name == "posix",
)
async def run_command(action: ClientAction) -> None:
"""Run only a locally configured argv list; never parse remote shell input."""
process = await asyncio.create_subprocess_exec(*action.argv, start_new_session=os.name == "posix")
try:
await process.wait()
except asyncio.CancelledError:
@@ -31,7 +23,6 @@ async def run_command(command: str) -> None:
pass
else:
process.terminate()
try:
await asyncio.wait_for(process.wait(), timeout=5)
except TimeoutError:
@@ -46,64 +37,27 @@ async def run_command(command: str) -> None:
raise
async def receive_actions(websocket: websockets.ClientConnection) -> None:
command_tasks: set[asyncio.Task[None]] = set()
try:
async for message in websocket:
action = Action.model_validate_json(message)
client_action = action_manager.get_action(action.name)
if client_action is not None:
task = asyncio.create_task(run_command(client_action.command))
command_tasks.add(task)
task.add_done_callback(command_tasks.discard)
except websockets.ConnectionClosed:
pass
finally:
for task in command_tasks:
task.cancel()
await asyncio.gather(*command_tasks, return_exceptions=True)
async def listen(config_path: str | Path = "~/.config/rex/client.toml") -> None:
config = ClientConfigManager(config_path).config
actions = {action.name: action for action in config.actions}
if not config.api_key:
raise RuntimeError("client.toml has no api_key")
async with websockets.connect(config.websocket_url(), additional_headers={"X-API-Key": config.api_key}) as websocket:
await websocket.send({"type": "register", "actions": sorted(actions)})
tasks: set[asyncio.Task[None]] = set()
try:
async for raw in websocket:
message = json.loads(raw)
action = actions.get(message.get("name")) if message.get("type") == "action" else None
if action:
task = asyncio.create_task(run_command(action))
tasks.add(task)
task.add_done_callback(tasks.discard)
finally:
for task in tasks:
task.cancel()
await asyncio.gather(*tasks, return_exceptions=True)
async def listen() -> None:
uri = f"ws://{client_config.get_endpoint()}:{client_config.get_port()}/ws/{client_config.get_name()}"
print(uri)
stop = asyncio.Event()
loop = asyncio.get_running_loop()
handles_sigterm = False
try:
loop.add_signal_handler(signal.SIGTERM, stop.set)
handles_sigterm = True
except (NotImplementedError, RuntimeError, ValueError):
pass
try:
async with websockets.connect(
uri, additional_headers={"X-API-Key": client_config.get_key()}
) as websocket:
print("Connected")
receiver = asyncio.create_task(receive_actions(websocket))
shutdown = asyncio.create_task(stop.wait())
try:
done, _ = await asyncio.wait(
(receiver, shutdown), return_when=asyncio.FIRST_COMPLETED
)
if receiver in done:
await receiver
finally:
receiver.cancel()
shutdown.cancel()
await asyncio.gather(receiver, shutdown, return_exceptions=True)
finally:
if handles_sigterm:
loop.remove_signal_handler(signal.SIGTERM)
def main() -> None:
print(f"Rex client name: {client_config.get_name()}")
print(f"API Key: {client_config.get_key()}")
asyncio.run(listen())
def main(config_path: str = "~/.config/rex/client.toml") -> None:
asyncio.run(listen(config_path))
src/rex/client/action_manager.py
-56
View File
@@ -1,56 +0,0 @@
from pathlib import Path
from pydantic import BaseModel
class ClientAction(BaseModel):
name: str
command: str
class ClientActionStore(BaseModel):
actions: list[ClientAction]
class ClientActionManager:
def __init__(self, file_path: str | Path) -> None:
self.file_path = Path(file_path).expanduser()
if not self.file_path.exists():
self.action_store = ClientActionStore(actions=[])
self.file_path.parent.mkdir(parents=True, exist_ok=True)
self.file_path.write_text(self.action_store.model_dump_json(indent=4))
self.action_store = ClientActionStore.model_validate_json(
self.file_path.read_text()
)
def create_action(self, action_name: str, command: str) -> ClientAction | None:
if self.get_action(action_name) is not None:
return None
new_action = ClientAction(name=action_name, command=command)
self.action_store.actions.append(new_action)
self._save()
return new_action
def _save(self) -> None:
self.file_path.parent.mkdir(parents=True, exist_ok=True)
self.file_path.write_text(self.action_store.model_dump_json(indent=4))
def get_action(self, action_name: str) -> ClientAction | None:
for action in self.action_store.actions:
if action_name == action.name:
return action
return None
def get_all_actions(self) -> list[ClientAction]:
return self.action_store.actions
def delete_action(self, action_name: str) -> bool:
action = self.get_action(action_name)
if action is None:
return False
self.action_store.actions.remove(action)
self._save()
return True
src/rex/client/config.py
+52
View File
@@ -0,0 +1,52 @@
from __future__ import annotations
from pathlib import Path
from urllib.parse import urlunsplit
from pydantic import BaseModel, Field, field_validator
from rex.config import read_toml, write_toml
class ClientAction(BaseModel):
name: str = Field(pattern=r"^[A-Za-z0-9][A-Za-z0-9_.-]{0,63}$")
argv: list[str] = Field(min_length=1, max_length=64)
@field_validator("argv")
@classmethod
def nonempty_arguments(cls, value: list[str]) -> list[str]:
if not value[0] or any("\x00" in item for item in value):
raise ValueError("argv must start with an executable and contain no NUL bytes")
return value
class ClientConfig(BaseModel):
device_name: str = Field(default="default-device", pattern=r"^[A-Za-z0-9][A-Za-z0-9_.-]{0,63}$")
api_key: str = ""
endpoint: str = "127.0.0.1"
scheme: str = "wss"
port: int | None = None
actions: list[ClientAction] = Field(default_factory=list)
@field_validator("scheme")
@classmethod
def valid_scheme(cls, value: str) -> str:
if value not in {"ws", "wss"}:
raise ValueError("scheme must be ws or wss")
return value
def websocket_url(self) -> str:
host = self.endpoint
if ":" in host and not host.startswith("["):
host = f"[{host}]"
authority = f"{host}:{self.port}" if self.port is not None else host
return urlunsplit((self.scheme, authority, f"/ws/{self.device_name}", "", ""))
class ClientConfigManager:
def __init__(self, file_path: str | Path) -> None:
self.file_path = Path(file_path).expanduser()
self.config = ClientConfig.model_validate(read_toml(self.file_path, ClientConfig().model_dump(mode="json")))
def save(self) -> None:
write_toml(self.file_path, self.config.model_dump(mode="json"))
src/rex/client/config_manager.py
-38
View File
@@ -1,38 +0,0 @@
from pathlib import Path
from pydantic import BaseModel
class ClientConfig(BaseModel):
device_name: str
api_key: str
endpoint: str
port: int
class ClientConfigManager:
def __init__(self, file_path: str | Path) -> None:
self.file_path = Path(file_path).expanduser()
if not self.file_path.exists():
self.client_config = ClientConfig(device_name="default_device")
self.file_path.parent.mkdir(parents=True, exist_ok=True)
self.file_path.write_text(self.client_config.model_dump_json(indent=4))
self.client_config = ClientConfig.model_validate_json(
self.file_path.read_text()
)
def _save(self) -> None:
self.file_path.parent.mkdir(parents=True, exist_ok=True)
self.file_path.write_text(self.client_config.model_dump_json(indent=4))
def get_name(self) -> str:
return self.client_config.device_name
def get_key(self) -> str:
return self.client_config.api_key
def get_endpoint(self) -> str:
return self.client_config.endpoint
def get_port(self) -> int:
return self.client_config.port
src/rex/config.py
+66
View File
@@ -0,0 +1,66 @@
"""Small, dependency-free TOML configuration helpers."""
from __future__ import annotations
import os
import tempfile
import tomllib
from pathlib import Path
from typing import Any
def read_toml(path: str | Path, default: dict[str, Any]) -> dict[str, Any]:
path = Path(path).expanduser()
if not path.exists():
write_toml(path, default)
return default
with path.open("rb") as file:
return tomllib.load(file)
def write_toml(path: str | Path, value: dict[str, Any]) -> None:
"""Atomically write the limited TOML shapes used by Rex's config files."""
path = Path(path).expanduser()
path.parent.mkdir(parents=True, exist_ok=True)
fd, temporary_name = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent)
try:
with os.fdopen(fd, "w", encoding="utf-8") as file:
for key, item in value.items():
if item is None:
continue
if isinstance(item, list) and not item:
continue
if isinstance(item, dict):
continue
if isinstance(item, list) and item and isinstance(item[0], dict):
continue
file.write(f"{key} = {_toml_value(item)}\n")
for key, item in value.items():
if isinstance(item, dict):
file.write(f"\n[{key}]\n")
for child_key, child_value in item.items():
if child_value is None:
continue
file.write(f"{child_key} = {_toml_value(child_value)}\n")
elif isinstance(item, list) and item and isinstance(item[0], dict):
for record in item:
file.write(f"\n[[{key}]]\n")
for child_key, child_value in record.items():
if child_value is None:
continue
file.write(f"{child_key} = {_toml_value(child_value)}\n")
os.replace(temporary_name, path)
path.chmod(0o600)
finally:
if os.path.exists(temporary_name):
os.unlink(temporary_name)
def _toml_value(value: Any) -> str:
if isinstance(value, bool):
return str(value).lower()
if isinstance(value, (str, int, float)):
return repr(value)
if isinstance(value, list):
return "[" + ", ".join(_toml_value(item) for item in value) + "]"
raise TypeError(f"Unsupported configuration value: {value!r}")
src/rex/management/__init__.py
-5
View File
@@ -1,5 +0,0 @@
from rex.management.app import RexManager
def main(url: str) -> None:
RexManager(url).run()
src/rex/management/app.py
-264
View File
@@ -1,264 +0,0 @@
from collections.abc import Callable, Coroutine
from typing import Any, ClassVar
from textual.app import App, ComposeResult
from textual.containers import Container
from textual.screen import ModalScreen
from textual.widgets import (
DataTable,
Footer,
Header,
Input,
Static,
TabbedContent,
TabPane,
Tree,
)
from rex.management.client import ManagementClient, ManagementError
class NamePrompt(ModalScreen[str | None]):
BINDINGS: ClassVar = [("escape", "cancel", "Cancel")]
CSS = """
NamePrompt { align: center middle; background: $background 60%; }
#prompt-dialog {
width: 60;
max-width: 90%;
height: auto;
padding: 1 2;
border: round $primary;
background: $surface;
}
#prompt-label { margin-bottom: 1; }
"""
def __init__(self, prompt: str) -> None:
super().__init__()
self.prompt = prompt
def compose(self) -> ComposeResult:
with Container(id="prompt-dialog"):
yield Static(self.prompt, id="prompt-label")
yield Input(placeholder="Type a name and press Enter", id="prompt-input")
def on_mount(self) -> None:
self.query_one(Input).focus()
def on_input_submitted(self, event: Input.Submitted) -> None:
name = event.value.strip()
if name:
self.dismiss(name)
def action_cancel(self) -> None:
self.dismiss(None)
class RexManager(App[None]):
TITLE = "Rex Manager"
SUB_TITLE = "Management API"
CSS = """
Screen { background: $surface; }
TabbedContent { margin: 1 2; }
TabPane { padding: 1 2; }
DataTable, Tree { height: 1fr; border: round $primary; }
#status { height: auto; padding: 0 3 1 3; color: $text-muted; }
"""
BINDINGS: ClassVar = [
("shift+h", "previous_tab", "Previous tab"),
("shift+l", "next_tab", "Next tab"),
("h", "move_left", "Move left"),
("l", "move_right", "Move right"),
("j", "move_down", "Move down"),
("k", "move_up", "Move up"),
("y", "yank", "Yank"),
("n", "new", "New"),
("a", "add_action", "Add action"),
("d", "delete", "Delete"),
("r", "refresh", "Refresh"),
("q", "quit", "Quit"),
]
def __init__(self, url: str) -> None:
super().__init__()
self.url = url.rstrip("/")
self.client = ManagementClient(self.url)
def compose(self) -> ComposeResult:
yield Header()
with TabbedContent():
with TabPane("API keys", id="keys-tab"):
yield DataTable(id="keys")
with TabPane("Devices", id="devices-tab"):
yield Tree("Devices", id="devices")
yield Static(f"Connecting to {self.url}", id="status")
yield Footer()
async def on_mount(self) -> None:
keys = self.query_one("#keys", DataTable)
keys.add_columns("Name", "Secret")
await self.action_refresh()
keys.focus()
async def on_unmount(self) -> None:
await self.client.close()
async def action_refresh(self) -> None:
try:
keys, devices = await self.client.keys(), await self.client.devices()
except ManagementError as error:
self._set_status(str(error), error=True)
return
key_table = self.query_one("#keys", DataTable)
key_table.clear()
for key in keys:
key_table.add_row(key["name"], key["secret"], key=key["name"])
tree = self.query_one("#devices", Tree)
tree.clear()
for device in devices:
device_name = device["name"]
device_node = tree.root.add(device_name, data=("device", device_name))
for action in device["actions"]:
device_node.add_leaf(
action["name"], data=("action", device_name, action["name"])
)
device_node.expand()
tree.root.expand()
self._set_status(f"Connected to {self.url}")
def action_previous_tab(self) -> None:
self._select_tab("keys-tab")
def action_next_tab(self) -> None:
self._select_tab("devices-tab")
def action_move_down(self) -> None:
self._active_widget().action_cursor_down()
def action_move_up(self) -> None:
self._active_widget().action_cursor_up()
def action_move_left(self) -> None:
if self.query_one(TabbedContent).active == "keys-tab":
self.query_one("#keys", DataTable).action_cursor_left()
def action_move_right(self) -> None:
if self.query_one(TabbedContent).active == "keys-tab":
self.query_one("#keys", DataTable).action_cursor_right()
def action_yank(self) -> None:
if self.query_one(TabbedContent).active == "keys-tab":
table = self.query_one("#keys", DataTable)
if table.row_count == 0:
self._set_status("Select a key first", error=True)
return
value = str(table.get_cell_at(table.cursor_coordinate))
else:
selection = self._tree_selection()
if selection is None:
self._set_status("Select a device or action first", error=True)
return
value = selection[-1]
self.copy_to_clipboard(value)
self._set_status("Copied selected value")
def action_new(self) -> None:
tabs = self.query_one(TabbedContent)
if tabs.active == "keys-tab":
self._prompt("New API key", self._create_key)
else:
self._prompt("New device", self._create_device)
def action_add_action(self) -> None:
if self.query_one(TabbedContent).active != "devices-tab":
self._set_status("Switch to Devices to add an action", error=True)
return
selection = self._tree_selection()
if selection is None:
self._set_status("Select a device first", error=True)
return
device = selection[1]
self._prompt(
f"New action for {device}",
lambda name: self._create_action(device, name),
)
def action_delete(self) -> None:
if self.query_one(TabbedContent).active == "keys-tab":
request = self._delete_key()
else:
request = self._delete_node()
if request is not None:
self.run_worker(request)
def _create_key(self, name: str) -> None:
self.run_worker(self._change(self.client.create_key(name), "Key created"))
def _create_device(self, name: str) -> None:
self.run_worker(self._change(self.client.create_device(name), "Device created"))
def _create_action(self, device: str, name: str) -> None:
self.run_worker(
self._change(self.client.create_action(device, name), "Action created")
)
def _delete_key(self) -> Coroutine[Any, Any, None] | None:
table = self.query_one("#keys", DataTable)
if table.row_count == 0:
self._set_status("Select a key first", error=True)
return None
name = str(table.get_row_at(table.cursor_row)[0])
return self._change(self.client.delete_key(name), "Key deleted")
def _delete_node(self) -> Coroutine[Any, Any, None] | None:
selection = self._tree_selection()
if selection is None:
self._set_status("Select a device or action first", error=True)
return None
if selection[0] == "device":
request = self.client.delete_device(selection[1])
message = "Device deleted"
else:
request = self.client.delete_action(selection[1], selection[2])
message = "Action deleted"
return self._change(request, message)
async def _change(self, request: Any, success: str) -> None:
try:
await request
except ManagementError as error:
self._set_status(str(error), error=True)
return
await self.action_refresh()
self._set_status(success)
def _prompt(self, label: str, callback: Callable[[str], None]) -> None:
def submitted(name: str | None) -> None:
if name is not None:
callback(name)
self.push_screen(NamePrompt(label), submitted)
def _select_tab(self, tab: str) -> None:
tabs = self.query_one(TabbedContent)
tabs.active = tab
self._active_widget().focus()
def _active_widget(self) -> DataTable[Any] | Tree[Any]:
if self.query_one(TabbedContent).active == "keys-tab":
return self.query_one("#keys", DataTable)
return self.query_one("#devices", Tree)
def _tree_selection(self) -> tuple[str, ...] | None:
return self.query_one("#devices", Tree).cursor_node.data
def _set_status(self, message: str, *, error: bool = False) -> None:
status = self.query_one("#status", Static)
status.update(message)
status.styles.color = "red" if error else "green"
src/rex/management/client.py
-65
View File
@@ -1,65 +0,0 @@
from typing import Any
from urllib.parse import quote
import httpx
class ManagementError(Exception):
pass
class ManagementClient:
def __init__(
self, base_url: str, *, transport: httpx.AsyncBaseTransport | None = None
) -> None:
self._client = httpx.AsyncClient(
base_url=base_url.rstrip("/"), timeout=5, transport=transport
)
async def close(self) -> None:
await self._client.aclose()
async def keys(self) -> list[dict[str, str]]:
return (await self._request("GET", "/keys"))["keys"]
async def create_key(self, name: str) -> None:
await self._request("POST", f"/key/{quote(name, safe='')}")
async def delete_key(self, name: str) -> None:
await self._request("DELETE", f"/key/{quote(name, safe='')}")
async def devices(self) -> list[dict[str, Any]]:
return (await self._request("GET", "/devices"))["devices"]
async def create_device(self, name: str) -> None:
await self._request("POST", f"/device/{quote(name, safe='')}")
async def delete_device(self, name: str) -> None:
await self._request("DELETE", f"/device/{quote(name, safe='')}")
async def create_action(self, device: str, name: str) -> None:
await self._request(
"POST", f"/device/{quote(device, safe='')}/action/{quote(name, safe='')}"
)
async def delete_action(self, device: str, name: str) -> None:
await self._request(
"DELETE", f"/device/{quote(device, safe='')}/action/{quote(name, safe='')}"
)
async def _request(self, method: str, path: str) -> Any:
try:
response = await self._client.request(method, path)
response.raise_for_status()
except httpx.HTTPStatusError as error:
try:
message = error.response.json().get("detail", str(error))
except ValueError:
message = str(error)
raise ManagementError(message) from error
except httpx.HTTPError as error:
raise ManagementError(f"Could not reach the server: {error}") from error
if not response.content:
return None
return response.json()
src/rex/server/__init__.py
+95 -193
View File
@@ -1,214 +1,116 @@
import asyncio
import time
from collections import defaultdict, deque
from pathlib import Path
import uvicorn
from fastapi import (
Depends,
FastAPI,
HTTPException,
Header,
WebSocket,
WebSocketDisconnect,
WebSocketException,
status,
)
from fastapi import Depends, FastAPI, HTTPException, Request, WebSocket, WebSocketDisconnect, WebSocketException, status
from pydantic import BaseModel, Field, field_validator
from rex.server.config import ServerConfigManager
from rex.server.connection_manager import ConnectionManager
from rex.server.device_manager import DeviceManager
from rex.server.key_manager import KeyManager
public_app = FastAPI()
management_app = FastAPI()
class Registration(BaseModel):
type: str
actions: set[str] = Field(max_length=256)
key_manager = KeyManager("~/.config/rex/keys.json")
device_manager = DeviceManager("~/.config/rex/devices.json")
connection_manager = ConnectionManager()
@field_validator("actions")
@classmethod
def action_names_are_safe(cls, actions: set[str]) -> set[str]:
import re
if any(not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_.-]{0,63}", action) for action in actions):
raise ValueError("invalid action name")
return actions
async def verify_api_key(x_api_key: str = Header()):
if not key_manager.compare_key(x_api_key):
raise HTTPException(
status_code=401,
detail="Invalid API key",
)
class RateLimiter:
def __init__(self, per_minute: int) -> None:
self.per_minute = per_minute
self.requests: dict[str, deque[float]] = defaultdict(deque)
def allowed(self, identity: str) -> bool:
now = time.monotonic()
requests = self.requests[identity]
while requests and requests[0] <= now - 60:
requests.popleft()
if len(requests) >= self.per_minute:
return False
requests.append(now)
return True
async def verify_api_key_ws(websocket: WebSocket) -> None:
api_key = websocket.headers.get("x-api-key")
def create_apps(config_path: str | Path = "~/.config/rex/server.toml") -> tuple[FastAPI, FastAPI]:
manager = ServerConfigManager(config_path)
connections = ConnectionManager()
limiter = RateLimiter(manager.config.rate_limit_per_minute)
if api_key is None or not key_manager.compare_key(api_key):
raise WebSocketException(
code=status.WS_1008_POLICY_VIOLATION,
reason="Invalid API key",
)
async def authorize(request: Request, permission: str) -> None:
host = request.client.host if request.client else None
if manager.config.is_blacklisted(host) or not limiter.allowed(host or "unknown"):
raise HTTPException(429, "Request denied")
if not manager.get_key(request.headers.get("x-api-key", ""), permission):
raise HTTPException(401, "Invalid API key or insufficient permission")
def require(permission: str):
async def dependency(request: Request) -> None:
await authorize(request, permission)
return dependency
public = FastAPI(title="Rex public API", docs_url=None, redoc_url=None, openapi_url=None)
management = FastAPI(title="Rex management API", docs_url=None, redoc_url=None, openapi_url=None)
@public.websocket("/ws/{device_name}")
async def websocket_endpoint(websocket: WebSocket, device_name: str):
host = websocket.client.host if websocket.client else None
key = websocket.headers.get("x-api-key", "")
denied = manager.config.is_blacklisted(host) or not limiter.allowed(host or "unknown")
if denied or device_name not in manager.config.devices or not manager.get_key(key, "connect"):
raise WebSocketException(status.WS_1008_POLICY_VIOLATION, "Connection denied")
await websocket.accept()
try:
registration = Registration.model_validate(await asyncio.wait_for(websocket.receive_json(), timeout=10))
if registration.type != "register":
raise ValueError("first message must register actions")
except (TimeoutError, ValueError):
await websocket.close(code=status.WS_1008_POLICY_VIOLATION, reason="Invalid registration")
return
await connections.connect(device_name, websocket, registration.actions)
try:
while True:
await websocket.receive_text()
except WebSocketDisconnect:
connections.disconnect(device_name, websocket)
@public.post("/device/{device_name}/action/{action_name}", dependencies=[Depends(require("execute"))])
async def send_action(device_name: str, action_name: str):
if not await connections.send(device_name, action_name):
raise HTTPException(404, "Device is offline or action is not registered")
return {"queued": True}
@management.get("/keys", dependencies=[Depends(require("admin"))])
async def keys():
return {"keys": [{"name": key.name, "permissions": sorted(key.permissions)} for key in manager.config.keys]}
@management.get("/devices", dependencies=[Depends(require("admin"))])
async def devices():
return {"devices": [{"name": name, "actions": connections.actions(name)} for name in manager.config.devices]}
return public, management
@public_app.websocket("/ws/{device_name}")
async def websocket_endpoint(
websocket: WebSocket, device_name: str, _: None = Depends(verify_api_key_ws)
):
if device_manager.get_device(device_name) is None:
raise WebSocketException(
code=status.WS_1008_POLICY_VIOLATION,
reason="Device not found",
)
await connection_manager.connect(device_name, websocket)
try:
while True:
await websocket.receive_json()
except WebSocketDisconnect:
connection_manager.disconnect(device_name)
@public_app.get("/device/{device_name}/action/{action_name}")
async def send_action(
device_name: str, action_name: str, _: None = Depends(verify_api_key)
):
action = device_manager.get_action(device_name, action_name)
if action is None:
raise HTTPException(
status_code=404,
detail="Action not found for this device",
)
await connection_manager.send(device_name, action)
@management_app.get("/tryauth")
async def tryauth(_: None = Depends(verify_api_key)):
return
@management_app.get("/keys")
async def get_keys():
return {"keys": key_manager.get_all_keys()}
@management_app.get("/key/{key_name}")
async def get_key_by_name(key_name: str):
key = key_manager.get_key(key_name)
if key is None:
raise HTTPException(
status_code=404,
detail="Key not found",
)
return {"key": key}
@management_app.post("/key/{key_name}")
async def create_key(key_name: str):
key = key_manager.create_key(key_name)
if key is None:
raise HTTPException(status_code=409, detail="Key already exists with this name")
return {"key": key}
@management_app.delete("/key/{key_name}")
async def delete_key(key_name: str):
val = key_manager.delete_key(key_name)
if not val:
raise HTTPException(status_code=404, detail="No key found with this name")
return
@management_app.get("/devices")
async def get_devices():
return {"devices": device_manager.get_all_devices()}
@management_app.get("/device/{device_name}")
async def get_device_by_name(device_name: str):
device = device_manager.get_device(device_name)
if device is None:
raise HTTPException(
status_code=404,
detail="Device not found",
)
return {"device": device}
@management_app.post("/device/{device_name}")
async def create_device(device_name: str):
device = device_manager.create_device(device_name)
if device is None:
raise HTTPException(
status_code=409, detail="Device already exists with this name"
)
return {"device": device}
@management_app.delete("/device/{device_name}")
async def delete_device(device_name: str):
val = device_manager.delete_device(device_name)
if not val:
raise HTTPException(status_code=404, detail="No device found with this name")
return
@management_app.get("/device/{device_name}/actions")
async def get_all_actions(device_name: str):
return {"devices": device_manager.get_all_actions(device_name)}
@management_app.get("/device/{device_name}/action/{action_name}")
async def get_action_by_name(device_name: str, action_name: str):
action = device_manager.get_action(device_name, action_name)
if action is None:
raise HTTPException(
status_code=404,
detail="Action not found on this device",
)
return {"device": action}
@management_app.post("/device/{device_name}/action/{action_name}")
async def create_action(device_name: str, action_name: str):
device = device_manager.create_action(device_name, action_name)
if device is None:
raise HTTPException(
status_code=409,
detail="Action already exists with this name on this device",
)
return {"device": device}
@management_app.delete("/device/{device_name}/action/{action_name}")
async def delete_action(device_name: str, action_name: str):
val = device_manager.delete_action(device_name, action_name)
if not val:
raise HTTPException(
status_code=404, detail="No action found with this name on this device"
)
return
async def run_servers() -> None:
public_server = uvicorn.Server(
uvicorn.Config(
public_app,
host="127.0.0.1",
port=8010,
)
)
management_server = uvicorn.Server(
uvicorn.Config(
management_app,
host="127.0.0.1",
port=8011,
)
)
async def run_servers(config_path: str | Path = "~/.config/rex/server.toml") -> None:
public, management = create_apps(config_path)
config = ServerConfigManager(config_path).config
tls = {"ssl_certfile": config.tls_certfile, "ssl_keyfile": config.tls_keyfile} if config.tls_certfile and config.tls_keyfile else {}
await asyncio.gather(
public_server.serve(),
management_server.serve(),
uvicorn.Server(uvicorn.Config(public, host=config.public_host, port=config.public_port, **tls)).serve(),
uvicorn.Server(uvicorn.Config(management, host=config.management_host, port=config.management_port, **tls)).serve(),
)
def main() -> None:
def main(config_path: str = "~/.config/rex/server.toml") -> None:
try:
asyncio.run(run_servers())
asyncio.run(run_servers(config_path))
except KeyboardInterrupt:
pass
src/rex/server/config.py
+76
View File
@@ -0,0 +1,76 @@
from __future__ import annotations
import ipaddress
import secrets
from pathlib import Path
from pydantic import BaseModel, Field, field_validator
from rex.config import read_toml, write_toml
class Key(BaseModel):
name: str
secret: str
permissions: set[str] = Field(min_length=1)
@field_validator("permissions")
@classmethod
def known_permissions(cls, value: set[str]) -> set[str]:
allowed = {"admin", "connect", "execute"}
if not value <= allowed:
raise ValueError(f"permissions must be drawn from {sorted(allowed)}")
return value
class ServerConfig(BaseModel):
public_host: str = "127.0.0.1"
public_port: int = 8010
management_host: str = "127.0.0.1"
management_port: int = 8011
tls_certfile: str | None = None
tls_keyfile: str | None = None
rate_limit_per_minute: int = Field(default=120, ge=1, le=10000)
blacklist: list[str] = Field(default_factory=list)
devices: list[str] = Field(default_factory=list)
keys: list[Key] = Field(default_factory=list)
def is_blacklisted(self, host: str | None) -> bool:
if not host:
return False
try:
address = ipaddress.ip_address(host)
return any(address in ipaddress.ip_network(item, strict=False) for item in self.blacklist)
except ValueError:
return False
class ServerConfigManager:
def __init__(self, file_path: str | Path) -> None:
self.file_path = Path(file_path).expanduser()
self.config = ServerConfig.model_validate(read_toml(self.file_path, ServerConfig().model_dump(mode="json")))
def save(self) -> None:
write_toml(self.file_path, self.config.model_dump(mode="json"))
def create_key(self, name: str, permissions: set[str]) -> Key | None:
if any(key.name == name for key in self.config.keys):
return None
key = Key(name=name, secret=secrets.token_urlsafe(32), permissions=permissions)
self.config.keys.append(key)
self.save()
return key
def delete_key(self, name: str) -> bool:
for key in self.config.keys:
if key.name == name:
self.config.keys.remove(key)
self.save()
return True
return False
def get_key(self, secret: str, permission: str) -> Key | None:
for key in self.config.keys:
if secrets.compare_digest(secret, key.secret) and (permission in key.permissions or "admin" in key.permissions):
return key
return None
src/rex/server/connection_manager.py
+16 -12
View File
@@ -1,27 +1,31 @@
from fastapi import WebSocket
from pydantic import BaseModel
class ConnectionManager:
def __init__(self) -> None:
self.connections: dict[str, WebSocket] = {}
self.connections: dict[str, tuple[WebSocket, set[str]]] = {}
async def connect(self, client_name: str, websocket: WebSocket) -> None:
await websocket.accept()
self.connections[client_name] = websocket
async def connect(self, client_name: str, websocket: WebSocket, actions: set[str]) -> None:
self.connections[client_name] = (websocket, actions)
def disconnect(self, client_name: str) -> None:
self.connections.pop(client_name, None)
def disconnect(self, client_name: str, websocket: WebSocket | None = None) -> None:
connection = self.connections.get(client_name)
if connection and (websocket is None or connection[0] is websocket):
self.connections.pop(client_name, None)
async def send(self, client_name: str, message: BaseModel) -> bool:
websocket = self.connections.get(client_name)
def actions(self, client_name: str) -> list[str]:
connection = self.connections.get(client_name)
return sorted(connection[1]) if connection else []
if websocket is None:
async def send(self, client_name: str, action_name: str) -> bool:
connection = self.connections.get(client_name)
if connection is None or action_name not in connection[1]:
return False
try:
await websocket.send_json(message.model_dump())
await connection[0].send_json({"type": "action", "name": action_name})
return True
except Exception:
self.disconnect(client_name)
self.disconnect(client_name, connection[0])
return False
src/rex/server/device_manager.py
-118
View File
@@ -1,118 +0,0 @@
from pathlib import Path
from pydantic import BaseModel
class Action(BaseModel):
name: str
class Device(BaseModel):
name: str
actions: list[Action]
def add_action(self, action_name: str) -> Action | None:
for action in self.actions:
if action.name == action_name:
return None
new_action = Action(name=action_name)
self.actions.append(new_action)
return new_action
def get_action(self, action_name: str) -> Action | None:
for action in self.actions:
if action.name == action_name:
return action
return None
def get_all_actions(self) -> list[Action]:
return self.actions
def delete_action(self, action_name: str) -> bool:
for action in self.actions:
if action.name == action_name:
self.actions.remove(action)
return True
return False
class DeviceStore(BaseModel):
devices: list[Device]
class DeviceManager:
def __init__(self, file_path: str | Path) -> None:
self.file_path = Path(file_path).expanduser()
if not self.file_path.exists():
self.device_store = DeviceStore(devices=[])
self.file_path.parent.mkdir(parents=True, exist_ok=True)
self.file_path.write_text(self.device_store.model_dump_json(indent=4))
self.device_store = DeviceStore.model_validate_json(self.file_path.read_text())
def create_device(self, device_name: str) -> Device | None:
if self.get_device(device_name) is not None:
return None
new_device = Device(name=device_name, actions=[])
self.device_store.devices.append(new_device)
self._save()
return new_device
def _save(self) -> None:
self.file_path.parent.mkdir(parents=True, exist_ok=True)
self.file_path.write_text(self.device_store.model_dump_json(indent=4))
def get_device(self, device_name: str) -> Device | None:
for device in self.device_store.devices:
if device_name == device.name:
return device
return None
def get_all_devices(self) -> list[Device]:
return self.device_store.devices
def delete_device(self, device_name: str) -> bool:
device = self.get_device(device_name)
if device is None:
return False
self.device_store.devices.remove(device)
self._save()
return True
def create_action(self, device_name: str, action_name: str) -> Action | None:
for device in self.device_store.devices:
if device.name == device_name:
action = device.add_action(action_name)
if action is not None:
self._save()
return action
return None
def get_action(self, device_name: str, action_name: str) -> Action | None:
for device in self.device_store.devices:
if device.name == device_name:
return device.get_action(action_name)
return None
def get_all_actions(self, device_name: str) -> list[Action]:
for device in self.device_store.devices:
if device.name == device_name:
return device.get_all_actions()
return []
def delete_action(self, device_name: str, action_name: str) -> bool:
for device in self.device_store.devices:
if device.name == device_name:
deleted = device.delete_action(action_name)
if deleted:
self._save()
return deleted
return False
src/rex/server/key_manager.py
-61
View File
@@ -1,61 +0,0 @@
from pathlib import Path
import secrets
from pydantic import BaseModel
class Key(BaseModel):
secret: str
name: str
class KeyStore(BaseModel):
keys: list[Key]
class KeyManager:
def __init__(self, file_path: str | Path) -> None:
self.file_path = Path(file_path).expanduser()
if not self.file_path.exists():
self.key_store = KeyStore(keys=[])
self.file_path.parent.mkdir(parents=True, exist_ok=True)
self.file_path.write_text(self.key_store.model_dump_json(indent=4))
self.key_store = KeyStore.model_validate_json(self.file_path.read_text())
def create_key(self, key_name: str) -> Key | None:
if self.get_key(key_name) is not None:
return None
new_key = Key(secret=secrets.token_urlsafe(32), name=key_name)
self.key_store.keys.append(new_key)
self._save()
return new_key
def _save(self) -> None:
self.file_path.parent.mkdir(parents=True, exist_ok=True)
self.file_path.write_text(self.key_store.model_dump_json(indent=4))
def get_key(self, key_name: str) -> Key | None:
for key in self.key_store.keys:
if key_name == key.name:
return key
return None
def get_all_keys(self) -> list[Key]:
return self.key_store.keys
def delete_key(self, key_name: str) -> bool:
key = self.get_key(key_name)
if key is None:
return False
self.key_store.keys.remove(key)
self._save()
return True
def compare_key(self, key: str) -> bool:
for k in self.key_store.keys:
if secrets.compare_digest(key, k.secret):
return True
return False
tests/test_client.py
-89
View File
@@ -1,89 +0,0 @@
import asyncio
import signal
import unittest
from unittest.mock import AsyncMock, patch
from rex.client import action_manager, receive_actions, run_command
from rex.client.action_manager import ClientAction
class FakeWebSocket:
def __init__(self, messages: list[str]) -> None:
self.messages = iter(messages)
def __aiter__(self):
return self
async def __anext__(self) -> str:
await asyncio.sleep(0)
try:
return next(self.messages)
except StopIteration:
raise StopAsyncIteration
class ClientTests(unittest.IsolatedAsyncioTestCase):
async def test_run_command_waits_for_process(self) -> None:
process = AsyncMock()
process.wait.return_value = 0
with patch("rex.client.asyncio.create_subprocess_shell", return_value=process):
await run_command("test command")
process.wait.assert_awaited_once()
async def test_cancelling_command_terminates_process_group(self) -> None:
process = AsyncMock()
process.pid = 123
process.returncode = None
wait_started = asyncio.Event()
wait_calls = 0
async def wait() -> int:
nonlocal wait_calls
wait_calls += 1
if wait_calls == 1:
wait_started.set()
await asyncio.Event().wait()
return 0
process.wait.side_effect = wait
with (
patch("rex.client.asyncio.create_subprocess_shell", return_value=process),
patch("rex.client.os.name", "posix"),
patch("rex.client.os.killpg") as killpg,
):
task = asyncio.create_task(run_command("test command"))
await wait_started.wait()
task.cancel()
await asyncio.gather(task, return_exceptions=True)
killpg.assert_called_once_with(123, signal.SIGTERM)
self.assertEqual(process.wait.await_count, 2)
async def test_disconnect_cancels_running_commands(self) -> None:
action = ClientAction(name="wake", command="test command")
command_started = asyncio.Event()
command_cancelled = asyncio.Event()
async def running_command(_: str) -> None:
command_started.set()
try:
await asyncio.Event().wait()
finally:
command_cancelled.set()
websocket = FakeWebSocket(['{"name":"wake"}'])
with (
patch.object(action_manager, "get_action", return_value=action),
patch("rex.client.run_command", side_effect=running_command),
):
await receive_actions(websocket) # type: ignore[arg-type]
self.assertTrue(command_started.is_set())
self.assertTrue(command_cancelled.is_set())
if __name__ == "__main__":
unittest.main()
tests/test_config.py
+46
View File
@@ -0,0 +1,46 @@
import tempfile
import unittest
from pathlib import Path
from rex.client.config import ClientAction, ClientConfig, ClientConfigManager
from rex.server.config import ServerConfigManager
from rex.server.connection_manager import ConnectionManager
class ConfigTests(unittest.TestCase):
def test_server_config_is_toml_and_persists_scoped_keys(self) -> None:
with tempfile.TemporaryDirectory() as directory:
path = Path(directory) / "server.toml"
manager = ServerConfigManager(path)
key = manager.create_key("device", {"connect"})
self.assertIsNotNone(key)
self.assertIn("[[keys]]", path.read_text())
self.assertIsNotNone(ServerConfigManager(path).get_key(key.secret, "connect"))
self.assertIsNone(ServerConfigManager(path).get_key(key.secret, "execute"))
def test_client_url_omits_optional_port(self) -> None:
self.assertEqual(ClientConfig().websocket_url(), "wss://127.0.0.1/ws/default-device")
self.assertEqual(ClientConfig(endpoint="rex.test", scheme="ws", port=8080).websocket_url(), "ws://rex.test:8080/ws/default-device")
def test_client_actions_are_local_argv_only(self) -> None:
action = ClientAction(name="lock", argv=["/bin/echo", "locked"])
self.assertEqual(action.argv[0], "/bin/echo")
with self.assertRaises(ValueError):
ClientAction(name="bad name", argv=["/bin/echo"])
class ConnectionManagerTests(unittest.IsolatedAsyncioTestCase):
async def test_only_registered_actions_are_sent(self) -> None:
class Socket:
def __init__(self) -> None:
self.messages: list[dict[str, str]] = []
async def accept(self) -> None: pass
async def send_json(self, message: dict[str, str]) -> None: self.messages.append(message)
socket = Socket()
manager = ConnectionManager()
await manager.connect("desk", socket, {"lock"}) # type: ignore[arg-type]
self.assertFalse(await manager.send("desk", "shell"))
self.assertTrue(await manager.send("desk", "lock"))
self.assertEqual(socket.messages, [{"type": "action", "name": "lock"}])
tests/test_management.py
-125
View File
@@ -1,125 +0,0 @@
import tempfile
import unittest
from pathlib import Path
from unittest.mock import patch
import httpx
from textual.widgets import DataTable, TabbedContent, Tree
from rex.management.app import RexManager
from rex.management.client import ManagementClient, ManagementError
from rex.server.device_manager import DeviceManager
class ManagementClientTests(unittest.IsolatedAsyncioTestCase):
async def test_reads_keys_and_encodes_names(self) -> None:
requests: list[httpx.Request] = []
def handler(request: httpx.Request) -> httpx.Response:
requests.append(request)
if request.method == "GET":
return httpx.Response(
200, json={"keys": [{"name": "home", "secret": "abc"}]}
)
return httpx.Response(200, json={"key": {}})
client = ManagementClient(
"http://rex.test", transport=httpx.MockTransport(handler)
)
self.addAsyncCleanup(client.close)
self.assertEqual((await client.keys())[0]["name"], "home")
await client.create_key("living room")
self.assertEqual(requests[1].url.raw_path, b"/key/living%20room")
async def test_reports_api_error_detail(self) -> None:
def handler(_: httpx.Request) -> httpx.Response:
return httpx.Response(409, json={"detail": "already exists"})
client = ManagementClient(
"http://rex.test", transport=httpx.MockTransport(handler)
)
self.addAsyncCleanup(client.close)
with self.assertRaisesRegex(ManagementError, "already exists"):
await client.create_device("desk")
class DeviceManagerTests(unittest.TestCase):
def test_action_changes_are_persisted(self) -> None:
with tempfile.TemporaryDirectory() as directory:
path = Path(directory) / "devices.json"
manager = DeviceManager(path)
manager.create_device("desk")
manager.create_action("desk", "wake")
reloaded = DeviceManager(path)
self.assertIsNotNone(reloaded.get_action("desk", "wake"))
reloaded.delete_action("desk", "wake")
self.assertIsNone(DeviceManager(path).get_action("desk", "wake"))
class FakeManagementClient:
def __init__(self) -> None:
self.created_keys: list[str] = []
async def close(self) -> None:
pass
async def keys(self) -> list[dict[str, str]]:
return [{"name": "home", "secret": "abc"}]
async def devices(self) -> list[dict[str, object]]:
return [{"name": "desk", "actions": [{"name": "wake"}]}]
async def create_key(self, name: str) -> None:
self.created_keys.append(name)
class ManagerAppTests(unittest.IsolatedAsyncioTestCase):
async def test_refresh_populates_keys_and_devices(self) -> None:
app = RexManager("http://rex.test")
app.client = FakeManagementClient() # type: ignore[assignment]
async with app.run_test() as pilot:
await pilot.pause()
keys = app.query_one("#keys", DataTable)
devices = app.query_one("#devices", Tree)
self.assertEqual(keys.get_row_at(0), ["home", "abc"])
self.assertEqual(devices.root.children[0].label.plain, "desk")
self.assertEqual(devices.root.children[0].children[0].label.plain, "wake")
async def test_vim_navigation_and_create_prompt(self) -> None:
app = RexManager("http://rex.test")
client = FakeManagementClient()
app.client = client # type: ignore[assignment]
async with app.run_test() as pilot:
await pilot.press("n")
await pilot.press(*"office", "enter")
await pilot.pause()
self.assertEqual(client.created_keys, ["office"])
await pilot.press("shift+l", "j")
self.assertEqual(app.query_one(TabbedContent).active, "devices-tab")
self.assertEqual(
app.query_one("#devices", Tree).cursor_node.data,
("device", "desk"),
)
async def test_moves_between_key_columns_and_yanks_value(self) -> None:
app = RexManager("http://rex.test")
app.client = FakeManagementClient() # type: ignore[assignment]
async with app.run_test() as pilot:
with patch.object(app, "copy_to_clipboard") as copy:
await pilot.press("l", "y")
self.assertEqual(app.query_one("#keys", DataTable).cursor_column, 1)
copy.assert_called_once_with("abc")
if __name__ == "__main__":
unittest.main()