"""SQLite database for token storage.""" from __future__ import annotations import aiosqlite import json import base64 import os from pathlib import Path from datetime import datetime from typing import Optional from cryptography.fernet import Fernet DATABASE_PATH = Path(__file__).parent.parent / "classlens.db" def get_encryption_key() -> bytes: """Get or generate encryption key.""" key = os.getenv("ENCRYPTION_KEY", "") if not key: # Generate a new key for development (should be set in production) key = Fernet.generate_key().decode() print(f"⚠️ No ENCRYPTION_KEY set. Generated temporary key: {key}") print("⚠️ Set this in your .env file for persistent token storage.") return key.encode() if isinstance(key, str) else key def encrypt_token(token_data: dict) -> str: """Encrypt token data for storage.""" key = get_encryption_key() f = Fernet(key) json_data = json.dumps(token_data) encrypted = f.encrypt(json_data.encode()) return base64.b64encode(encrypted).decode() def decrypt_token(encrypted_data: str) -> dict: """Decrypt stored token data.""" key = get_encryption_key() f = Fernet(key) encrypted = base64.b64decode(encrypted_data.encode()) decrypted = f.decrypt(encrypted) return json.loads(decrypted.decode()) async def init_database(): """Initialize the database with required tables.""" async with aiosqlite.connect(DATABASE_PATH) as db: await db.execute(""" CREATE TABLE IF NOT EXISTS oauth_tokens ( teacher_email TEXT PRIMARY KEY, encrypted_tokens TEXT NOT NULL, created_at TEXT NOT NULL, updated_at TEXT NOT NULL ) """) await db.execute(""" CREATE TABLE IF NOT EXISTS analysis_reports ( id INTEGER PRIMARY KEY AUTOINCREMENT, teacher_email TEXT NOT NULL, exam_title TEXT, report_markdown TEXT, report_json TEXT, created_at TEXT NOT NULL ) """) await db.commit() async def save_oauth_tokens(teacher_email: str, tokens: dict): """Save encrypted OAuth tokens for a teacher.""" encrypted = encrypt_token(tokens) now = datetime.utcnow().isoformat() async with aiosqlite.connect(DATABASE_PATH) as db: await db.execute(""" INSERT INTO oauth_tokens (teacher_email, encrypted_tokens, created_at, updated_at) VALUES (?, ?, ?, ?) ON CONFLICT(teacher_email) DO UPDATE SET encrypted_tokens = excluded.encrypted_tokens, updated_at = excluded.updated_at """, (teacher_email, encrypted, now, now)) await db.commit() async def get_oauth_tokens(teacher_email: str) -> Optional[dict]: """Retrieve OAuth tokens for a teacher.""" async with aiosqlite.connect(DATABASE_PATH) as db: async with db.execute( "SELECT encrypted_tokens FROM oauth_tokens WHERE teacher_email = ?", (teacher_email,) ) as cursor: row = await cursor.fetchone() if row: return decrypt_token(row[0]) return None async def delete_oauth_tokens(teacher_email: str): """Delete OAuth tokens for a teacher.""" async with aiosqlite.connect(DATABASE_PATH) as db: await db.execute( "DELETE FROM oauth_tokens WHERE teacher_email = ?", (teacher_email,) ) await db.commit() async def save_report(teacher_email: str, exam_title: str, report_markdown: str, report_json: str): """Save an analysis report.""" now = datetime.utcnow().isoformat() async with aiosqlite.connect(DATABASE_PATH) as db: await db.execute(""" INSERT INTO analysis_reports (teacher_email, exam_title, report_markdown, report_json, created_at) VALUES (?, ?, ?, ?, ?) """, (teacher_email, exam_title, report_markdown, report_json, now)) await db.commit() async def get_teacher_reports(teacher_email: str, limit: int = 10) -> list[dict]: """Get recent reports for a teacher.""" async with aiosqlite.connect(DATABASE_PATH) as db: db.row_factory = aiosqlite.Row async with db.execute(""" SELECT id, exam_title, report_markdown, report_json, created_at FROM analysis_reports WHERE teacher_email = ? ORDER BY created_at DESC LIMIT ? """, (teacher_email, limit)) as cursor: rows = await cursor.fetchall() return [dict(row) for row in rows]