lpra/internal/server/db.go

package server

import (
	"database/sql"
	"fmt"
	"strings"
	"time"

	_ "modernc.org/sqlite"
)

// DB wraps the SQLite connection and all queries used by the server.
type DB struct {
	sql *sql.DB
}

// OpenDB opens (creating if needed) the SQLite database at path and runs
// migrations. The pure-Go driver ("sqlite") needs no CGO.
func OpenDB(path string) (*DB, error) {
	// _pragma busy_timeout avoids "database is locked" under the concurrent
	// long-poll load; foreign_keys keeps the schema honest.
	dsn := fmt.Sprintf("file:%s?_pragma=busy_timeout(5000)&_pragma=foreign_keys(1)&_pragma=journal_mode(WAL)", path)
	sdb, err := sql.Open("sqlite", dsn)
	if err != nil {
		return nil, err
	}
	// SQLite writes serialize; a single connection avoids lock churn while the
	// busy_timeout still covers the long-poll readers.
	sdb.SetMaxOpenConns(1)
	if err := sdb.Ping(); err != nil {
		return nil, err
	}
	db := &DB{sql: sdb}
	if err := db.migrate(); err != nil {
		return nil, err
	}
	return db, nil
}

func (db *DB) Close() error { return db.sql.Close() }

func (db *DB) migrate() error {
	const schema = `
CREATE TABLE IF NOT EXISTS config (
    key   TEXT PRIMARY KEY,
    value TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS admins (
    id         INTEGER PRIMARY KEY AUTOINCREMENT,
    username   TEXT UNIQUE NOT NULL,
    pass_hash  TEXT NOT NULL,
    created_at INTEGER NOT NULL
);
CREATE TABLE IF NOT EXISTS agents (
    id          TEXT PRIMARY KEY,
    name        TEXT NOT NULL,
    hostname    TEXT NOT NULL,
    os          TEXT NOT NULL,
    arch        TEXT NOT NULL,
    token_hash  TEXT NOT NULL,
    version     TEXT NOT NULL DEFAULT '',
    run_user    TEXT NOT NULL DEFAULT '',
    login_user  TEXT NOT NULL DEFAULT '',
    enrolled_at INTEGER NOT NULL,
    last_seen   INTEGER NOT NULL
);
CREATE TABLE IF NOT EXISTS admin_sessions (
    id         TEXT PRIMARY KEY,
    admin_id   INTEGER NOT NULL REFERENCES admins(id) ON DELETE CASCADE,
    created_at INTEGER NOT NULL,
    expires_at INTEGER NOT NULL
);
CREATE TABLE IF NOT EXISTS audit (
    id       INTEGER PRIMARY KEY AUTOINCREMENT,
    admin_id INTEGER,
    agent_id TEXT,
    action   TEXT NOT NULL,
    at       INTEGER NOT NULL
);
CREATE TABLE IF NOT EXISTS ufw_rules (
    id         INTEGER PRIMARY KEY AUTOINCREMENT,
    port       INTEGER NOT NULL,
    proto      TEXT NOT NULL,
    rule_id    INTEGER NOT NULL DEFAULT 0,
    created_at INTEGER NOT NULL,
    UNIQUE(port, proto)
);
CREATE TABLE IF NOT EXISTS forward_rules (
    id          INTEGER PRIMARY KEY AUTOINCREMENT,
    agent_id    TEXT NOT NULL,
    client_port INTEGER NOT NULL,
    host_port   INTEGER NOT NULL,
    listen_addr TEXT NOT NULL DEFAULT '0.0.0.0',
    proto       TEXT NOT NULL DEFAULT 'http',
    created_at  INTEGER NOT NULL
);`
	if _, err := db.sql.Exec(schema); err != nil {
		return err
	}
	// Add columns to databases created before they existed. SQLite has no
	// "IF NOT EXISTS" for columns, so ignore the "duplicate column" error.
	for _, alter := range []string{
		`ALTER TABLE forward_rules ADD COLUMN proto TEXT NOT NULL DEFAULT 'http'`,
		`ALTER TABLE agents ADD COLUMN version TEXT NOT NULL DEFAULT ''`,
		`ALTER TABLE ufw_rules ADD COLUMN rule_id INTEGER NOT NULL DEFAULT 0`,
		`ALTER TABLE agents ADD COLUMN run_user TEXT NOT NULL DEFAULT ''`,
		`ALTER TABLE agents ADD COLUMN login_user TEXT NOT NULL DEFAULT ''`,
		`ALTER TABLE forward_rules ADD COLUMN no_ufw INTEGER NOT NULL DEFAULT 0`,
	} {
		if _, err := db.sql.Exec(alter); err != nil && !strings.Contains(err.Error(), "duplicate column") {
			return err
		}
	}
	return nil
}

