Spaces:
Running
Running
File size: 4,081 Bytes
2f1e7ff | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 | """Explicit credentials in the system PostgreSQL registry; hashes only at rest."""
from contextlib import contextmanager
from datetime import datetime, timezone
from uuid import uuid4
import threading
import secrets
import psycopg2.extras
from src.services.auth import AuthError, token_digest
PERMISSIONS = frozenset({"memory:read", "memory:write", "brain:read", "brain:write"})
_DDL = """
CREATE TABLE IF NOT EXISTS access_tokens (
id TEXT PRIMARY KEY,
label TEXT NOT NULL,
brain_id TEXT NOT NULL,
brain_resource_id TEXT NOT NULL,
token_prefix TEXT NOT NULL,
token_hash TEXT UNIQUE NOT NULL,
permissions TEXT[] NOT NULL,
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
last_used_at TIMESTAMPTZ,
expires_at TIMESTAMPTZ,
revoked_at TIMESTAMPTZ,
legacy BOOLEAN NOT NULL DEFAULT FALSE,
created_by TEXT NOT NULL
);
CREATE INDEX IF NOT EXISTS access_tokens_brain ON access_tokens (brain_id);
"""
def public_metadata(row):
return {key: value for key, value in row.items() if key not in {"token_hash", "brain_resource_id"}}
class PostgreSQLTokenRepository:
def __init__(self):
self._ready = False
self._lock = threading.Lock()
@contextmanager
def connection(self):
from src.lib.postgresql._provisioning import borrow, get_system_pool
with borrow(get_system_pool()) as connection:
yield connection
def initialize(self):
if self._ready:
return
with self._lock:
if self._ready:
return
with self.connection() as conn:
with conn.cursor() as cur:
cur.execute(_DDL)
conn.commit()
self._ready = True
def _query(self, sql, args=(), many=False, write=False):
self.initialize()
with self.connection() as conn:
with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
cur.execute(sql, args)
rows = cur.fetchall() if many else cur.fetchone()
if write:
conn.commit()
return rows
def find_by_hash(self, digest):
return self._query("SELECT * FROM access_tokens WHERE token_hash=%s", (digest,))
def get(self, token_id):
return self._query("SELECT * FROM access_tokens WHERE id=%s", (token_id,))
def list(self, brain_id=None):
sql = "SELECT * FROM access_tokens"
args = ()
if brain_id is not None:
sql += " WHERE brain_id=%s"
args = (brain_id,)
return [public_metadata(row) for row in self._query(sql + " ORDER BY created_at DESC", args, many=True)]
def create(self, *, brain, label, permissions, created_by, expires_at=None):
if not permissions or not set(permissions) <= PERMISSIONS:
raise ValueError("Unsupported token permissions")
if expires_at and expires_at <= datetime.now(timezone.utc):
raise ValueError("Token expiry must be in the future")
plaintext = "bpat_" + secrets.token_urlsafe(32)
row = self._query("""
INSERT INTO access_tokens
(id,label,brain_id,brain_resource_id,token_prefix,token_hash,permissions,created_by,expires_at)
VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s) RETURNING *
""", (str(uuid4()), label, brain.name_key, brain.id, plaintext[:12],
token_digest(plaintext), sorted(set(permissions)), created_by, expires_at), write=True)
return {**public_metadata(row), "token": plaintext}
def revoke(self, token_id):
row = self._query("UPDATE access_tokens SET revoked_at=COALESCE(revoked_at,NOW()) WHERE id=%s RETURNING *",
(token_id,), write=True)
if not row:
raise AuthError(404, "TOKEN_NOT_FOUND")
return public_metadata(row)
def touch(self, token_id):
self._query("UPDATE access_tokens SET last_used_at=NOW() WHERE id=%s RETURNING id", (token_id,), write=True)
token_repository = PostgreSQLTokenRepository()
|