feat: secure agent auth and heartbeat endpoints

This commit is contained in:
2025-10-17 17:15:02 -06:00
parent 78a5d3d7f9
commit a72bff5e8e
6 changed files with 420 additions and 1 deletions

View File

@@ -0,0 +1 @@

View File

@@ -0,0 +1,125 @@
from __future__ import annotations
import hashlib
import sqlite3
from datetime import datetime, timezone
from typing import Callable
from flask import Blueprint, jsonify, request
from Modules.auth.dpop import DPoPValidator, DPoPVerificationError, DPoPReplayError
def register(
app,
*,
db_conn_factory: Callable[[], sqlite3.Connection],
jwt_service,
dpop_validator: DPoPValidator,
) -> None:
blueprint = Blueprint("tokens", __name__)
def _hash_token(token: str) -> str:
return hashlib.sha256(token.encode("utf-8")).hexdigest()
def _iso_now() -> str:
return datetime.now(tz=timezone.utc).isoformat()
def _parse_iso(ts: str) -> datetime:
return datetime.fromisoformat(ts)
@blueprint.route("/api/agent/token/refresh", methods=["POST"])
def refresh():
payload = request.get_json(force=True, silent=True) or {}
guid = str(payload.get("guid") or "").strip()
refresh_token = str(payload.get("refresh_token") or "").strip()
if not guid or not refresh_token:
return jsonify({"error": "invalid_request"}), 400
conn = db_conn_factory()
try:
cur = conn.cursor()
cur.execute(
"""
SELECT id, guid, token_hash, dpop_jkt, created_at, expires_at, revoked_at
FROM refresh_tokens
WHERE guid = ?
AND token_hash = ?
""",
(guid, _hash_token(refresh_token)),
)
row = cur.fetchone()
if not row:
return jsonify({"error": "invalid_refresh_token"}), 401
record_id, row_guid, _token_hash, stored_jkt, created_at, expires_at, revoked_at = row
if row_guid != guid:
return jsonify({"error": "invalid_refresh_token"}), 401
if revoked_at:
return jsonify({"error": "refresh_token_revoked"}), 401
if expires_at:
try:
if _parse_iso(expires_at) <= datetime.now(tz=timezone.utc):
return jsonify({"error": "refresh_token_expired"}), 401
except Exception:
pass
cur.execute(
"""
SELECT guid, ssl_key_fingerprint, token_version, status
FROM devices
WHERE guid = ?
""",
(guid,),
)
device_row = cur.fetchone()
if not device_row:
return jsonify({"error": "device_not_found"}), 404
device_guid, fingerprint, token_version, status = device_row
status_norm = (status or "active").strip().lower()
if status_norm in {"revoked", "decommissioned"}:
return jsonify({"error": "device_revoked"}), 403
dpop_proof = request.headers.get("DPoP")
jkt = stored_jkt or ""
if dpop_proof:
try:
jkt = dpop_validator.verify(request.method, request.url, dpop_proof, access_token=None)
except DPoPReplayError:
return jsonify({"error": "dpop_replayed"}), 400
except DPoPVerificationError:
return jsonify({"error": "dpop_invalid"}), 400
elif stored_jkt:
return jsonify({"error": "dpop_required"}), 400
new_access_token = jwt_service.issue_access_token(
guid,
fingerprint or "",
token_version or 1,
)
cur.execute(
"""
UPDATE refresh_tokens
SET last_used_at = ?,
dpop_jkt = COALESCE(NULLIF(?, ''), dpop_jkt)
WHERE id = ?
""",
(_iso_now(), jkt, record_id),
)
conn.commit()
finally:
conn.close()
return jsonify(
{
"access_token": new_access_token,
"expires_in": 900,
"token_type": "Bearer",
}
)
app.register_blueprint(blueprint)