backend/security/key_usage_tracker.py

import sqlite3
import uuid
import logging
from datetime import datetime, timezone
from backend.services.usb_monitor import get_db_path
from backend.services.usb_vault import KeyDomain

def track_token_usage(domain: KeyDomain, tokens: int):
    """Logs token usage for a specific KeyDomain."""
    try:
        db_path = get_db_path()
        with sqlite3.connect(db_path) as conn:
            conn.execute('PRAGMA journal_mode=WAL')
            conn.execute(
                "INSERT INTO key_usage (id, domain, tokens_used, timestamp) VALUES (?, ?, ?, ?)",
                (str(uuid.uuid4()), domain.value, tokens, datetime.now(timezone.utc).isoformat())
            )
            conn.commit()
    except Exception as e:
        logging.error(f"Failed to track token usage for {domain.value}: {e}")

def get_today_usage_by_domain():
    """Returns aggregated token usage per domain for the current day."""
    try:
        db_path = get_db_path()
        today = datetime.now(timezone.utc).strftime('%Y-%m-%d')
        with sqlite3.connect(db_path) as conn:
            conn.execute('PRAGMA journal_mode=WAL')
            cursor = conn.cursor()
            # SQLite DATE() extracts YYYY-MM-DD from ISO format string
            cursor.execute(
                "SELECT domain, SUM(tokens_used) FROM key_usage WHERE DATE(timestamp) = ? GROUP BY domain",
                (today,)
            )
            rows = cursor.fetchall()
            return {row[0]: row[1] for row in rows}
    except Exception as e:
        logging.error(f"Failed to fetch key usage: {e}")
        return {}