utm_storage_backends_oss.py

#!/usr/bin/env python3
"""
UTM Protocol Storage Backends - Open Source Edition
Multi-storage backend system for UTM Protocol (open-source version).
No Cloudflare dependencies - uses local storage only.
"""

import time
import json
import os
import pickle
import logging
import sqlite3
from typing import Dict, List, Any, Optional, Union
from abc import ABC, abstractmethod
from dataclasses import dataclass, asdict
import hashlib
import base64
import asyncio
import threading
import queue

# Quantum-safe cryptography imports (optional for OSS)
try:
    import oqs
    QUANTUM_SAFE_AVAILABLE = True
except ImportError:
    QUANTUM_SAFE_AVAILABLE = False

@dataclass
class UTMStorageConfig:
    """Configuration for UTM storage backends (OSS version)"""
    backend_type: str  # local, file, direct
    primary_backend: str
    fallback_backends: List[str]
    local_config: Optional[Dict[str, Any]] = None
    file_config: Optional[Dict[str, Any]] = None
    direct_config: Optional[Dict[str, Any]] = None
    encryption_enabled: bool = True
    compression_enabled: bool = True
    quantum_safe_enabled: bool = False  # Disabled by default in OSS
    version: str = "1.0.0-oss"

@dataclass
class UTMStorageResult:
    """Result of UTM storage operation"""
    success: bool
    data: Optional[Any] = None
    error: Optional[str] = None
    backend_used: Optional[str] = None
    timestamp: Optional[float] = None
    metadata: Optional[Dict[str, Any]] = None

class UTMStorageBackend(ABC):
    """Abstract base class for UTM storage backends"""
    
    def __init__(self, config: Dict[str, Any]):
        self.config = config
        self.logger = logging.getLogger(__name__)
    
    @abstractmethod
    def store(self, key: str, data: Any, metadata: Optional[Dict[str, Any]] = None) -> UTMStorageResult:
        """Store data with the given key"""
        pass
    
    @abstractmethod
    def retrieve(self, key: str) -> UTMStorageResult:
        """Retrieve data by key"""
        pass
    
    @abstractmethod
    def delete(self, key: str) -> UTMStorageResult:
        """Delete data by key"""
        pass
    
    @abstractmethod
    def list_keys(self, prefix: str = "") -> UTMStorageResult:
        """List all keys with optional prefix"""
        pass
    
    def _serialize_data(self, data: Any) -> str:
        """Serialize data for storage"""
        if self.config.get('compression_enabled', True):
            return base64.b64encode(pickle.dumps(data)).decode('utf-8')
        else:
            return json.dumps(data, default=str)
    
    def _deserialize_data(self, serialized_data: str) -> Any:
        """Deserialize data from storage"""
        if self.config.get('compression_enabled', True):
            return pickle.loads(base64.b64decode(serialized_data))
        else:
            return json.loads(serialized_data)

