opsxshell/electron/services/database-service.ts

import initSqlJs, { Database as SqlJsDatabase } from 'sql.js'
import path from 'path'
import { app } from 'electron'
import fs from 'fs'
import { simpleEncrypt, simpleDecrypt } from './crypto'
import { keyService } from './key-service'
import { createLogger } from '../utils/logger'

const logger = createLogger('DatabaseService')

// sql.js statement wrapper to match better-sqlite3 API
interface StatementResult {
  id: number
  [key: string]: any
}

class PreparedStatement {
  private db: SqlJsDatabase
  private sql: string

  constructor(db: SqlJsDatabase, sql: string) {
    this.db = db
    this.sql = sql
  }

  run(...params: any[]): StatementResult {
    this.db.run(this.sql, params as any)
    return { id: 0 } as StatementResult
  }

  get(...params: any[]): any {
    // 使用 exec 模式获取单行
    const stmt = this.db.prepare(this.sql)
    stmt.bind(params as any)
    let result: any = null
    if (stmt.step()) {
      result = stmt.getAsObject()
    }
    stmt.free()
    return result
  }

  all(...params: any[]): any[] {
    const stmt = this.db.prepare(this.sql)
    stmt.bind(params as any)
    const results: any[] = []
    while (stmt.step()) {
      results.push(stmt.getAsObject())
    }
    stmt.free()
    return results
  }
}

export class DatabaseService {
  private db: SqlJsDatabase | null = null
  private dbPath: string
  private masterKey: string | null = null
  private initialized = false

  constructor() {
    const userDataPath = app.getPath('userData')
    this.dbPath = path.join(userDataPath, 'opsxshell.db')
  }

  async initialize(): Promise<void> {
    try {
      const dbDir = path.dirname(this.dbPath)
      if (!fs.existsSync(dbDir)) {
        fs.mkdirSync(dbDir, { recursive: true })
      }

      this.masterKey = keyService.getMasterKey()

      // 初始化 sql.js (WASM, 纯 JS 无需原生编译)
      const SQL = await initSqlJs()

      // 如果已有数据库文件，加载它
      if (fs.existsSync(this.dbPath)) {
        const buffer = fs.readFileSync(this.dbPath)
        this.db = new SQL.Database(buffer)
      } else {
        this.db = new SQL.Database()
      }

      this.createTables()
      this.save()

      this.initialized = true
      logger.info('Database initialized at:', this.dbPath)
    } catch (error) {
      logger.error('Failed to initialize database:', error)
      throw error
    }
  }

  private createTables(): void {
    if (!this.db) return

    this.db.run(`
      CREATE TABLE IF NOT EXISTS sessions (
        id TEXT PRIMARY KEY,
        name TEXT NOT NULL,
        host TEXT NOT NULL,
        port INTEGER DEFAULT 22,
        username TEXT NOT NULL,
        auth_type TEXT CHECK(auth_type IN ('password', 'key', 'agent')),
        password TEXT,
        private_key_path TEXT,
        passphrase_encrypted TEXT,
        group_name TEXT DEFAULT 'default',
        tags TEXT,
        created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
        updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
      )
    `)

    // 迁移：如果 password 列不存在，添加它
    try {
      this.db.run(`ALTER TABLE sessions ADD COLUMN password TEXT`)
    } catch (e) {
      // 列可能已存在，忽略错误
    }

    this.db.run(`
      CREATE TABLE IF NOT EXISTS connection_history (
        id TEXT PRIMARY KEY,
        session_id TEXT REFERENCES sessions(id),
        connected_at DATETIME DEFAULT CURRENT_TIMESTAMP,
        disconnected_at DATETIME,
        duration INTEGER,
        status TEXT
      )
    `)

    this.db.run(`
      CREATE TABLE IF NOT EXISTS settings (
        key TEXT PRIMARY KEY,
        value TEXT,
        updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
      )
    `)

    this.db.run(`
      CREATE TABLE IF NOT EXISTS keys (
        id TEXT PRIMARY KEY,
        name TEXT NOT NULL,
        type TEXT CHECK(type IN ('rsa', 'ed25519', 'ecdsa')),
        private_key_encrypted TEXT NOT NULL,
        public_key TEXT,
        passphrase_encrypted TEXT,
        created_at DATETIME DEFAULT CURRENT_TIMESTAMP
      )
    `)

    this.db.run(`
      CREATE TABLE IF NOT EXISTS port_forwards (
        id TEXT PRIMARY KEY,
        session_id TEXT REFERENCES sessions(id),
        type TEXT CHECK(type IN ('local', 'remote', 'dynamic')),
        local_host TEXT,
        local_port INTEGER,
        remote_host TEXT,
        remote_port INTEGER,
        is_active BOOLEAN DEFAULT FALSE
      )
    `)
  }

  // 持久化到磁盘
  private save(): void {
    if (!this.db) return
    try {
      const data = this.db.export()
      const buffer = Buffer.from(data)
      fs.writeFileSync(this.dbPath, buffer)
    } catch (error) {
      logger.error('Failed to save database:', error)
    }
  }

  private prepare(sql: string): PreparedStatement {
    if (!this.db) throw new Error('Database not initialized')
    return new PreparedStatement(this.db, sql)
  }

  // 加密字段
  private encryptField(value: string | null | undefined): string | null {
    if (!value) return null
    if (!this.masterKey) throw new Error('Master key not initialized')
    return simpleEncrypt(value, this.masterKey)
  }