// --- config kv ---

func (db *DB) GetConfig(key string) (string, error) {
	var v string
	err := db.sql.QueryRow(`SELECT value FROM config WHERE key = ?`, key).Scan(&v)
	if err == sql.ErrNoRows {
		return "", nil
	}
	return v, err
}

func (db *DB) SetConfig(key, value string) error {
	_, err := db.sql.Exec(
		`INSERT INTO config(key, value) VALUES(?, ?)
		 ON CONFLICT(key) DO UPDATE SET value = excluded.value`, key, value)
	return err
}

// --- admins ---

func (db *DB) CountAdmins() (int, error) {
	var n int
	err := db.sql.QueryRow(`SELECT COUNT(*) FROM admins`).Scan(&n)
	return n, err
}

func (db *DB) CreateAdmin(username, passHash string) error {
	_, err := db.sql.Exec(
		`INSERT INTO admins(username, pass_hash, created_at) VALUES(?, ?, ?)`,
		username, passHash, time.Now().Unix())
	return err
}

// Admin is one row of the admins table.
type Admin struct {
	ID       int64
	Username string
	PassHash string
}

func (db *DB) AdminByUsername(username string) (*Admin, error) {
	a := &Admin{}
	err := db.sql.QueryRow(
		`SELECT id, username, pass_hash FROM admins WHERE username = ?`, username).
		Scan(&a.ID, &a.Username, &a.PassHash)
	if err == sql.ErrNoRows {
		return nil, nil
	}
	return a, err
}