class LocalStorageBackend(UTMStorageBackend):
    """SQLite-based local storage backend for UTM tokens"""
    
    def __init__(self, config: Dict[str, Any]):
        super().__init__(config)
        self.db_path = config.get('db_path', 'utm_tokens.db')
        self._init_database()
    
    def _init_database(self):
        """Initialize SQLite database with UTM token tables"""
        try:
            with sqlite3.connect(self.db_path) as conn:
                conn.execute('''
                    CREATE TABLE IF NOT EXISTS utm_tokens (
                        token_id TEXT PRIMARY KEY,
                        token_data TEXT NOT NULL,
                        metadata TEXT,
                        created_at REAL NOT NULL,
                        expires_at REAL,
                        last_accessed REAL,
                        access_count INTEGER DEFAULT 0
                    )
                ''')
                
                conn.execute('''
                    CREATE INDEX IF NOT EXISTS idx_expires_at 
                    ON utm_tokens(expires_at)
                ''')
                
                conn.execute('''
                    CREATE INDEX IF NOT EXISTS idx_created_at 
                    ON utm_tokens(created_at)
                ''')
                
                conn.commit()
                self.logger.info(f"Initialized UTM token database: {self.db_path}")
        except Exception as e:
            self.logger.error(f"Failed to initialize database: {e}")
            raise
    
    def store(self, key: str, data: Any, metadata: Optional[Dict[str, Any]] = None) -> UTMStorageResult:
        """Store UTM token in SQLite database"""
        try:
            serialized_data = self._serialize_data(data)
            serialized_metadata = json.dumps(metadata) if metadata else None
            
            # Calculate expiration time
            expires_at = time.time() + data.get('ttl', 86400) if isinstance(data, dict) else time.time() + 86400
            
            with sqlite3.connect(self.db_path) as conn:
                conn.execute('''
                    INSERT OR REPLACE INTO utm_tokens 
                    (token_id, token_data, metadata, created_at, expires_at, last_accessed, access_count)
                    VALUES (?, ?, ?, ?, ?, ?, ?)
                ''', (key, serialized_data, serialized_metadata, time.time(), expires_at, time.time(), 0))
                conn.commit()
            
            return UTMStorageResult(
                success=True,
                data=data,
                backend_used="sqlite_local",
                timestamp=time.time(),
                metadata=metadata
            )
        except Exception as e:
            self.logger.error(f"Failed to store token {key}: {e}")
            return UTMStorageResult(
                success=False,
                error=str(e),
                backend_used="sqlite_local"
            )
    
    def retrieve(self, key: str) -> UTMStorageResult:
        """Retrieve UTM token from SQLite database"""
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.execute('''
                    SELECT token_data, metadata, expires_at 
                    FROM utm_tokens 
                    WHERE token_id = ?
                ''', (key,))
                
                row = cursor.fetchone()
                if not row:
                    return UTMStorageResult(
                        success=False,
                        error="Token not found",
                        backend_used="sqlite_local"
                    )
                
                token_data, metadata_str, expires_at = row
                
                # Check expiration
                if expires_at and time.time() > expires_at:
                    return UTMStorageResult(
                        success=False,
                        error="Token expired",
                        backend_used="sqlite_local"
                    )
                
                # Update access statistics
                conn.execute('''
                    UPDATE utm_tokens 
                    SET last_accessed = ?, access_count = access_count + 1 
                    WHERE token_id = ?
                ''', (time.time(), key))
                conn.commit()
                
                data = self._deserialize_data(token_data)
                metadata = json.loads(metadata_str) if metadata_str else None
                
                return UTMStorageResult(
                    success=True,
                    data=data,
                    backend_used="sqlite_local",
                    timestamp=time.time(),
                    metadata=metadata
                )
        except Exception as e:
            self.logger.error(f"Failed to retrieve token {key}: {e}")
            return UTMStorageResult(
                success=False,
                error=str(e),
                backend_used="sqlite_local"
            )
    
    def delete(self, key: str) -> UTMStorageResult:
        """Delete UTM token from SQLite database"""
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.execute('DELETE FROM utm_tokens WHERE token_id = ?', (key,))
                conn.commit()
                
                if cursor.rowcount == 0:
                    return UTMStorageResult(
                        success=False,
                        error="Token not found",
                        backend_used="sqlite_local"
                    )
                
                return UTMStorageResult(
                    success=True,
                    backend_used="sqlite_local",
                    timestamp=time.time()
                )
        except Exception as e:
            self.logger.error(f"Failed to delete token {key}: {e}")
            return UTMStorageResult(
                success=False,
                error=str(e),
                backend_used="sqlite_local"
            )
    
    def list_keys(self, prefix: str = "") -> UTMStorageResult:
        """List UTM token keys with optional prefix"""
        try:
            with sqlite3.connect(self.db_path) as conn:
                if prefix:
                    cursor = conn.execute('''
                        SELECT token_id, created_at, expires_at 
                        FROM utm_tokens 
                        WHERE token_id LIKE ? 
                        ORDER BY created_at DESC
                    ''', (f"{prefix}%",))
                else:
                    cursor = conn.execute('''
                        SELECT token_id, created_at, expires_at 
                        FROM utm_tokens 
                        ORDER BY created_at DESC
                    ''')
                
                rows = cursor.fetchall()
                keys = [{"key": row[0], "created_at": row[1], "expires_at": row[2]} for row in rows]
                
                return UTMStorageResult(
                    success=True,
                    data=keys,
                    backend_used="sqlite_local",
                    timestamp=time.time()
                )
        except Exception as e:
            self.logger.error(f"Failed to list keys: {e}")
            return UTMStorageResult(
                success=False,
                error=str(e),
                backend_used="sqlite_local"
            )
    
    def cleanup_expired(self) -> UTMStorageResult:
        """Remove expired tokens from database"""
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.execute('DELETE FROM utm_tokens WHERE expires_at < ?', (time.time(),))
                deleted_count = cursor.rowcount
                conn.commit()
                
                return UTMStorageResult(
                    success=True,
                    data={"deleted_count": deleted_count},
                    backend_used="sqlite_local",
                    timestamp=time.time()
                )
        except Exception as e:
            self.logger.error(f"Failed to cleanup expired tokens: {e}")
            return UTMStorageResult(
                success=False,
                error=str(e),
                backend_used="sqlite_local"
            )

