528 lines
22 KiB
Python
528 lines
22 KiB
Python
"""
|
|
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
|