Support multi-use installer codes and reuse

tests/test_enrollment_install_codes.py

@@ -0,0 +1,213 @@
+import base64
+import os
+import pathlib
+import sqlite3
+import sys
+import uuid
+
+import pytest
+try:  # pragma: no cover - optional dependency
+    from cryptography.hazmat.primitives import serialization
+    from cryptography.hazmat.primitives.asymmetric import ed25519
+    _CRYPTO_IMPORT_ERROR: Exception | None = None
+except Exception as exc:  # pragma: no cover - dependency unavailable
+    serialization = None  # type: ignore
+    ed25519 = None  # type: ignore
+    _CRYPTO_IMPORT_ERROR = exc
+try:  # pragma: no cover - optional dependency
+    from flask import Flask
+    _FLASK_IMPORT_ERROR: Exception | None = None
+except Exception as exc:  # pragma: no cover - dependency unavailable
+    Flask = None  # type: ignore
+    _FLASK_IMPORT_ERROR = exc
+
+ROOT = pathlib.Path(__file__).resolve().parents[1]
+if str(ROOT) not in sys.path:
+    sys.path.insert(0, str(ROOT))
+
+from Data.Server.Modules import db_migrations
+from Data.Server.Modules.auth.rate_limit import SlidingWindowRateLimiter
+from Data.Server.Modules.enrollment.nonce_store import NonceCache
+
+if Flask is not None:  # pragma: no cover - dependency unavailable
+    from Data.Server.Modules.enrollment import routes as enrollment_routes
+else:  # pragma: no cover - dependency unavailable
+    enrollment_routes = None  # type: ignore
+
+
+class _DummyJWTService:
+    def issue_access_token(self, guid: str, fingerprint: str, token_version: int, expires_in: int = 900, extra_claims=None):
+        return f"token-{guid}"
+
+
+class _DummySigner:
+    def public_base64_spki(self) -> str:
+        return ""
+
+
+def _make_app(db_path: str, tls_path: str):
+    if Flask is None or enrollment_routes is None:  # pragma: no cover - dependency unavailable
+        pytest.skip(f"flask unavailable: {_FLASK_IMPORT_ERROR}")
+    app = Flask(__name__)
+
+    def _factory():
+        conn = sqlite3.connect(db_path)
+        conn.row_factory = sqlite3.Row
+        return conn
+
+    enrollment_routes.register(
+        app,
+        db_conn_factory=_factory,
+        log=lambda channel, message: None,
+        jwt_service=_DummyJWTService(),
+        tls_bundle_path=tls_path,
+        ip_rate_limiter=SlidingWindowRateLimiter(),
+        fp_rate_limiter=SlidingWindowRateLimiter(),
+        nonce_cache=NonceCache(ttl_seconds=30.0),
+        script_signer=_DummySigner(),
+    )
+    return app, _factory
+
+
+def _create_install_code(conn: sqlite3.Connection, code: str, *, max_uses: int = 2):
+    cur = conn.cursor()
+    record_id = str(uuid.uuid4())
+    cur.execute(
+        """
+        INSERT INTO enrollment_install_codes (
+            id, code, expires_at, created_by_user_id, max_uses, use_count
+        )
+        VALUES (?, ?, datetime('now', '+6 hours'), 'test-user', ?, 0)
+        """,
+        (record_id, code, max_uses),
+    )
+    conn.commit()
+    return record_id
+
+
+def _perform_enrollment_cycle(app, factory, code: str, private_key):
+    client = app.test_client()
+    public_der = private_key.public_key().public_bytes(
+        serialization.Encoding.DER,
+        serialization.PublicFormat.SubjectPublicKeyInfo,
+    )
+    public_b64 = base64.b64encode(public_der).decode("ascii")
+    client_nonce = os.urandom(32)
+    payload = {
+        "hostname": "unit-test-host",
+        "enrollment_code": code,
+        "agent_pubkey": public_b64,
+        "client_nonce": base64.b64encode(client_nonce).decode("ascii"),
+    }
+    request_resp = client.post("/api/agent/enroll/request", json=payload)
+    assert request_resp.status_code == 200
+    request_data = request_resp.get_json()
+    approval_reference = request_data["approval_reference"]
+
+    with factory() as conn:
+        cur = conn.cursor()
+        cur.execute(
+            """
+            UPDATE device_approvals
+               SET status = 'approved',
+                   approved_by_user_id = 'tester'
+             WHERE approval_reference = ?
+            """,
+            (approval_reference,),
+        )
+        cur.execute(
+            """
+            SELECT server_nonce, client_nonce
+              FROM device_approvals
+             WHERE approval_reference = ?
+            """,
+            (approval_reference,),
+        )
+        row = cur.fetchone()
+        assert row is not None
+        server_nonce_b64 = row["server_nonce"]
+
+    server_nonce = base64.b64decode(server_nonce_b64)
+    proof_message = server_nonce + approval_reference.encode("utf-8") + client_nonce
+    proof_sig = private_key.sign(proof_message)
+    poll_payload = {
+        "approval_reference": approval_reference,
+        "client_nonce": base64.b64encode(client_nonce).decode("ascii"),
+        "proof_sig": base64.b64encode(proof_sig).decode("ascii"),
+    }
+    poll_resp = client.post("/api/agent/enroll/poll", json=poll_payload)
+    assert poll_resp.status_code == 200
+    return poll_resp.get_json()
+
+
+@pytest.mark.parametrize("max_uses", [2])
+@pytest.mark.skipif(ed25519 is None, reason=f"cryptography unavailable: {_CRYPTO_IMPORT_ERROR}")
+@pytest.mark.skipif(Flask is None, reason=f"flask unavailable: {_FLASK_IMPORT_ERROR}")
+def test_install_code_allows_multiple_and_reuse(tmp_path, max_uses):
+    db_path = tmp_path / "test.db"
+    conn = sqlite3.connect(db_path)
+    db_migrations.apply_all(conn)
+    _create_install_code(conn, "TEST-CODE-1234", max_uses=max_uses)
+    conn.close()
+
+    tls_path = tmp_path / "tls.pem"
+    tls_path.write_text("TEST CERT")
+
+    app, factory = _make_app(str(db_path), str(tls_path))
+
+    private_key = ed25519.Ed25519PrivateKey.generate()
+
+    # First enrollment consumes one use but keeps the code active.
+    first = _perform_enrollment_cycle(app, factory, "TEST-CODE-1234", private_key)
+    assert first["status"] == "approved"
+
+    with factory() as conn:
+        cur = conn.cursor()
+        cur.execute(
+            "SELECT use_count, max_uses, used_at, last_used_at FROM enrollment_install_codes WHERE code = ?",
+            ("TEST-CODE-1234",),
+        )
+        row = cur.fetchone()
+        assert row is not None
+        assert row["use_count"] == 1
+        assert row["max_uses"] == max_uses
+        assert row["used_at"] is None
+        assert row["last_used_at"] is not None
+
+    # Second enrollment hits the configured max uses.
+    second = _perform_enrollment_cycle(app, factory, "TEST-CODE-1234", private_key)
+    assert second["status"] == "approved"
+
+    with factory() as conn:
+        cur = conn.cursor()
+        cur.execute(
+            "SELECT use_count, used_at, last_used_at, used_by_guid FROM enrollment_install_codes WHERE code = ?",
+            ("TEST-CODE-1234",),
+        )
+        row = cur.fetchone()
+        assert row is not None
+        assert row["use_count"] == max_uses
+        assert row["used_at"] is not None
+        assert row["last_used_at"] is not None
+        consumed_guid = row["used_by_guid"]
+        assert consumed_guid
+
+    # Additional enrollments from the same identity reuse the stored GUID even after consumption.
+    third = _perform_enrollment_cycle(app, factory, "TEST-CODE-1234", private_key)
+    assert third["status"] == "approved"
+
+    with factory() as conn:
+        cur = conn.cursor()
+        cur.execute(
+            "SELECT use_count, used_at, last_used_at, used_by_guid FROM enrollment_install_codes WHERE code = ?",
+            ("TEST-CODE-1234",),
+        )
+        row = cur.fetchone()
+        assert row is not None
+        assert row["use_count"] == max_uses + 1
+        assert row["used_by_guid"] == consumed_guid
+        assert row["used_at"] is not None
+        assert row["last_used_at"] is not None
+
+        cur.execute("SELECT COUNT(*) FROM devices WHERE guid = ?", (consumed_guid,))
+        assert cur.fetchone()[0] == 1