class FileStorageBackend(UTMStorageBackend):
    """File-based storage backend for UTM tokens (JSON files)"""
    
    def __init__(self, config: Dict[str, Any]):
        super().__init__(config)
        self.storage_dir = config.get('storage_dir', './utm_tokens')
        self._ensure_storage_dir()
    
    def _ensure_storage_dir(self):
        """Ensure storage directory exists"""
        try:
            os.makedirs(self.storage_dir, exist_ok=True)
            self.logger.info(f"Using file storage directory: {self.storage_dir}")
        except Exception as e:
            self.logger.error(f"Failed to create storage directory: {e}")
            raise
    
    def _get_file_path(self, key: str) -> str:
        """Get file path for a given key"""
        # Use hash of key for filename to avoid filesystem issues
        key_hash = hashlib.sha256(key.encode()).hexdigest()
        return os.path.join(self.storage_dir, f"{key_hash}.json")
    
    def store(self, key: str, data: Any, metadata: Optional[Dict[str, Any]] = None) -> UTMStorageResult:
        """Store UTM token as JSON file"""
        try:
            file_path = self._get_file_path(key)
            
            token_record = {
                "key": key,
                "data": data,
                "metadata": metadata,
                "created_at": time.time(),
                "expires_at": time.time() + (data.get('ttl', 86400) if isinstance(data, dict) else 86400)
            }
            
            with open(file_path, 'w') as f:
                json.dump(token_record, f, indent=2, default=str)
            
            return UTMStorageResult(
                success=True,
                data=data,
                backend_used="file_json",
                timestamp=time.time(),
                metadata=metadata
            )
        except Exception as e:
            self.logger.error(f"Failed to store token {key}: {e}")
            return UTMStorageResult(
                success=False,
                error=str(e),
                backend_used="file_json"
            )
    
    def retrieve(self, key: str) -> UTMStorageResult:
        """Retrieve UTM token from JSON file"""
        try:
            file_path = self._get_file_path(key)
            
            if not os.path.exists(file_path):
                return UTMStorageResult(
                    success=False,
                    error="Token not found",
                    backend_used="file_json"
                )
            
            with open(file_path, 'r') as f:
                token_record = json.load(f)
            
            # Check expiration
            if time.time() > token_record.get('expires_at', 0):
                return UTMStorageResult(
                    success=False,
                    error="Token expired",
                    backend_used="file_json"
                )
            
            return UTMStorageResult(
                success=True,
                data=token_record['data'],
                backend_used="file_json",
                timestamp=time.time(),
                metadata=token_record.get('metadata')
            )
        except Exception as e:
            self.logger.error(f"Failed to retrieve token {key}: {e}")
            return UTMStorageResult(
                success=False,
                error=str(e),
                backend_used="file_json"
            )
    
    def delete(self, key: str) -> UTMStorageResult:
        """Delete UTM token JSON file"""
        try:
            file_path = self._get_file_path(key)
            
            if not os.path.exists(file_path):
                return UTMStorageResult(
                    success=False,
                    error="Token not found",
                    backend_used="file_json"
                )
            
            os.remove(file_path)
            
            return UTMStorageResult(
                success=True,
                backend_used="file_json",
                timestamp=time.time()
            )
        except Exception as e:
            self.logger.error(f"Failed to delete token {key}: {e}")
            return UTMStorageResult(
                success=False,
                error=str(e),
                backend_used="file_json"
            )
    
    def list_keys(self, prefix: str = "") -> UTMStorageResult:
        """List UTM token keys from JSON files"""
        try:
            keys = []
            
            for filename in os.listdir(self.storage_dir):
                if filename.endswith('.json'):
                    file_path = os.path.join(self.storage_dir, filename)
                    try:
                        with open(file_path, 'r') as f:
                            token_record = json.load(f)
                        
                        key = token_record.get('key', '')
                        if not prefix or key.startswith(prefix):
                            keys.append({
                                "key": key,
                                "created_at": token_record.get('created_at'),
                                "expires_at": token_record.get('expires_at')
                            })
                    except Exception as e:
                        self.logger.warning(f"Failed to read token file {filename}: {e}")
                        continue
            
            # Sort by creation time
            keys.sort(key=lambda x: x.get('created_at', 0), reverse=True)
            
            return UTMStorageResult(
                success=True,
                data=keys,
                backend_used="file_json",
                timestamp=time.time()
            )
        except Exception as e:
            self.logger.error(f"Failed to list keys: {e}")
            return UTMStorageResult(
                success=False,
                error=str(e),
                backend_used="file_json"
            )

