""" Database module - CredentialsManager for SQLite operations. """ import os import sqlite3 from typing import Any, Dict, List, Optional from cryptography.fernet import Fernet from .config import ENCRYPTION_KEY, STORAGE_FILE, logger class CredentialsManager: """Manages encrypted storage of SSH credentials per user.""" def __init__(self, db_path: str, encryption_key: str): self.db_path = db_path self.cipher = Fernet(encryption_key.encode() if isinstance(encryption_key, str) else encryption_key) # Ensure directory exists before initializing DB db_dir = os.path.dirname(db_path) if db_dir: os.makedirs(db_dir, exist_ok=True) self._init_db() def _init_db(self): """Initialize SQLite database.""" conn = sqlite3.connect(self.db_path) cursor = conn.cursor() cursor.execute(''' CREATE TABLE IF NOT EXISTS connections ( id TEXT PRIMARY KEY, username TEXT NOT NULL, name TEXT NOT NULL, host TEXT NOT NULL, port INTEGER NOT NULL, ssh_username TEXT NOT NULL, auth_type TEXT NOT NULL, password_encrypted BLOB, key_encrypted BLOB, ws_url TEXT, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ) ''') conn.commit() # Migration: Create recordings table if not exists cursor.execute(''' CREATE TABLE IF NOT EXISTS recordings ( id TEXT PRIMARY KEY, username TEXT NOT NULL, connection_id TEXT NOT NULL, connection_name TEXT, session_id TEXT NOT NULL, host TEXT NOT NULL, ssh_username TEXT NOT NULL, started_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, ended_at TIMESTAMP, duration_seconds INTEGER, file_size_bytes INTEGER, file_path TEXT NOT NULL, width INTEGER DEFAULT 80, height INTEGER DEFAULT 24, FOREIGN KEY (connection_id) REFERENCES connections(id) ) ''') # Migration: Add connection_name column if it doesn't exist (for existing databases) try: cursor.execute('ALTER TABLE recordings ADD COLUMN connection_name TEXT') logger.info("Added connection_name column to recordings table") except sqlite3.OperationalError: # Column already exists pass cursor.execute('CREATE INDEX IF NOT EXISTS idx_recordings_username ON recordings(username)') cursor.execute('CREATE INDEX IF NOT EXISTS idx_recordings_connection ON recordings(connection_id)') cursor.execute('CREATE INDEX IF NOT EXISTS idx_recordings_started_at ON recordings(started_at)') # Migration: Add lock_timeout and disconnect_timeout columns if they don't exist try: cursor.execute('ALTER TABLE connections ADD COLUMN lock_timeout INTEGER DEFAULT 15') logger.info("Added lock_timeout column to connections table") except sqlite3.OperationalError: pass try: cursor.execute('ALTER TABLE connections ADD COLUMN disconnect_timeout INTEGER DEFAULT 60') logger.info("Added disconnect_timeout column to connections table") except sqlite3.OperationalError: pass cursor.execute(''' CREATE TABLE IF NOT EXISTS telemetry_state ( id INTEGER PRIMARY KEY CHECK (id = 1), installation_uuid TEXT, register_request_id TEXT, credential_id TEXT, credential_secret_encrypted BLOB, expires_at TEXT, last_heartbeat_at REAL, last_heartbeat_request_id TEXT, last_reported_version TEXT, update_request_id TEXT, renewal_request_id TEXT ) ''') conn.commit() conn.close() logger.info(f"Credentials database initialized at {self.db_path}") def _encrypt(self, data: str) -> bytes: """Encrypt string data.""" if not data: return None return self.cipher.encrypt(data.encode()) def _decrypt(self, data: bytes) -> str: """Decrypt bytes to string.""" if not data: return None try: return self.cipher.decrypt(data).decode() except Exception as e: logger.warning(f"Failed to decrypt data: {e}") return None def save_connection(self, username: str, conn_data: Dict[str, Any]) -> bool: """Save or update a connection.""" try: conn = sqlite3.connect(self.db_path) cursor = conn.cursor() conn_id = conn_data.get('id') name = conn_data.get('name', 'Unnamed') host = conn_data.get('host', '') port = conn_data.get('port', 22) ssh_username = conn_data.get('username', '') auth_type = conn_data.get('authType', 'password') ws_url = conn_data.get('wsUrl', 'ws://localhost:2222') lock_timeout = conn_data.get('lockTimeout', 15) disconnect_timeout = conn_data.get('disconnectTimeout', 60) # Encrypt sensitive data password = conn_data.get('password', '') key = conn_data.get('key', '') password_encrypted = self._encrypt(password) if password else None key_encrypted = self._encrypt(key) if key else None cursor.execute(''' INSERT OR REPLACE INTO connections (id, username, name, host, port, ssh_username, auth_type, password_encrypted, key_encrypted, ws_url, lock_timeout, disconnect_timeout, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, CURRENT_TIMESTAMP) ''', (conn_id, username, name, host, port, ssh_username, auth_type, password_encrypted, key_encrypted, ws_url, lock_timeout, disconnect_timeout)) conn.commit() conn.close() logger.info(f"Saved connection {name} for user {username}") return True except Exception as e: logger.error(f"Failed to save connection: {e}") return False def get_connections(self, username: str) -> list: """Get all connections for a user.""" try: conn = sqlite3.connect(self.db_path) cursor = conn.cursor() logger.info(f"Querying connections for username: '{username}'") # DEBUG # First, let's see all connections to debug cursor.execute('SELECT DISTINCT username FROM connections') all_users = cursor.fetchall() logger.info(f"All users in DB: {all_users}") # DEBUG cursor.execute(''' SELECT id, name, host, port, ssh_username, auth_type, password_encrypted, key_encrypted, ws_url, lock_timeout, disconnect_timeout FROM connections WHERE username = ? ''', (username,)) rows = cursor.fetchall() logger.info(f"Found {len(rows)} rows for username '{username}'") # DEBUG conn.close() connections = [] for i, row in enumerate(rows): try: logger.info(f"Processing row {i}: id={row[0]}, name={row[1]}") # DEBUG conn_dict = { 'id': row[0], 'name': row[1], 'host': row[2], 'port': row[3], 'username': row[4], 'authType': row[5], 'password': self._decrypt(row[6]) if row[6] else '', 'key': self._decrypt(row[7]) if row[7] else '', 'wsUrl': row[8] if row[8] else 'ws://localhost:2222', 'lockTimeout': row[9] if row[9] is not None else 15, 'disconnectTimeout': row[10] if row[10] is not None else 60, } connections.append(conn_dict) logger.info(f"Successfully processed row {i}") # DEBUG except Exception as row_error: logger.error(f"Error processing row {i}: {row_error}") import traceback logger.error(traceback.format_exc()) return connections except Exception as e: logger.error(f"Failed to get connections: {e}") import traceback logger.error(traceback.format_exc()) return [] def delete_connection(self, username: str, conn_id: str) -> bool: """Delete a connection.""" try: conn = sqlite3.connect(self.db_path) cursor = conn.cursor() cursor.execute( "DELETE FROM connections WHERE id = ? AND username = ?", (conn_id, username) ) deleted = cursor.rowcount conn.commit() conn.close() if deleted != 1: logger.warning(f"Connection {conn_id} not found for user {username}") return False logger.info(f"Deleted connection {conn_id} for user {username}") return True except Exception as e: logger.error(f"Failed to delete connection: {e}") return False # Session Recording Methods def save_recording(self, recording_data: Dict[str, Any]) -> bool: """Save recording metadata to database.""" try: conn = sqlite3.connect(self.db_path) cursor = conn.cursor() cursor.execute(''' INSERT INTO recordings (id, username, connection_id, connection_name, session_id, host, ssh_username, started_at, ended_at, duration_seconds, file_size_bytes, file_path, width, height) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ''', ( recording_data['id'], recording_data['username'], recording_data['connection_id'], recording_data.get('connection_name', recording_data['host']), recording_data['session_id'], recording_data['host'], recording_data['ssh_username'], recording_data.get('started_at'), recording_data.get('ended_at'), recording_data.get('duration_seconds'), recording_data.get('file_size_bytes'), recording_data['file_path'], recording_data.get('width', 80), recording_data.get('height', 24) )) conn.commit() conn.close() logger.info(f"Saved recording metadata: {recording_data['id']}") return True except Exception as e: logger.error(f"Failed to save recording: {e}") return False def update_recording_end(self, recording_id: str, ended_at: str, duration: int, file_size: int) -> bool: """Update recording with end time and file info.""" try: conn = sqlite3.connect(self.db_path) cursor = conn.cursor() cursor.execute(''' UPDATE recordings SET ended_at = ?, duration_seconds = ?, file_size_bytes = ? WHERE id = ? ''', (ended_at, duration, file_size, recording_id)) conn.commit() conn.close() logger.info(f"Updated recording end: {recording_id}") return True except Exception as e: logger.error(f"Failed to update recording: {e}") return False def get_recordings(self, username: str, connection_id: str = None) -> list: """Get all recordings for a user, optionally filtered by connection.""" try: conn = sqlite3.connect(self.db_path) cursor = conn.cursor() if connection_id: cursor.execute(''' SELECT id, connection_id, connection_name, session_id, host, ssh_username, started_at, ended_at, duration_seconds, file_size_bytes, file_path, width, height FROM recordings WHERE username = ? AND connection_id = ? ORDER BY started_at DESC ''', (username, connection_id)) else: cursor.execute(''' SELECT id, connection_id, connection_name, session_id, host, ssh_username, started_at, ended_at, duration_seconds, file_size_bytes, file_path, width, height FROM recordings WHERE username = ? ORDER BY started_at DESC ''', (username,)) rows = cursor.fetchall() conn.close() recordings = [] for row in rows: recordings.append({ 'id': row[0], 'connection_id': row[1], 'connection_name': row[2], 'session_id': row[3], 'host': row[4], 'ssh_username': row[5], 'started_at': row[6], 'ended_at': row[7], 'duration_seconds': row[8], 'file_size_bytes': row[9], 'file_path': row[10], 'width': row[11], 'height': row[12] }) return recordings except Exception as e: logger.error(f"Failed to get recordings: {e}") return [] def get_recording_by_id(self, recording_id: str, username: str) -> Optional[Dict]: """Get a specific recording by ID.""" try: conn = sqlite3.connect(self.db_path) cursor = conn.cursor() cursor.execute(''' SELECT id, connection_id, connection_name, session_id, host, ssh_username, started_at, ended_at, duration_seconds, file_size_bytes, file_path, width, height FROM recordings WHERE id = ? AND username = ? ''', (recording_id, username)) row = cursor.fetchone() conn.close() if row: return { 'id': row[0], 'connection_id': row[1], 'connection_name': row[2], 'session_id': row[3], 'host': row[4], 'ssh_username': row[5], 'started_at': row[6], 'ended_at': row[7], 'duration_seconds': row[8], 'file_size_bytes': row[9], 'file_path': row[10], 'width': row[11], 'height': row[12] } return None except Exception as e: logger.error(f"Failed to get recording: {e}") return None def delete_recording(self, recording_id: str, username: str) -> bool: """Delete a recording from database.""" try: conn = sqlite3.connect(self.db_path) cursor = conn.cursor() cursor.execute( "DELETE FROM recordings WHERE id = ? AND username = ?", (recording_id, username) ) conn.commit() conn.close() logger.info(f"Deleted recording {recording_id} for user {username}") return True except Exception as e: logger.error(f"Failed to delete recording: {e}") return False def get_connections_with_recordings(self, username: str) -> list: """Get unique connections that have recordings for a user.""" try: conn = sqlite3.connect(self.db_path) cursor = conn.cursor() cursor.execute(''' SELECT DISTINCT connection_id, connection_name, host, ssh_username, COUNT(*) as recording_count, SUM(file_size_bytes) as total_size, MAX(started_at) as last_recording FROM recordings WHERE username = ? AND connection_id IS NOT NULL GROUP BY connection_id, connection_name, host, ssh_username ORDER BY last_recording DESC ''', (username,)) rows = cursor.fetchall() conn.close() return [{ 'connection_id': row[0], 'connection_name': row[1], 'host': row[2], 'ssh_username': row[3], 'recording_count': row[4], 'total_size_bytes': row[5] or 0, 'last_recording': row[6] } for row in rows] except Exception as e: logger.error(f"Failed to get connections with recordings: {e}") return [] def delete_recordings_by_connection(self, username: str, connection_id: str) -> tuple: """Delete all recordings for a specific connection. Returns (success_count, failed_count).""" try: # First get all recordings to delete files recordings = self.get_recordings(username, connection_id) success_count = 0 failed_count = 0 conn = sqlite3.connect(self.db_path) cursor = conn.cursor() for rec in recordings: try: # Delete file if exists file_path = rec.get('file_path') if file_path and os.path.exists(file_path): os.remove(file_path) # Delete from database cursor.execute( "DELETE FROM recordings WHERE id = ? AND username = ?", (rec['id'], username) ) success_count += 1 except Exception as e: logger.error(f"Failed to delete recording {rec.get('id')}: {e}") failed_count += 1 conn.commit() conn.close() logger.info(f"Deleted {success_count} recordings for connection {connection_id}, failed: {failed_count}") return (success_count, failed_count) except Exception as e: logger.error(f"Failed to delete recordings by connection: {e}") return (0, 0) def get_old_recordings(self, before_date: str) -> list: """Get recordings older than a specific date for cleanup.""" try: conn = sqlite3.connect(self.db_path) cursor = conn.cursor() cursor.execute(''' SELECT id, username, file_path FROM recordings WHERE started_at < ? ''', (before_date,)) rows = cursor.fetchall() conn.close() return [{'id': row[0], 'username': row[1], 'file_path': row[2]} for row in rows] except Exception as e: logger.error(f"Failed to get old recordings: {e}") return [] def get_telemetry_state(self) -> Optional[Dict[str, Any]]: """Return the encrypted telemetry state, if it exists.""" try: conn = sqlite3.connect(self.db_path) row = conn.execute('SELECT * FROM telemetry_state WHERE id = 1').fetchone() conn.close() if not row: return None return { 'installation_uuid': row[1], 'register_request_id': row[2], 'credential_id': row[3], 'credential_secret': self._decrypt(row[4]) if row[4] else None, 'expires_at': row[5], 'last_heartbeat_at': row[6], 'last_heartbeat_request_id': row[7], 'last_reported_version': row[8], 'update_request_id': row[9], 'renewal_request_id': row[10], } except Exception as e: logger.warning(f"Failed to read telemetry state: {e}") return None def save_telemetry_state(self, state: Dict[str, Any]) -> bool: """Persist telemetry state with the credential secret encrypted.""" try: conn = sqlite3.connect(self.db_path) conn.execute(''' INSERT OR REPLACE INTO telemetry_state (id, installation_uuid, register_request_id, credential_id, credential_secret_encrypted, expires_at, last_heartbeat_at, last_heartbeat_request_id, last_reported_version, update_request_id, renewal_request_id) VALUES (1, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ''', ( state.get('installation_uuid'), state.get('register_request_id'), state.get('credential_id'), self._encrypt(state['credential_secret']) if state.get('credential_secret') else None, state.get('expires_at'), state.get('last_heartbeat_at'), state.get('last_heartbeat_request_id'), state.get('last_reported_version'), state.get('update_request_id'), state.get('renewal_request_id'), )) conn.commit() conn.close() return True except Exception as e: logger.warning(f"Failed to save telemetry state: {e}") return False # Initialize credentials manager (singleton) try: creds_manager = CredentialsManager(STORAGE_FILE, ENCRYPTION_KEY) logger.info("Credentials manager initialized successfully") except Exception as e: logger.error(f"Failed to initialize credentials manager: {e}") import traceback logger.error(traceback.format_exc()) creds_manager = None # Will fail later when used