ansible-ui/database.py

import sqlite3
from contextlib import contextmanager
import os
from crypto_utils import CryptoUtils

class Database:
    def __init__(self, db_path="db/ansible.db"):
        self.db_path = db_path
        self.crypto = CryptoUtils()

        db_dir = os.path.dirname(self.db_path)
        if not os.path.exists(db_dir):
            os.makedirs(db_dir, exist_ok=True)

        self.init_database()

    def init_database(self):
        """初始化数据库表"""
        with self.get_connection() as conn:
            conn.execute("""
                CREATE TABLE IF NOT EXISTS hosts (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    comment TEXT NOT NULL,
                    address TEXT NOT NULL,
                    username TEXT NOT NULL,
                    port INTEGER NOT NULL,
                    password TEXT,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    auth_method TEXT NOT NULL DEFAULT 'password'
                )
            """)
            
            conn.execute("""
                CREATE TABLE IF NOT EXISTS command_logs (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    host_id INTEGER,
                    command TEXT NOT NULL,
                    output TEXT,
                    status TEXT NOT NULL,
                    executed_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    FOREIGN KEY (host_id) REFERENCES hosts (id)
                )
            """)

            conn.execute("""
                CREATE TABLE IF NOT EXISTS access_logs (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    ip_address TEXT NOT NULL,
                    path TEXT NOT NULL,
                    status TEXT NOT NULL,
                    status_code INTEGER NOT NULL,
                    access_time TIMESTAMP DEFAULT (datetime('now', '+8 hours'))
                )
            """)

    def init_users_table(self):
        """初始化用户表"""
        with self.get_connection() as conn:
            conn.execute("""
                CREATE TABLE IF NOT EXISTS users (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    username TEXT NOT NULL,
                    password TEXT NOT NULL,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            """)

    @contextmanager
    def get_connection(self):
        """获取数据库连接的上下文管理器"""
        conn = sqlite3.connect(self.db_path)
        conn.row_factory = sqlite3.Row
        try:
            yield conn
            conn.commit()
        except Exception as e:
            conn.rollback()
            raise e
        finally:
            conn.close()

    def add_host(self, host_data):
        """添加单个主机"""
        with self.get_connection() as conn:
            encrypted_password = None
            auth_method = host_data.get('auth_method', 'password')
            
            if auth_method == 'password' and host_data.get('password'):
                encrypted_password = self.crypto.encrypt(host_data['password'])

            cursor = conn.execute("""
                INSERT INTO hosts (comment, address, username, port, password, auth_method)
                VALUES (?, ?, ?, ?, ?, ?)
            """, (
                host_data['comment'],
                host_data['address'],
                host_data['username'],
                host_data['port'],
                encrypted_password,
                auth_method
            ))
            return cursor.lastrowid

    def add_hosts_batch(self, hosts_data):
        """批量添加主机"""
        with self.get_connection() as conn:
            processed_hosts = []
            for host in hosts_data:
                encrypted_password = None
                auth_method = host.get('auth_method', 'password')
                
                if auth_method == 'password' and host.get('password'):
                    encrypted_password = self.crypto.encrypt(host['password'])
                
                processed_hosts.append((
                    host['comment'],
                    host['address'],
                    host['username'],
                    host['port'],
                    encrypted_password,
                    auth_method
                ))
                
            cursor = conn.executemany("""
                INSERT INTO hosts (comment, address, username, port, password, auth_method)
                VALUES (?, ?, ?, ?, ?, ?)
            """, processed_hosts)
            return cursor.rowcount

    def get_hosts(self):
        """获取所有主机"""
        with self.get_connection() as conn:
            cursor = conn.execute("SELECT * FROM hosts ORDER BY created_at DESC")
            hosts = [dict(row) for row in cursor.fetchall()]
            
            for host in hosts:
                host['encrypted_password'] = host['password']
                if host['auth_method'] == 'password' and host['password']:
                    host['password'] = self.crypto.decrypt(host['password'])
                else:
                    host['password'] = None
            return hosts

    def get_host(self, host_id):
        """获取单个主机信息"""
        with self.get_connection() as conn:
            cursor = conn.execute("SELECT * FROM hosts WHERE id = ?", (host_id,))
            row = cursor.fetchone()
            if row:
                host = dict(row)
                host['encrypted_password'] = host['password']
                if host['auth_method'] == 'password' and host['password']:
                    host['password'] = self.crypto.decrypt(host['password'])
                else:
                    host['password'] = None
                return host
            return None

    def update_host(self, host_id, host_data):
        """更新主机信息"""
        with self.get_connection() as conn:
            auth_method = host_data.get('auth_method', 'password')
            current_host = conn.execute(
                "SELECT password FROM hosts WHERE id = ?",
                (host_id,)
            ).fetchone()
            current_password = current_host["password"] if current_host else None

            if auth_method == 'key':
                encrypted_password = None
            elif host_data.get('password'):
                encrypted_password = self.crypto.encrypt(host_data['password'])
            else:
                encrypted_password = current_password
            
            conn.execute("""
                UPDATE hosts 
                SET comment = ?, address = ?, username = ?, port = ?, password = ?, auth_method = ?
                WHERE id = ?
            """, (
                host_data['comment'],
                host_data['address'],
                host_data['username'],
                host_data['port'],
                encrypted_password,
                auth_method,
                host_id
            ))

    def delete_host(self, host_id):
        """删除主机"""
        with self.get_connection() as conn:
            conn.execute("DELETE FROM command_logs WHERE host_id = ?", (host_id,))
            conn.execute("DELETE FROM hosts WHERE id = ?", (host_id,))

    def log_command(self, host_id, command, output, status):
        """记录命令执行日志"""
        with self.get_connection() as conn:
            conn.execute("""
                INSERT INTO command_logs (host_id, command, output, status)
                VALUES (?, ?, ?, ?)
            """, (host_id, command, output, status))

    def get_command_logs(self, limit=100):
        """获取命令执行日志"""
        with self.get_connection() as conn:
            cursor = conn.execute("""
                SELECT cl.*, h.comment, h.address 
                FROM command_logs cl
                LEFT JOIN hosts h ON cl.host_id = h.id
                ORDER BY cl.executed_at DESC
                LIMIT ?
            """, (limit,))
            return [dict(row) for row in cursor.fetchall()]

    def add_access_log(self, ip_address, path, status, status_code):
        """添加访问日志"""
        with self.get_connection() as conn:
            conn.execute("""
                INSERT INTO access_logs (ip_address, path, status, status_code)
                VALUES (?, ?, ?, ?)
            """, (ip_address, path, status, status_code))

    def get_access_logs(self, limit=100, ip_filter='', path_filter=''):
        """获取访问日志"""
        with self.get_connection() as conn:
            clauses = []
            params = []

            if ip_filter:
                clauses.append("ip_address LIKE ?")
                params.append(f"%{ip_filter}%")

            if path_filter:
                clauses.append("path LIKE ?")
                params.append(f"%{path_filter}%")

            where_sql = f"WHERE {' AND '.join(clauses)}" if clauses else ""
            query = f"""
                SELECT * FROM access_logs
                {where_sql}
                ORDER BY access_time DESC
                LIMIT ?
            """
            params.append(limit)

            cursor = conn.execute(query, tuple(params))
            return [dict(row) for row in cursor.fetchall()]

    def cleanup_old_logs(self):
        """清理7天前的日志（使用北京时间）"""
        with self.get_connection() as conn:
            conn.execute("""
                DELETE FROM access_logs 
                WHERE access_time < datetime('now', '+8 hours', '-7 days')
            """)