class UTMStorageManager:
    """UTM Storage Manager for open-source edition"""
    
    def __init__(self, config: UTMStorageConfig):
        self.config = config
        self.logger = logging.getLogger(__name__)
        self.backend = self._create_backend()
        self.fallback_backends = [self._create_fallback_backend(backend_type) for backend_type in config.fallback_backends]
    
    def _create_backend(self) -> UTMStorageBackend:
        """Create primary storage backend"""
        if self.config.primary_backend == "local":
            return LocalStorageBackend(self.config.local_config or {})
        elif self.config.primary_backend == "file":
            return FileStorageBackend(self.config.file_config or {})
        else:
            raise ValueError(f"Unsupported primary backend: {self.config.primary_backend}")
    
    def _create_fallback_backend(self, backend_type: str) -> UTMStorageBackend:
        """Create fallback storage backend"""
        if backend_type == "local":
            return LocalStorageBackend(self.config.local_config or {})
        elif backend_type == "file":
            return FileStorageBackend(self.config.file_config or {})
        else:
            raise ValueError(f"Unsupported fallback backend: {backend_type}")
    
    def store_utm_token(self, token_id: str, token_data: Dict[str, Any]) -> UTMStorageResult:
        """Store UTM token with fallback support"""
        result = self.backend.store(token_id, token_data)
        
        if not result.success and self.fallback_backends:
            self.logger.warning(f"Primary backend failed, trying fallback: {result.error}")
            for fallback in self.fallback_backends:
                result = fallback.store(token_id, token_data)
                if result.success:
                    break
        
        return result
    
    def retrieve_utm_token(self, token_id: str) -> UTMStorageResult:
        """Retrieve UTM token with fallback support"""
        result = self.backend.retrieve(token_id)
        
        if not result.success and self.fallback_backends:
            self.logger.warning(f"Primary backend failed, trying fallback: {result.error}")
            for fallback in self.fallback_backends:
                result = fallback.retrieve(token_id)
                if result.success:
                    break
        
        return result
    
    def delete_utm_token(self, token_id: str) -> UTMStorageResult:
        """Delete UTM token from all backends"""
        results = []
        
        # Try primary backend
        result = self.backend.delete(token_id)
        results.append(result)
        
        # Try fallback backends
        for fallback in self.fallback_backends:
            result = fallback.delete(token_id)
            results.append(result)
        
        # Return success if any backend succeeded
        for result in results:
            if result.success:
                return result
        
        return UTMStorageResult(
            success=False,
            error="Failed to delete from all backends",
            backend_used="all"
        )
    
    def list_utm_tokens(self, prefix: str = "") -> UTMStorageResult:
        """List UTM tokens from primary backend"""
        return self.backend.list_keys(prefix)
    
    def cleanup_expired_tokens(self) -> UTMStorageResult:
        """Cleanup expired tokens from all backends"""
        results = []
        
        # Try primary backend
        if hasattr(self.backend, 'cleanup_expired'):
            result = self.backend.cleanup_expired()
            results.append(result)
        
        # Try fallback backends
        for fallback in self.fallback_backends:
            if hasattr(fallback, 'cleanup_expired'):
                result = fallback.cleanup_expired()
                results.append(result)
        
        # Aggregate results
        total_deleted = sum(r.data.get('deleted_count', 0) for r in results if r.success and r.data)
        
        return UTMStorageResult(
            success=True,
            data={"total_deleted": total_deleted, "backend_results": results},
            backend_used="all",
            timestamp=time.time()
        )

def create_utm_storage_manager(config: UTMStorageConfig) -> UTMStorageManager:
    """Factory function to create UTM Storage Manager"""
    return UTMStorageManager(config)

def create_oss_config(
    primary_backend: str = "local",
    fallback_backends: List[str] = None,
    db_path: str = "utm_tokens.db",
    storage_dir: str = "./utm_tokens"
) -> UTMStorageConfig:
    """Create configuration for open-source UTM storage"""
    if fallback_backends is None:
        fallback_backends = ["file"]
    
    return UTMStorageConfig(
        backend_type="oss",
        primary_backend=primary_backend,
        fallback_backends=fallback_backends,
        local_config={"db_path": db_path},
        file_config={"storage_dir": storage_dir},
        encryption_enabled=True,
        compression_enabled=True,
        quantum_safe_enabled=False,  # Disabled in OSS
        version="1.0.0-oss"
    )