backend/vault.py

import os
import sqlite3
import logging
from typing import Optional

def get_db_path():
    if "JARVIS_APP_DATA_DIR" in os.environ:
        return os.path.join(os.environ["JARVIS_APP_DATA_DIR"], "memory.db")
    project_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
    return os.path.join(project_root, "memory.db")

async def resolve_vault_key(key_name: str) -> Optional[str]:
    """Retrieve a token from the SQLite vault secrets."""
    try:
        from backend.services.usb_vault import get_secret
        return get_secret(key_name)
    except Exception as e:
        logging.error(f"Failed to fetch vault key {key_name}: {e}")
        return None

async def exchange_refresh_for_access(refresh_token: str) -> Optional[str]:
    """Exchanges a long-lived Google refresh token for a short-lived access token."""
    import aiohttp
    client_id = os.environ.get("GMAIL_OAUTH_CLIENT_ID")
    client_secret = os.environ.get("GMAIL_OAUTH_CLIENT_SECRET")
    
    if not refresh_token or not client_id or not client_secret:
        logging.error("Missing Google OAuth credentials.")
        return None

    try:
        async with aiohttp.ClientSession() as session:
            data = {
                "client_id": client_id,
                "client_secret": client_secret,
                "refresh_token": refresh_token,
                "grant_type": "refresh_token"
            }
            async with session.post("https://oauth2.googleapis.com/token", data=data) as response:
                if response.status == 200:
                    json_data = await response.json()
                    return json_data.get("access_token")
                else:
                    logging.error(f"Failed to refresh Google token: {await response.text()}")
                    return None
    except Exception as e:
        logging.error(f"Error during token exchange: {e}")
        return None

async def resolve_vault_key_oauth(key_name: str) -> Optional[str]:
    """
    Convenience method for connectors expecting an OAuth access token.
    For GMAIL_OAUTH, it fetches the refresh token and exchanges it.
    """
    if key_name == "GMAIL_OAUTH":
        refresh_token = await resolve_vault_key("GMAIL_REFRESH_TOKEN")
        # If not in vault, try environment variable as fallback
        if not refresh_token:
            refresh_token = os.environ.get("GMAIL_REFRESH_TOKEN")
        return await exchange_refresh_for_access(refresh_token)
    return None