118 lines
3.4 KiB
Python
118 lines
3.4 KiB
Python
|
"""Box (WireGuard peer) lifecycle: registration, rotation, deregistration."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import sqlite3
|
||
|
|
import uuid
|
||
|
|
from datetime import UTC, datetime
|
||
|
|
|
||
|
|
from control_plane import routes, tokens
|
||
|
|
from control_plane.db import Database
|
||
|
|
from control_plane.passwd import hash_password, verify_password
|
||
|
|
from control_plane.wgeasy import WgEasyClient
|
||
|
|
|
||
|
|
|
||
|
|
class BoxLimitExceeded(Exception):
|
||
|
|
pass
|
||
|
|
|
||
|
|
|
||
|
|
def _box_count(db: Database, account_id: str) -> int:
|
||
|
|
row = db.query_one("SELECT COUNT(*) AS n FROM boxes WHERE account_id = ?", (account_id,))
|
||
|
|
return row["n"]
|
||
|
|
|
||
|
|
|
||
|
|
def register_box(
|
||
|
|
db: Database,
|
||
|
|
wgeasy: WgEasyClient,
|
||
|
|
account_id: str,
|
||
|
|
box_name: str,
|
||
|
|
box_limit: int,
|
||
|
|
) -> dict:
|
||
|
|
"""Create a wg-easy peer for a new box and persist it.
|
||
|
|
|
||
|
|
Returns everything the box needs to bring its tunnel up itself: a box
|
||
|
|
token for future gateway API calls, and the full WireGuard interface
|
||
|
|
config (including the private key — see wgeasy.py's module docstring
|
||
|
|
for why the gateway ends up handling that at all).
|
||
|
|
"""
|
||
|
|
if _box_count(db, account_id) >= box_limit:
|
||
|
|
raise BoxLimitExceeded(account_id)
|
||
|
|
|
||
|
|
peer = wgeasy.create_client(box_name)
|
||
|
|
|
||
|
|
box_id = uuid.uuid4().hex
|
||
|
|
box_token, box_secret = tokens.issue(box_id)
|
||
|
|
now = datetime.now(UTC).isoformat()
|
||
|
|
db.execute(
|
||
|
|
"""
|
||
|
|
INSERT INTO boxes
|
||
|
|
(id, account_id, name, wg_public_key, wg_peer_id, wg_allowed_ip,
|
||
|
|
box_token_hash, registered_at, last_seen_at)
|
||
|
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?, NULL)
|
||
|
|
""",
|
||
|
|
(
|
||
|
|
box_id,
|
||
|
|
account_id,
|
||
|
|
box_name,
|
||
|
|
peer["public_key"],
|
||
|
|
peer["id"],
|
||
|
|
peer["address"],
|
||
|
|
hash_password(box_secret),
|
||
|
|
now,
|
||
|
|
),
|
||
|
|
)
|
||
|
|
return {
|
||
|
|
"box_id": box_id,
|
||
|
|
"box_token": box_token,
|
||
|
|
"wg": {
|
||
|
|
"private_key": peer["private_key"],
|
||
|
|
"public_key": peer["public_key"],
|
||
|
|
"address": peer["address"],
|
||
|
|
"server_public_key": peer["server_public_key"],
|
||
|
|
"endpoint": peer["endpoint"],
|
||
|
|
"allowed_ips": peer["allowed_ips"],
|
||
|
|
},
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def authenticate_box_token(db: Database, token: str) -> sqlite3.Row | None:
|
||
|
|
parsed = tokens.split(token)
|
||
|
|
if parsed is None:
|
||
|
|
return None
|
||
|
|
box_id, secret = parsed
|
||
|
|
row = db.query_one("SELECT * FROM boxes WHERE id = ?", (box_id,))
|
||
|
|
if row is None:
|
||
|
|
return None
|
||
|
|
if not verify_password(secret, row["box_token_hash"]):
|
||
|
|
return None
|
||
|
|
return row
|
||
|
|
|
||
|
|
|
||
|
|
def touch_last_seen(db: Database, box_id: str) -> None:
|
||
|
|
db.execute(
|
||
|
|
"UPDATE boxes SET last_seen_at = ? WHERE id = ?",
|
||
|
|
(datetime.now(UTC).isoformat(), box_id),
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def rotate_box_token(db: Database, box_id: str) -> str | None:
|
||
|
|
row = db.query_one("SELECT id FROM boxes WHERE id = ?", (box_id,))
|
||
|
|
if row is None:
|
||
|
|
return None
|
||
|
|
box_token, box_secret = tokens.issue(box_id)
|
||
|
|
db.execute(
|
||
|
|
"UPDATE boxes SET box_token_hash = ? WHERE id = ?",
|
||
|
|
(hash_password(box_secret), box_id),
|
||
|
|
)
|
||
|
|
return box_token
|
||
|
|
|
||
|
|
|
||
|
|
def deregister_box(db: Database, wgeasy: WgEasyClient, box_id: str) -> bool:
|
||
|
|
row = db.query_one("SELECT wg_peer_id FROM boxes WHERE id = ?", (box_id,))
|
||
|
|
if row is None:
|
||
|
|
return False
|
||
|
|
routes.unpublish_all_for_box(db, box_id)
|
||
|
|
wgeasy.delete_client(row["wg_peer_id"])
|
||
|
|
db.execute("DELETE FROM boxes WHERE id = ?", (box_id,))
|
||
|
|
return True
|