diff --git a/src/proxy_pool/accounts/service.py b/src/proxy_pool/accounts/service.py new file mode 100644 index 0000000..f11034a --- /dev/null +++ b/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