  // 解密字段
  private decryptField(value: string | null | undefined): string | null {
    if (!value) return null
    if (!this.masterKey) {
      logger.error('Decrypt failed: Master key not initialized')
      return null
    }
    try {
      const decrypted = simpleDecrypt(value, this.masterKey)
      logger.info('Decrypted field length:', decrypted?.length || 0)
      return decrypted
    } catch (error) {
      logger.error('Decrypt failed:', error)
      return null
    }
  }

  // 会话 CRUD 操作
  createSession(session: any): any {
    if (!this.db) throw new Error('Database not initialized')

    logger.info('Creating session with password:', session.password ? `len=${session.password.length}` : 'null/undefined')

    const stmt = this.prepare(`
      INSERT INTO sessions (id, name, host, port, username, auth_type, password, private_key_path, passphrase_encrypted, group_name, tags)
      VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
    `)

    const id = `session-${Date.now()}`
    const encryptedPassword = this.encryptField(session.password)
    logger.info('Encrypted password:', encryptedPassword ? `len=${encryptedPassword.length}` : 'null')

    stmt.run(
      id,
      session.name,
      session.host,
      session.port || 22,
      session.username,
      session.authType,
      encryptedPassword,
      session.privateKey || null,
      this.encryptField(session.passphrase),
      session.group || 'default',
      session.tags ? JSON.stringify(session.tags) : null
    )

    this.save()
    return this.getSessionById(id)
  }

  getAllSessions(): any[] {
    if (!this.db) throw new Error('Database not initialized')

    const stmt = this.prepare('SELECT * FROM sessions ORDER BY created_at DESC')
    const sessions = stmt.all()

    return sessions.map((s: any) => ({
      id: s.id,
      name: s.name,
      host: s.host,
      port: s.port,
      username: s.username,
      authType: s.auth_type,
      password: this.decryptField(s.password),
      privateKey: s.private_key_path,
      passphrase: this.decryptField(s.passphrase_encrypted),
      group: s.group_name,
      tags: s.tags ? JSON.parse(s.tags) : [],
    }))
  }

  getSessionById(id: string): any {
    if (!this.db) throw new Error('Database not initialized')

    const stmt = this.prepare('SELECT * FROM sessions WHERE id = ?')
    const session = stmt.get(id) as any

    if (session) {
      return {
        id: session.id,
        name: session.name,
        host: session.host,
        port: session.port,
        username: session.username,
        authType: session.auth_type,
        password: this.decryptField(session.password),
        privateKey: session.private_key_path,
        passphrase: this.decryptField(session.passphrase_encrypted),
        group: session.group_name,
        tags: session.tags ? JSON.parse(session.tags) : [],
      }
    }
    return null
  }

  updateSession(id: string, updates: any): any {
    if (!this.db) throw new Error('Database not initialized')

    const fields: string[] = []
    const values: any[] = []

    if (updates.name !== undefined) {
      fields.push('name = ?')
      values.push(updates.name)
    }
    if (updates.host !== undefined) {
      fields.push('host = ?')
      values.push(updates.host)
    }
    if (updates.port !== undefined) {
      fields.push('port = ?')
      values.push(updates.port)
    }
    if (updates.username !== undefined) {
      fields.push('username = ?')
      values.push(updates.username)
    }
    if (updates.authType !== undefined) {
      fields.push('auth_type = ?')
      values.push(updates.authType)
    }
    if (updates.password !== undefined) {
      fields.push('password = ?')
      values.push(this.encryptField(updates.password))
    }
    if (updates.passphrase !== undefined) {
      fields.push('passphrase_encrypted = ?')
      values.push(this.encryptField(updates.passphrase))
    }
    if (updates.group !== undefined) {
      fields.push('group_name = ?')
      values.push(updates.group)
    }
    if (updates.tags !== undefined) {
      fields.push('tags = ?')
      values.push(JSON.stringify(updates.tags))
    }

    fields.push('updated_at = CURRENT_TIMESTAMP')
    values.push(id)

    const stmt = this.prepare(`
      UPDATE sessions SET ${fields.join(', ')} WHERE id = ?
    `)
    stmt.run(...values)

    this.save()
    return this.getSessionById(id)
  }

  deleteSession(id: string): void {
    if (!this.db) throw new Error('Database not initialized')

    const stmt = this.prepare('DELETE FROM sessions WHERE id = ?')
    stmt.run(id)
    this.save()
  }

  // 设置操作
  getSetting(key: string): any {
    if (!this.db) throw new Error('Database not initialized')

    const stmt = this.prepare('SELECT value FROM settings WHERE key = ?')
    const row: any = stmt.get(key)
    return row ? JSON.parse(row.value) : null
  }

  setSetting(key: string, value: any): void {
    if (!this.db) throw new Error('Database not initialized')

    const stmt = this.prepare(`
      INSERT INTO settings (key, value, updated_at)
      VALUES (?, ?, CURRENT_TIMESTAMP)
      ON CONFLICT(key) DO UPDATE SET
        value = excluded.value,
        updated_at = CURRENT_TIMESTAMP
    `)
    stmt.run(key, JSON.stringify(value))
    this.save()
  }

  close(): void {
    if (this.db) {
      this.save()
      this.db.close()
      this.db = null
    }
  }
}