func (db *DB) ListAdmins() ([]Admin, error) {
	rows, err := db.sql.Query(`SELECT id, username FROM admins ORDER BY username`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []Admin
	for rows.Next() {
		var a Admin
		if err := rows.Scan(&a.ID, &a.Username); err != nil {
			return nil, err
		}
		out = append(out, a)
	}
	return out, rows.Err()
}

func (db *DB) DeleteAdmin(id int64) error {
	_, err := db.sql.Exec(`DELETE FROM admins WHERE id = ?`, id)
	return err
}

// --- admin sessions ---

func (db *DB) CreateSession(id string, adminID int64, ttl time.Duration) error {
	now := time.Now()
	_, err := db.sql.Exec(
		`INSERT INTO admin_sessions(id, admin_id, created_at, expires_at) VALUES(?, ?, ?, ?)`,
		id, adminID, now.Unix(), now.Add(ttl).Unix())
	return err
}

// SessionAdmin returns the admin for a live (non-expired) session, or nil.
func (db *DB) SessionAdmin(sessionID string) (*Admin, error) {
	a := &Admin{}
	err := db.sql.QueryRow(
		`SELECT a.id, a.username, a.pass_hash
		   FROM admin_sessions s JOIN admins a ON a.id = s.admin_id
		  WHERE s.id = ? AND s.expires_at > ?`, sessionID, time.Now().Unix()).
		Scan(&a.ID, &a.Username, &a.PassHash)
	if err == sql.ErrNoRows {
		return nil, nil
	}
	return a, err
}

func (db *DB) DeleteSession(id string) error {
	_, err := db.sql.Exec(`DELETE FROM admin_sessions WHERE id = ?`, id)
	return err
}

// --- agents ---

// Agent is one row of the agents table.
type Agent struct {
	ID        string
	Name      string
	Hostname  string
	OS        string
	Arch      string
	TokenHash string
	Version   string
	RunUser   string // OS user the agent process runs as (reported by the agent)
	LoginUser string // OS user terminals are opened as; "" = RunUser
	LastSeen  time.Time
}

// TerminalUser is the OS user a terminal on this agent runs as. Agents are
// installed as a root systemd service, so that is the fallback for agents too
// old to report their user.
func (a Agent) TerminalUser() string {
	return firstNonEmpty(a.LoginUser, a.RunUser, "root")
}

// recentlySeen is the fallback for "online" right after a server start, while
// the hub has not seen the agents' polls yet (see Server.online).
func (a Agent) recentlySeen() bool {
	return time.Since(a.LastSeen) < 45*time.Second
}

func (db *DB) CreateAgent(a Agent) error {
	now := time.Now().Unix()
	_, err := db.sql.Exec(
		`INSERT INTO agents(id, name, hostname, os, arch, token_hash, version, run_user, enrolled_at, last_seen)
		 VALUES(?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
		a.ID, a.Name, a.Hostname, a.OS, a.Arch, a.TokenHash, a.Version, a.RunUser, now, now)
	return err
}

func (db *DB) GetAgent(id string) (*Agent, error) {
	a := &Agent{}
	var last int64
	err := db.sql.QueryRow(
		`SELECT id, name, hostname, os, arch, token_hash, version, run_user, login_user, last_seen FROM agents WHERE id = ?`, id).
		Scan(&a.ID, &a.Name, &a.Hostname, &a.OS, &a.Arch, &a.TokenHash, &a.Version, &a.RunUser, &a.LoginUser, &last)
	if err == sql.ErrNoRows {
		return nil, nil
	}
	if err != nil {
		return nil, err
	}
	a.LastSeen = time.Unix(last, 0)
	return a, nil
}

func (db *DB) ListAgents() ([]Agent, error) {
	rows, err := db.sql.Query(
		`SELECT id, name, hostname, os, arch, version, run_user, login_user, last_seen FROM agents ORDER BY name`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []Agent
	for rows.Next() {
		var a Agent
		var last int64
		if err := rows.Scan(&a.ID, &a.Name, &a.Hostname, &a.OS, &a.Arch, &a.Version, &a.RunUser, &a.LoginUser, &last); err != nil {
			return nil, err
		}
		a.LastSeen = time.Unix(last, 0)
		out = append(out, a)
	}
	return out, rows.Err()
}

func (db *DB) TouchAgent(id string) error {
	_, err := db.sql.Exec(`UPDATE agents SET last_seen = ? WHERE id = ?`, time.Now().Unix(), id)
	return err
}

// TouchAgentInfo updates last_seen together with the reported version and OS user.
func (db *DB) TouchAgentInfo(id, version, runUser string) error {
	_, err := db.sql.Exec(`UPDATE agents SET last_seen = ?, version = ?, run_user = ? WHERE id = ?`,
		time.Now().Unix(), version, runUser, id)
	return err
}

// SetAgentLoginUser sets the OS user terminals are opened as ("" = the agent's own).
func (db *DB) SetAgentLoginUser(id, user string) error {
	_, err := db.sql.Exec(`UPDATE agents SET login_user = ? WHERE id = ?`, user, id)
	return err
}

func (db *DB) DeleteAgent(id string) error {
	_, err := db.sql.Exec(`DELETE FROM agents WHERE id = ?`, id)
	return err
}

func (db *DB) RenameAgent(id, name string) error {
	_, err := db.sql.Exec(`UPDATE agents SET name = ? WHERE id = ?`, name, id)
	return err
}

// --- forward rules ---

// ForwardRule maps a port on the agent's localhost to a port on the host.
// Proto is "http" (long-poll), "tcp" (direct mux) or "udp" (direct).
type ForwardRule struct {
	ID         int64
	AgentID    string
	ClientPort int
	HostPort   int
	ListenAddr string
	Proto      string
	NoUfw      bool // the panel opens no ufw rule for this forward
}

func (db *DB) CreateForwardRule(r ForwardRule) (int64, error) {
	if r.Proto == "" {
		r.Proto = "http"
	}
	res, err := db.sql.Exec(
		`INSERT INTO forward_rules(agent_id, client_port, host_port, listen_addr, proto, no_ufw, created_at)
		 VALUES(?, ?, ?, ?, ?, ?, ?)`,
		r.AgentID, r.ClientPort, r.HostPort, r.ListenAddr, r.Proto, r.NoUfw, time.Now().Unix())
	if err != nil {
		return 0, err
	}
	return res.LastInsertId()
}

func (db *DB) ListForwardRules() ([]ForwardRule, error) {
	rows, err := db.sql.Query(
		`SELECT id, agent_id, client_port, host_port, listen_addr, proto, no_ufw FROM forward_rules ORDER BY id`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []ForwardRule
	for rows.Next() {
		var r ForwardRule
		if err := rows.Scan(&r.ID, &r.AgentID, &r.ClientPort, &r.HostPort, &r.ListenAddr, &r.Proto, &r.NoUfw); err != nil {
			return nil, err
		}
		out = append(out, r)
	}
	return out, rows.Err()
}

func (db *DB) ListForwardRulesByAgent(agentID string) ([]ForwardRule, error) {
	all, err := db.ListForwardRules()
	if err != nil {
		return nil, err
	}
	var out []ForwardRule
	for _, r := range all {
		if r.AgentID == agentID {
			out = append(out, r)
		}
	}
	return out, nil
}

func (db *DB) GetForwardRule(id int64) (*ForwardRule, error) {
	r := &ForwardRule{}
	err := db.sql.QueryRow(
		`SELECT id, agent_id, client_port, host_port, listen_addr, proto, no_ufw FROM forward_rules WHERE id = ?`, id).
		Scan(&r.ID, &r.AgentID, &r.ClientPort, &r.HostPort, &r.ListenAddr, &r.Proto, &r.NoUfw)
	if err == sql.ErrNoRows {
		return nil, nil
	}
	return r, err
}

func (db *DB) DeleteForwardRule(id int64) error {
	_, err := db.sql.Exec(`DELETE FROM forward_rules WHERE id = ?`, id)
	return err
}

// --- ufw rules created by the panel ---

type UfwRule struct {
	ID     int64
	Port   int
	Proto  string
	RuleID int64 // bound forward rule id, or 0 for shared/manual (e.g. tunnel 8446)
}

func (db *DB) AddUfwRule(port int, proto string, ruleID int64) error {
	_, err := db.sql.Exec(
		`INSERT OR IGNORE INTO ufw_rules(port, proto, rule_id, created_at) VALUES(?, ?, ?, ?)`,
		port, proto, ruleID, time.Now().Unix())
	return err
}

func (db *DB) HasUfwRule(port int, proto string) bool {
	var n int
	_ = db.sql.QueryRow(`SELECT COUNT(*) FROM ufw_rules WHERE port = ? AND proto = ?`, port, proto).Scan(&n)
	return n > 0
}

func (db *DB) scanUfwRules(query string, args ...any) ([]UfwRule, error) {
	rows, err := db.sql.Query(query, args...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []UfwRule
	for rows.Next() {
		var r UfwRule
		if err := rows.Scan(&r.ID, &r.Port, &r.Proto, &r.RuleID); err != nil {
			return nil, err
		}
		out = append(out, r)
	}
	return out, rows.Err()
}

func (db *DB) ListUfwRules() ([]UfwRule, error) {
	return db.scanUfwRules(`SELECT id, port, proto, rule_id FROM ufw_rules ORDER BY port, proto`)
}

// ListUfwRulesByRule returns ufw rules bound to a specific forward rule.
func (db *DB) ListUfwRulesByRule(ruleID int64) ([]UfwRule, error) {
	return db.scanUfwRules(`SELECT id, port, proto, rule_id FROM ufw_rules WHERE rule_id = ?`, ruleID)
}

func (db *DB) GetUfwRule(id int64) (*UfwRule, error) {
	r := &UfwRule{}
	err := db.sql.QueryRow(`SELECT id, port, proto, rule_id FROM ufw_rules WHERE id = ?`, id).
		Scan(&r.ID, &r.Port, &r.Proto, &r.RuleID)
	if err == sql.ErrNoRows {
		return nil, nil
	}
	return r, err
}

func (db *DB) DeleteUfwRule(id int64) error {
	_, err := db.sql.Exec(`DELETE FROM ufw_rules WHERE id = ?`, id)
	return err
}

// HostPortUsed reports whether any forward rule already uses this host port.
func (db *DB) HostPortUsed(hostPort int) bool {
	var n int
	_ = db.sql.QueryRow(`SELECT COUNT(*) FROM forward_rules WHERE host_port = ?`, hostPort).Scan(&n)
	return n > 0
}

// --- audit ---

func (db *DB) Audit(adminID int64, agentID, action string) {
	// Best-effort: an audit failure must not block the operation.
	_, _ = db.sql.Exec(
		`INSERT INTO audit(admin_id, agent_id, action, at) VALUES(?, ?, ?, ?)`,
		adminID, agentID, action, time.Now().Unix())
}