feat: add account service with API key generation and credit ledger

src/proxy_pool/accounts/service.py

@@ -0,0 +1,137 @@
+from __future__ import annotations
+
+import hashlib
+import secrets
+from uuid import UUID
+
+from sqlalchemy import func, select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from proxy_pool.accounts.models import ApiKey, CreditLedger, CreditTxType, User
+from proxy_pool.config import AccountSettings
+
+
+def generate_api_key(prefix: str) -> tuple[str, str, str]:
+    """Generate an API key. Returns (raw_key, key_hash, prefix_fragment)."""
+    random_part = secrets.token_urlsafe(36)
+    raw_key = f"{prefix}{random_part}"
+    key_hash = hashlib.sha256(raw_key.encode()).hexdigest()
+    prefix_fragment = raw_key[:8]
+    return raw_key, key_hash, prefix_fragment
+
+
+def hash_api_key(raw_key: str) -> str:
+    return hashlib.sha256(raw_key.encode()).hexdigest()
+
+
+async def create_user(
+    db: AsyncSession,
+    email: str,
+    display_name: str | None,
+    settings: AccountSettings,
+) -> tuple[User, str]:
+    """Create a user with an initial API key and starting credits.
+
+    Returns (user, raw_api_key). The raw key is only available at creation time.
+    """
+    user = User(email=email, display_name=display_name)
+    db.add(user)
+    await db.flush()
+
+    raw_key, key_hash, prefix = generate_api_key(settings.api_key_prefix)
+    api_key = ApiKey(
+        user_id=user.id,
+        key_hash=key_hash,
+        prefix=prefix,
+        label="default",
+    )
+    db.add(api_key)
+
+    if settings.default_credits > 0:
+        ledger_entry = CreditLedger(
+            user_id=user.id,
+            amount=settings.default_credits,
+            tx_type=CreditTxType.PURCHASE,
+            description="Initial credits",
+        )
+        db.add(ledger_entry)
+
+    await db.commit()
+    await db.refresh(user)
+
+    return user, raw_key
+
+
+async def verify_api_key(
+    db: AsyncSession,
+    raw_key: str,
+) -> User | None:
+    """Verify an API key and return the associated user, or None."""
+    prefix = raw_key[:8]
+    key_hash = hash_api_key(raw_key)
+
+    stmt = (
+        select(ApiKey)
+        .where(
+            ApiKey.prefix == prefix,
+            ApiKey.key_hash == key_hash,
+            ApiKey.is_active.is_(True),
+        )
+        .where((ApiKey.expires_at.is_(None)) | (ApiKey.expires_at >= func.now()))
+    )
+
+    result = await db.execute(stmt)
+    api_key = result.scalar_one_or_none()
+
+    if api_key is None:
+        return None
+
+    api_key.last_used_at = func.now()
+
+    user = await db.get(User, api_key.user_id)
+    if user is None or not user.is_active:
+        return None
+
+    return user
+
+
+async def get_credit_balance(
+    db: AsyncSession,
+    user_id: UUID,
+) -> int:
+    """Compute the current credit balance from the ledger."""
+    stmt = select(func.coalesce(func.sum(CreditLedger.amount), 0)).where(
+        CreditLedger.user_id == user_id,
+    )
+    result = await db.execute(stmt)
+    return result.scalar_one()
+
+
+async def debit_credits(
+    db: AsyncSession,
+    user_id: UUID,
+    amount: int,
+    tx_type: CreditTxType,
+    description: str,
+    reference_id: UUID | None = None,
+) -> int:
+    """Debit credits and return new balance.
+
+    Raises ValueError if insufficient credits.
+    """
+    balance = await get_credit_balance(db, user_id)
+    if balance < amount:
+        msg = f"Insufficient credits: need {amount}, have {balance}"
+        raise ValueError(msg)
+
+    entry = CreditLedger(
+        user_id=user_id,
+        amount=-amount,
+        tx_type=tx_type,
+        description=description,
+        reference_id=reference_id,
+    )
+    db.add(entry)
+    await db.flush()
+
+    return balance - amount