Files
opnsense-sftp/database.py
T
jamie 68a8bdfe9e
Release / release (pull_request) Successful in 20s
feat: add ability to change username, password and setup totp
2026-03-23 11:28:56 +00:00

546 lines
22 KiB
Python

"""
Database connection and schema management for OPNsense backup system.
"""
import mysql.connector
from mysql.connector import Error
import os
from contextlib import contextmanager
from datetime import datetime
from typing import Optional, List, Dict, Any
from logger_config import get_logger
logger = get_logger(__name__)
class Database:
"""Handle MariaDB database operations."""
def __init__(self):
self.host = os.getenv('DB_HOST', 'localhost')
self.port = int(os.getenv('DB_PORT', '3306'))
self.database = os.getenv('DB_NAME', 'opnsense_backup')
self.user = os.getenv('DB_USER', 'opnsense_backup')
self.password = os.getenv('DB_PASSWORD', 'changeme')
@contextmanager
def get_connection(self):
"""Get database connection with context manager."""
conn = None
try:
conn = mysql.connector.connect(
host=self.host,
port=self.port,
database=self.database,
user=self.user,
password=self.password
)
yield conn
except Error as e:
logger.error(f"Database connection error: {e}")
raise
finally:
if conn and conn.is_connected():
conn.close()
def init_database(self):
"""Initialize database schema."""
try:
with self.get_connection() as conn:
cursor = conn.cursor()
# Create users table
cursor.execute("""
CREATE TABLE IF NOT EXISTS users (
id INT AUTO_INCREMENT PRIMARY KEY,
username VARCHAR(50) UNIQUE NOT NULL,
password_hash VARCHAR(255) NOT NULL,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
)
""")
# Backward-compatible user auth columns.
cursor.execute("SHOW COLUMNS FROM users LIKE 'is_admin'")
if not cursor.fetchone():
cursor.execute("ALTER TABLE users ADD COLUMN is_admin BOOLEAN NOT NULL DEFAULT FALSE")
cursor.execute("SHOW COLUMNS FROM users LIKE 'totp_secret'")
if not cursor.fetchone():
cursor.execute("ALTER TABLE users ADD COLUMN totp_secret VARCHAR(64) NULL")
cursor.execute("SHOW COLUMNS FROM users LIKE 'totp_enabled'")
if not cursor.fetchone():
cursor.execute("ALTER TABLE users ADD COLUMN totp_enabled BOOLEAN NOT NULL DEFAULT FALSE")
# Create opnsense_instances table
cursor.execute("""
CREATE TABLE IF NOT EXISTS opnsense_instances (
id INT AUTO_INCREMENT PRIMARY KEY,
name VARCHAR(100) NOT NULL,
identifier VARCHAR(100) UNIQUE NOT NULL,
ssh_key_id VARCHAR(50) NOT NULL,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
last_backup TIMESTAMP NULL,
description TEXT
)
""")
# Create ssh_keys table
cursor.execute("""
CREATE TABLE IF NOT EXISTS ssh_keys (
id INT AUTO_INCREMENT PRIMARY KEY,
key_id VARCHAR(50) UNIQUE NOT NULL,
instance_id INT NOT NULL,
public_key TEXT NOT NULL,
private_key_path VARCHAR(255) NOT NULL,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
FOREIGN KEY (instance_id) REFERENCES opnsense_instances(id) ON DELETE CASCADE
)
""")
# Create backups table
cursor.execute("""
CREATE TABLE IF NOT EXISTS backups (
id INT AUTO_INCREMENT PRIMARY KEY,
instance_id INT NOT NULL,
filename VARCHAR(255) NOT NULL,
file_path VARCHAR(500) NOT NULL,
file_size BIGINT NOT NULL,
uploaded_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
FOREIGN KEY (instance_id) REFERENCES opnsense_instances(id) ON DELETE CASCADE
)
""")
# Create backup pruning settings table (single policy row).
cursor.execute("""
CREATE TABLE IF NOT EXISTS backup_prune_settings (
id INT PRIMARY KEY,
enabled BOOLEAN NOT NULL DEFAULT FALSE,
scope_type VARCHAR(10) NOT NULL DEFAULT 'all', /* 'all' or 'instance' */
scope_instance_id INT NULL,
keep_days INT NULL,
keep_count INT NULL,
interval_seconds INT NOT NULL DEFAULT 86400,
last_run_at TIMESTAMP NULL,
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
FOREIGN KEY (scope_instance_id) REFERENCES opnsense_instances(id) ON DELETE SET NULL
)
""")
# Ensure we have exactly one settings row.
cursor.execute("""
INSERT INTO backup_prune_settings
(id, enabled, scope_type, scope_instance_id, keep_days, keep_count, interval_seconds, last_run_at)
VALUES
(1, FALSE, 'all', NULL, NULL, NULL, 86400, NULL)
ON DUPLICATE KEY UPDATE
id = id
""")
conn.commit()
cursor.close()
logger.info("Database schema initialized successfully")
except Error as e:
logger.error(f"Error initializing database: {e}")
raise
def create_user(self, username: str, password_hash: str, is_admin: bool = False) -> Optional[int]:
"""Create a new user."""
try:
with self.get_connection() as conn:
cursor = conn.cursor()
cursor.execute(
"INSERT INTO users (username, password_hash, is_admin) VALUES (%s, %s, %s)",
(username, password_hash, is_admin)
)
conn.commit()
user_id = cursor.lastrowid
cursor.close()
return user_id
except Error as e:
logger.error(f"Error creating user: {e}")
return None
def get_user_by_username(self, username: str) -> Optional[Dict[str, Any]]:
"""Get user by username."""
try:
with self.get_connection() as conn:
cursor = conn.cursor(dictionary=True)
cursor.execute("SELECT * FROM users WHERE username = %s", (username,))
user = cursor.fetchone()
cursor.close()
return user
except Error as e:
logger.error(f"Error getting user: {e}")
return None
def get_user_by_id(self, user_id: int) -> Optional[Dict[str, Any]]:
"""Get user by ID."""
try:
with self.get_connection() as conn:
cursor = conn.cursor(dictionary=True)
cursor.execute("SELECT * FROM users WHERE id = %s", (user_id,))
user = cursor.fetchone()
cursor.close()
return user
except Error as e:
logger.error(f"Error getting user by id: {e}")
return None
def get_all_users(self) -> List[Dict[str, Any]]:
"""Get all users."""
try:
with self.get_connection() as conn:
cursor = conn.cursor(dictionary=True)
cursor.execute(
"""
SELECT id, username, is_admin, totp_enabled, created_at
FROM users
ORDER BY created_at ASC
"""
)
users = cursor.fetchall()
cursor.close()
return users
except Error as e:
logger.error(f"Error getting users: {e}")
return []
def update_user_username(self, user_id: int, username: str) -> bool:
"""Update username for a user."""
try:
with self.get_connection() as conn:
cursor = conn.cursor()
cursor.execute("UPDATE users SET username = %s WHERE id = %s", (username, user_id))
conn.commit()
cursor.close()
return True
except Error as e:
logger.error(f"Error updating username: {e}")
return False
def update_user_password(self, user_id: int, password_hash: str) -> bool:
"""Update password hash for a user."""
try:
with self.get_connection() as conn:
cursor = conn.cursor()
cursor.execute("UPDATE users SET password_hash = %s WHERE id = %s", (password_hash, user_id))
conn.commit()
cursor.close()
return True
except Error as e:
logger.error(f"Error updating password hash: {e}")
return False
def update_user_totp(self, user_id: int, totp_secret: Optional[str], totp_enabled: bool) -> bool:
"""Update TOTP settings for a user."""
try:
with self.get_connection() as conn:
cursor = conn.cursor()
cursor.execute(
"UPDATE users SET totp_secret = %s, totp_enabled = %s WHERE id = %s",
(totp_secret, totp_enabled, user_id),
)
conn.commit()
cursor.close()
return True
except Error as e:
logger.error(f"Error updating TOTP settings: {e}")
return False
def update_user_admin(self, user_id: int, is_admin: bool) -> bool:
"""Update admin flag for a user."""
try:
with self.get_connection() as conn:
cursor = conn.cursor()
cursor.execute("UPDATE users SET is_admin = %s WHERE id = %s", (is_admin, user_id))
conn.commit()
cursor.close()
return True
except Error as e:
logger.error(f"Error updating user admin flag: {e}")
return False
def delete_user(self, user_id: int) -> bool:
"""Delete user by ID."""
try:
with self.get_connection() as conn:
cursor = conn.cursor()
cursor.execute("DELETE FROM users WHERE id = %s", (user_id,))
conn.commit()
cursor.close()
return True
except Error as e:
logger.error(f"Error deleting user: {e}")
return False
def create_instance(self, name: str, identifier: str, ssh_key_id: str, description: str = "") -> Optional[int]:
"""Create a new OPNsense instance."""
try:
with self.get_connection() as conn:
cursor = conn.cursor()
cursor.execute(
"""INSERT INTO opnsense_instances (name, identifier, ssh_key_id, description)
VALUES (%s, %s, %s, %s)""",
(name, identifier, ssh_key_id, description)
)
conn.commit()
instance_id = cursor.lastrowid
cursor.close()
return instance_id
except Error as e:
logger.error(f"Error creating instance: {e}")
return None
def get_instance_by_identifier(self, identifier: str) -> Optional[Dict[str, Any]]:
"""Get instance by identifier."""
try:
with self.get_connection() as conn:
cursor = conn.cursor(dictionary=True)
cursor.execute(
"SELECT * FROM opnsense_instances WHERE identifier = %s",
(identifier,)
)
instance = cursor.fetchone()
cursor.close()
return instance
except Error as e:
logger.error(f"Error getting instance by identifier: {e}")
return None
def get_all_instances(self) -> List[Dict[str, Any]]:
"""Get all instances."""
try:
with self.get_connection() as conn:
cursor = conn.cursor(dictionary=True)
cursor.execute("SELECT * FROM opnsense_instances ORDER BY created_at DESC")
instances = cursor.fetchall()
cursor.close()
return instances
except Error as e:
logger.error(f"Error getting instances: {e}")
return []
def save_ssh_key(self, key_id: str, instance_id: int, public_key: str, private_key_path: str) -> bool:
"""Save SSH key to database."""
try:
with self.get_connection() as conn:
cursor = conn.cursor()
cursor.execute(
"""INSERT INTO ssh_keys (key_id, instance_id, public_key, private_key_path)
VALUES (%s, %s, %s, %s)""",
(key_id, instance_id, public_key, private_key_path)
)
conn.commit()
cursor.close()
return True
except Error as e:
logger.error(f"Error saving SSH key: {e}")
return False
def get_ssh_key_by_key_id(self, key_id: str) -> Optional[Dict[str, Any]]:
"""Get SSH key by key_id."""
try:
with self.get_connection() as conn:
cursor = conn.cursor(dictionary=True)
cursor.execute("SELECT * FROM ssh_keys WHERE key_id = %s", (key_id,))
key = cursor.fetchone()
cursor.close()
return key
except Error as e:
logger.error(f"Error getting SSH key: {e}")
return None
def record_backup(self, instance_id: int, filename: str, file_path: str, file_size: int) -> bool:
"""Record a backup in the database."""
try:
with self.get_connection() as conn:
cursor = conn.cursor()
cursor.execute(
"""INSERT INTO backups (instance_id, filename, file_path, file_size)
VALUES (%s, %s, %s, %s)""",
(instance_id, filename, file_path, file_size)
)
cursor.execute(
"UPDATE opnsense_instances SET last_backup = CURRENT_TIMESTAMP WHERE id = %s",
(instance_id,)
)
conn.commit()
cursor.close()
return True
except Error as e:
logger.error(f"Error recording backup: {e}")
return False
def get_backups_for_instance(self, instance_id: int) -> List[Dict[str, Any]]:
"""Get all backups for an instance."""
try:
with self.get_connection() as conn:
cursor = conn.cursor(dictionary=True)
cursor.execute(
"""SELECT * FROM backups WHERE instance_id = %s
ORDER BY uploaded_at DESC""",
(instance_id,)
)
backups = cursor.fetchall()
cursor.close()
return backups
except Error as e:
logger.error(f"Error getting backups: {e}")
return []
def get_all_backups(self) -> List[Dict[str, Any]]:
"""Get all backups with instance information."""
try:
with self.get_connection() as conn:
cursor = conn.cursor(dictionary=True)
cursor.execute("""
SELECT b.*, o.name as instance_name, o.identifier as instance_identifier
FROM backups b
JOIN opnsense_instances o ON b.instance_id = o.id
ORDER BY b.uploaded_at DESC
""")
backups = cursor.fetchall()
cursor.close()
return backups
except Error as e:
logger.error(f"Error getting all backups: {e}")
return []
def get_instance_by_id(self, instance_id: int) -> Optional[Dict[str, Any]]:
"""Get instance by ID."""
try:
with self.get_connection() as conn:
cursor = conn.cursor(dictionary=True)
cursor.execute("SELECT * FROM opnsense_instances WHERE id = %s", (instance_id,))
instance = cursor.fetchone()
cursor.close()
return instance
except Error as e:
logger.error(f"Error getting instance by ID: {e}")
return None
def get_latest_backup_per_instance(self) -> List[Dict[str, Any]]:
"""Get the latest backup date and time for each instance."""
try:
with self.get_connection() as conn:
cursor = conn.cursor(dictionary=True)
cursor.execute("""
SELECT
o.id as instance_id,
o.name as instance_name,
o.identifier as instance_identifier,
MAX(b.uploaded_at) as latest_backup
FROM opnsense_instances o
LEFT JOIN backups b ON o.id = b.instance_id
GROUP BY o.id, o.name, o.identifier
ORDER BY o.name
""")
results = cursor.fetchall()
cursor.close()
return results
except Error as e:
logger.error(f"Error getting latest backup per instance: {e}")
return []
def get_backup_prune_settings(self) -> Dict[str, Any]:
"""Get automated backup prune settings (single row)."""
try:
with self.get_connection() as conn:
cursor = conn.cursor(dictionary=True)
cursor.execute("SELECT * FROM backup_prune_settings WHERE id = %s", (1,))
row = cursor.fetchone()
cursor.close()
if row:
return row
except Error as e:
logger.error(f"Error getting backup prune settings: {e}")
# Safe defaults if table/row doesn't exist yet.
return {
"id": 1,
"enabled": False,
"scope_type": "all",
"scope_instance_id": None,
"keep_days": None,
"keep_count": None,
"interval_seconds": 86400,
"last_run_at": None,
"updated_at": None,
}
def upsert_backup_prune_settings(
self,
enabled: bool,
scope_type: str,
scope_instance_id: Optional[int],
keep_days: Optional[int],
keep_count: Optional[int],
interval_seconds: int,
) -> None:
"""Upsert automated backup prune settings (single row)."""
try:
with self.get_connection() as conn:
cursor = conn.cursor()
cursor.execute(
"""
INSERT INTO backup_prune_settings
(id, enabled, scope_type, scope_instance_id, keep_days, keep_count, interval_seconds)
VALUES
(1, %s, %s, %s, %s, %s, %s)
ON DUPLICATE KEY UPDATE
enabled = VALUES(enabled),
scope_type = VALUES(scope_type),
scope_instance_id = VALUES(scope_instance_id),
keep_days = VALUES(keep_days),
keep_count = VALUES(keep_count),
interval_seconds = VALUES(interval_seconds)
""",
(
bool(enabled),
scope_type,
scope_instance_id,
keep_days,
keep_count,
interval_seconds,
),
)
conn.commit()
cursor.close()
except Error as e:
logger.error(f"Error updating backup prune settings: {e}")
def set_backup_prune_last_run_at(self, last_run_at: datetime) -> None:
"""Update last_run_at after a prune run."""
try:
with self.get_connection() as conn:
cursor = conn.cursor()
cursor.execute(
"UPDATE backup_prune_settings SET last_run_at = %s WHERE id = %s",
(last_run_at, 1),
)
conn.commit()
cursor.close()
except Error as e:
logger.error(f"Error updating backup prune last_run_at: {e}")
def delete_backups_by_ids(self, backup_ids: List[int]) -> int:
"""Delete backup records by IDs (returns number of deleted rows)."""
if not backup_ids:
return 0
try:
placeholders = ", ".join(["%s"] * len(backup_ids))
with self.get_connection() as conn:
cursor = conn.cursor()
cursor.execute(
f"DELETE FROM backups WHERE id IN ({placeholders})",
tuple(backup_ids),
)
affected = cursor.rowcount or 0
conn.commit()
cursor.close()
return affected
except Error as e:
logger.error(f"Error deleting backups by ids: {e}")
return 0