137 lines
4.8 KiB
Python
137 lines
4.8 KiB
Python
"""Secrets and package bundling for Agent distribution.
|
|
|
|
Two independent secrets guard the distribution flow, both random(256 bits),
|
|
single-use and short-TTL, both stored only as SHA-256 hashes:
|
|
|
|
- the *download* token: embedded in the shareable link an operator hands to
|
|
whoever installs the Agent (or clicked directly, authenticated, from the
|
|
panel);
|
|
- the *enrollment* token: embedded inside the downloaded bundle as
|
|
`enrollment.json`, consumed exactly once by `POST /v1/agents/enroll`.
|
|
|
|
Neither token is ever the Agent's own permanent Bearer token (issued only
|
|
after a successful enrollment) and neither is ever logged or audited in
|
|
plaintext.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import hmac
|
|
import io
|
|
import json
|
|
import re
|
|
import secrets
|
|
import zipfile
|
|
from dataclasses import dataclass
|
|
|
|
|
|
class EnrollmentRejected(Exception):
|
|
def __init__(self, reason: str):
|
|
super().__init__(reason)
|
|
self.reason = reason
|
|
|
|
|
|
def new_token() -> str:
|
|
"""256 bits of randomness, URL-safe — usable directly in a link or a JSON field."""
|
|
return secrets.token_urlsafe(32)
|
|
|
|
|
|
# ---- installation code (XXXX-XXXX-XXXX-XXXX) ----
|
|
#
|
|
# A 16-character Crockford Base32 code = 16 * 5 = 80 bits of entropy. The
|
|
# alphabet deliberately excludes I/L/O/U (ambiguous when retyped by a human);
|
|
# the same alphabet is used so a code cannot contain those letters.
|
|
|
|
INSTALL_CODE_ALPHABET = "0123456789ABCDEFGHJKMNPQRSTVWXYZ"
|
|
_INSTALL_CODE_GROUPS = 4
|
|
_INSTALL_CODE_PER_GROUP = 4
|
|
_INSTALL_CODE_LENGTH = _INSTALL_CODE_GROUPS * _INSTALL_CODE_PER_GROUP
|
|
# Crockford transliteration for input normalization: O -> 0, I/L -> 1.
|
|
_INSTALL_CODE_TRANSLIT = str.maketrans({"O": "0", "I": "1", "L": "1"})
|
|
_INSTALL_CODE_RE = re.compile(
|
|
r"^[" + re.escape(INSTALL_CODE_ALPHABET) + r"]{" + str(_INSTALL_CODE_LENGTH) + r"}$"
|
|
)
|
|
|
|
|
|
def generate_install_code() -> str:
|
|
"""Return a fresh installation code formatted XXXX-XXXX-XXXX-XXXX."""
|
|
chars = "".join(secrets.choice(INSTALL_CODE_ALPHABET) for _ in range(_INSTALL_CODE_LENGTH))
|
|
return "-".join(
|
|
chars[i : i + _INSTALL_CODE_PER_GROUP] for i in range(0, _INSTALL_CODE_LENGTH, _INSTALL_CODE_PER_GROUP)
|
|
)
|
|
|
|
|
|
def normalize_install_code(text: str | None) -> str:
|
|
"""Upper-case, drop hyphen/space separators and map Crockford lookalikes,
|
|
so a human-typed code normalizes to the same digest regardless of how it
|
|
was entered. Raises ValueError if the result isn't a valid code."""
|
|
if not text:
|
|
raise ValueError("empty install code")
|
|
clean = "".join(ch for ch in text.upper().translate(_INSTALL_CODE_TRANSLIT) if ch.isalnum())
|
|
if not _INSTALL_CODE_RE.match(clean):
|
|
raise ValueError(f"invalid install code: {text!r}")
|
|
return clean
|
|
|
|
|
|
def install_code_digest(code: str) -> str:
|
|
"""Hash a normalized install code. Only the hash is ever persisted; the
|
|
plaintext code exists only in the operator's clipboard/UI and the EXE's
|
|
memory for the brief instant it is sent to the Hub."""
|
|
return token_digest(normalize_install_code(code))
|
|
|
|
|
|
def valid_install_code(text: str | None) -> bool:
|
|
try:
|
|
normalize_install_code(text)
|
|
return True
|
|
except ValueError:
|
|
return False
|
|
|
|
|
|
def token_digest(token: str) -> str:
|
|
return hashlib.sha256(token.encode()).hexdigest()
|
|
|
|
|
|
def token_matches(token: str | None, stored_hash: str | None) -> bool:
|
|
if not token or not stored_hash:
|
|
return False
|
|
return hmac.compare_digest(token_digest(token), stored_hash)
|
|
|
|
|
|
ENROLLMENT_UUID_RE = re.compile(r"^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$")
|
|
|
|
|
|
def valid_uuid(value: str | None) -> bool:
|
|
return bool(value) and bool(ENROLLMENT_UUID_RE.match(value))
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class BundlePlan:
|
|
hub_url: str
|
|
enrollment_uuid: str
|
|
enrollment_token: str
|
|
enroll_path: str = "/v1/agents/enroll"
|
|
|
|
|
|
def build_bundle(release_zip_bytes: bytes, plan: BundlePlan) -> bytes:
|
|
"""Return a new zip = the release zip's entries plus one
|
|
`<package-root>/enrollment.json` holding only temporary, single-use data.
|
|
Never a permanent credential — that is issued only once enrollment succeeds."""
|
|
with zipfile.ZipFile(io.BytesIO(release_zip_bytes)) as source:
|
|
names = source.namelist()
|
|
if not names:
|
|
raise EnrollmentRejected("EMPTY_RELEASE_PACKAGE")
|
|
root = names[0].split("/", 1)[0]
|
|
out = io.BytesIO()
|
|
with zipfile.ZipFile(out, "w", zipfile.ZIP_DEFLATED) as dest:
|
|
for name in names:
|
|
dest.writestr(name, source.read(name))
|
|
payload = {
|
|
"installation_id": plan.enrollment_uuid,
|
|
"enrollment_token": plan.enrollment_token,
|
|
"hub_url": plan.hub_url,
|
|
"enroll_path": plan.enroll_path,
|
|
}
|
|
dest.writestr(f"{root}/enrollment.json", json.dumps(payload, indent=1))
|
|
return out.getvalue()
|