Initial commit
Build and Push Multi-Arch Docker Image / build (push) Has been cancelled
Build and Push Multi-Arch Docker Image / build (push) Has been cancelled
This commit is contained in:
dido committed
commit
8ce01cb00a
34 files changed
+8346
No files matched your search
@@ -0,0 +1,458 @@
|
||||
"""
|
||||
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
|
||||
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)
|
||||
)
|
||||
conn.commit()
|
||||
conn.close()
|
||||
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 []
|
||||
|
||||
|
||||
# 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
|
||||
Reference in new issue
Block a user