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  *dbx
	path string // the file; the key of saved secrets sits beside it
}

// 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: &dbx{DB: sdb}, path: path}
	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 admin_keys (
    id          INTEGER PRIMARY KEY AUTOINCREMENT,
    admin_id    INTEGER NOT NULL REFERENCES admins(id) ON DELETE CASCADE,
    cred_id     BLOB UNIQUE NOT NULL,
    public_key  BLOB NOT NULL,
    sign_count  INTEGER NOT NULL DEFAULT 0,
    name        TEXT NOT NULL DEFAULT '',
    created_at  INTEGER NOT NULL,
    last_used   INTEGER NOT NULL DEFAULT 0
);
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`,
		`ALTER TABLE forward_rules ADD COLUMN label TEXT NOT NULL DEFAULT ''`,
		`ALTER TABLE admins ADD COLUMN password_off INTEGER NOT NULL DEFAULT 0`,
		`ALTER TABLE agents ADD COLUMN transport TEXT NOT NULL DEFAULT 'auto'`,
		`ALTER TABLE admins ADD COLUMN pin_hash TEXT NOT NULL DEFAULT ''`,
		`ALTER TABLE admins ADD COLUMN role TEXT NOT NULL DEFAULT 'admin'`,
		`ALTER TABLE admins ADD COLUMN pw_locked INTEGER NOT NULL DEFAULT 0`,
		`CREATE TABLE IF NOT EXISTS ai_messages (
		    id INTEGER PRIMARY KEY AUTOINCREMENT,
		    admin_id INTEGER NOT NULL REFERENCES admins(id) ON DELETE CASCADE,
		    role TEXT NOT NULL, content TEXT NOT NULL, at INTEGER NOT NULL)`,
		`CREATE INDEX IF NOT EXISTS ai_messages_admin ON ai_messages(admin_id, id)`,
		`ALTER TABLE ai_messages ADD COLUMN chat_id INTEGER NOT NULL DEFAULT 0`,
		`ALTER TABLE ai_messages ADD COLUMN think_ms INTEGER NOT NULL DEFAULT 0`,
		`CREATE TABLE IF NOT EXISTS ai_chats (
		    id INTEGER PRIMARY KEY AUTOINCREMENT,
		    admin_id INTEGER NOT NULL REFERENCES admins(id) ON DELETE CASCADE,
		    title TEXT NOT NULL, created INTEGER NOT NULL, updated INTEGER NOT NULL)`,
		`CREATE INDEX IF NOT EXISTS ai_messages_chat ON ai_messages(chat_id, id)`,
		`CREATE TABLE IF NOT EXISTS ssh_hosts (
		    id INTEGER PRIMARY KEY AUTOINCREMENT,
		    admin_id INTEGER NOT NULL REFERENCES admins(id) ON DELETE CASCADE,
		    label TEXT NOT NULL, host TEXT NOT NULL, port INTEGER NOT NULL, login TEXT NOT NULL,
		    auth TEXT NOT NULL, secret TEXT NOT NULL)`,
		`CREATE TABLE IF NOT EXISTS admin_agents (
		    admin_id INTEGER NOT NULL REFERENCES admins(id) ON DELETE CASCADE,
		    agent_id TEXT NOT NULL,
		    PRIMARY KEY (admin_id, agent_id))`,
	} {
		if _, err := db.sql.Exec(alter); err != nil && !strings.Contains(err.Error(), "duplicate column") {
			return err
		}
	}
	// ssh_hosts.user (a word PostgreSQL keeps for itself) became login.
	if !db.sql.pg {
		_, _ = db.sql.Exec(`ALTER TABLE ssh_hosts RENAME COLUMN user TO login`)
	}
	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.CreateUser(username, passHash, RoleAdmin)
	return err
}

// CreateUser adds a panel user with the role given and returns its id.
func (db *DB) CreateUser(username, passHash, role string) (int64, error) {
	return db.sql.insert(
		`INSERT INTO admins(username, pass_hash, created_at, role) VALUES(?, ?, ?, ?)`,
		username, passHash, time.Now().Unix(), role)
}

// Roles of panel users: an admin sees and runs everything; a user only the
// machines given to them, and none of the settings.
const (
	RoleAdmin = "admin"
	RoleUser  = "user"
)

// Admin is one row of the admins table: a panel user, an admin or not.
type Admin struct {
	ID       int64
	Username string
	PassHash string
	Role     string
	// PwLocked keeps the user from turning the password sign-in on or off
	// themselves: an admin set it.
	PwLocked bool
}

// IsAdmin reports whether the user has the run of the whole panel.
func (a *Admin) IsAdmin() bool { return a != nil && a.Role != RoleUser }

const adminCols = `id, username, pass_hash, role, pw_locked`

func scanAdmin(row interface{ Scan(...any) error }) (*Admin, error) {
	a := &Admin{}
	var locked int
	if err := row.Scan(&a.ID, &a.Username, &a.PassHash, &a.Role, &locked); err != nil {
		return nil, err
	}
	a.PwLocked = locked == 1
	return a, nil
}

func (db *DB) AdminByUsername(username string) (*Admin, error) {
	a, err := scanAdmin(db.sql.QueryRow(`SELECT `+adminCols+` FROM admins WHERE username = ?`, username))
	if err == sql.ErrNoRows {
		return nil, nil
	}
	return a, err
}

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

// CountRole is how many users have the role.
func (db *DB) CountRole(role string) int {
	var n int
	_ = db.sql.QueryRow(`SELECT COUNT(*) FROM admins WHERE role = ?`, role).Scan(&n)
	return n
}

func (db *DB) SetUserRole(id int64, role string) error {
	_, err := db.sql.Exec(`UPDATE admins SET role = ? WHERE id = ?`, role, id)
	return err
}

func (db *DB) SetPwLocked(id int64, locked bool) error {
	_, err := db.sql.Exec(`UPDATE admins SET pw_locked = ? WHERE id = ?`, boolInt(locked), id)
	return err
}

// UserAgents are the machines a user may use.
func (db *DB) UserAgents(adminID int64) map[string]bool {
	out := map[string]bool{}
	rows, err := db.sql.Query(`SELECT agent_id FROM admin_agents WHERE admin_id = ?`, adminID)
	if err != nil {
		return out
	}
	defer rows.Close()
	for rows.Next() {
		var id string
		if rows.Scan(&id) == nil {
			out[id] = true
		}
	}
	return out
}

// SetUserAgents gives the user exactly these machines.
func (db *DB) SetUserAgents(adminID int64, agentIDs []string) error {
	tx, err := db.sql.Begin()
	if err != nil {
		return err
	}
	defer tx.Rollback()
	if _, err := tx.Exec(`DELETE FROM admin_agents WHERE admin_id = ?`, adminID); err != nil {
		return err
	}
	for _, id := range agentIDs {
		if _, err := tx.Exec(`INSERT OR IGNORE INTO admin_agents(admin_id, agent_id) VALUES(?, ?)`, adminID, id); err != nil {
			return err
		}
	}
	return tx.Commit()
}

// CanUseAgent reports whether the user may use the machine.
func (db *DB) CanUseAgent(a *Admin, agentID string) bool {
	if a == nil {
		return false
	}
	if a.IsAdmin() {
		return true
	}
	var n int
	_ = db.sql.QueryRow(`SELECT COUNT(*) FROM admin_agents WHERE admin_id = ? AND agent_id = ?`, a.ID, agentID).Scan(&n)
	return n > 0
}

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

// AdminKey is a passkey an admin signs in with.
type AdminKey struct {
	ID        int64
	AdminID   int64
	CredID    []byte
	PublicKey []byte // PKIX DER
	SignCount uint32
	Name      string
	CreatedAt int64
	LastUsed  int64
}

func (db *DB) AdminByID(id int64) (*Admin, error) {
	a, err := scanAdmin(db.sql.QueryRow(`SELECT `+adminCols+` FROM admins WHERE id = ?`, id))
	if err == sql.ErrNoRows {
		return nil, nil
	}
	return a, err
}

func (db *DB) AddAdminKey(k AdminKey) error {
	_, err := db.sql.Exec(
		`INSERT INTO admin_keys(admin_id, cred_id, public_key, sign_count, name, created_at) VALUES(?, ?, ?, ?, ?, ?)`,
		k.AdminID, k.CredID, k.PublicKey, k.SignCount, k.Name, time.Now().Unix())
	return err
}

func (db *DB) ListAdminKeys(adminID int64) ([]AdminKey, error) {
	rows, err := db.sql.Query(
		`SELECT id, admin_id, cred_id, public_key, sign_count, name, created_at, last_used
		   FROM admin_keys WHERE admin_id = ? ORDER BY id`, adminID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []AdminKey
	for rows.Next() {
		var k AdminKey
		if err := rows.Scan(&k.ID, &k.AdminID, &k.CredID, &k.PublicKey, &k.SignCount, &k.Name, &k.CreatedAt, &k.LastUsed); err != nil {
			return nil, err
		}
		out = append(out, k)
	}
	return out, rows.Err()
}

// PasswordOff tells an admin who turned the password off: then only a
// passkey signs them in. It holds only while they have one.
func (db *DB) PasswordOff(adminID int64) bool {
	var off int
	_ = db.sql.QueryRow(`SELECT password_off FROM admins WHERE id = ?`, adminID).Scan(&off)
	return off == 1 && db.CountAdminKeys(adminID) > 0
}

// AdminPin is the hash of the admin's PIN, "" when there is none.
func (db *DB) AdminPin(adminID int64) string {
	var h string
	_ = db.sql.QueryRow(`SELECT pin_hash FROM admins WHERE id = ?`, adminID).Scan(&h)
	return h
}

func (db *DB) SetAdminPin(adminID int64, pinHash string) error {
	_, err := db.sql.Exec(`UPDATE admins SET pin_hash = ? WHERE id = ?`, pinHash, adminID)
	return err
}

func (db *DB) SetAdminPassword(adminID int64, passHash string) error {
	_, err := db.sql.Exec(`UPDATE admins SET pass_hash = ? WHERE id = ?`, passHash, adminID)
	return err
}

func (db *DB) SetPasswordOff(adminID int64, off bool) error {
	v := 0
	if off {
		v = 1
	}
	_, err := db.sql.Exec(`UPDATE admins SET password_off = ? WHERE id = ?`, v, adminID)
	return err
}

// CountAdminKeys is how many passkeys an admin has.
func (db *DB) CountAdminKeys(adminID int64) int {
	var n int
	_ = db.sql.QueryRow(`SELECT COUNT(*) FROM admin_keys WHERE admin_id = ?`, adminID).Scan(&n)
	return n
}

func (db *DB) AdminKeyByCredID(credID []byte) (*AdminKey, error) {
	k := &AdminKey{}
	err := db.sql.QueryRow(
		`SELECT id, admin_id, cred_id, public_key, sign_count, name, created_at, last_used
		   FROM admin_keys WHERE cred_id = ?`, credID).
		Scan(&k.ID, &k.AdminID, &k.CredID, &k.PublicKey, &k.SignCount, &k.Name, &k.CreatedAt, &k.LastUsed)
	if err == sql.ErrNoRows {
		return nil, nil
	}
	return k, err
}

func (db *DB) UseAdminKey(id int64, signCount uint32) error {
	_, err := db.sql.Exec(`UPDATE admin_keys SET sign_count = ?, last_used = ? WHERE id = ?`, signCount, time.Now().Unix(), id)
	return err
}

// DeleteAdminKey removes a passkey; with the last one gone, the password
// signs the admin in again.
func (db *DB) DeleteAdminKey(adminID, id int64) error {
	if _, err := db.sql.Exec(`DELETE FROM admin_keys WHERE admin_id = ? AND id = ?`, adminID, id); err != nil {
		return err
	}
	if db.CountAdminKeys(adminID) == 0 {
		return db.SetPasswordOff(adminID, false)
	}
	return nil
}

func (db *DB) DeleteAdminKeys(adminID int64) error {
	if _, err := db.sql.Exec(`DELETE FROM admin_keys WHERE admin_id = ?`, adminID); err != nil {
		return err
	}
	return db.SetPasswordOff(adminID, false)
}

// --- 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, err := scanAdmin(db.sql.QueryRow(
		`SELECT a.id, a.username, a.pass_hash, a.role, a.pw_locked
		   FROM admin_sessions s JOIN admins a ON a.id = s.admin_id
		  WHERE s.id = ? AND s.expires_at > ?`, sessionID, time.Now().Unix()))
	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
	Transport string // auto, ws or lp: how the agent is to reach us
	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, transport, 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, &a.Transport, &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) SetAgentTransport(id, transport string) error {
	_, err := db.sql.Exec(`UPDATE agents SET transport = ? WHERE id = ?`, transport, id)
	return 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)
	if err == nil {
		_, _ = db.sql.Exec(`DELETE FROM admin_agents WHERE agent_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
	Label      string // optional note shown next to the rule
}

func (db *DB) CreateForwardRule(r ForwardRule) (int64, error) {
	if r.Proto == "" {
		r.Proto = "http"
	}
	return db.sql.insert(
		`INSERT INTO forward_rules(agent_id, client_port, host_port, listen_addr, proto, no_ufw, label, created_at)
		 VALUES(?, ?, ?, ?, ?, ?, ?, ?)`,
		r.AgentID, r.ClientPort, r.HostPort, r.ListenAddr, r.Proto, boolInt(r.NoUfw), r.Label, time.Now().Unix())
}

func (db *DB) ListForwardRules() ([]ForwardRule, error) {
	rows, err := db.sql.Query(
		`SELECT id, agent_id, client_port, host_port, listen_addr, proto, no_ufw, label 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, &r.Label); 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, label FROM forward_rules WHERE id = ?`, id).
		Scan(&r.ID, &r.AgentID, &r.ClientPort, &r.HostPort, &r.ListenAddr, &r.Proto, &r.NoUfw, &r.Label)
	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, 0 for shared (the tunnel port), ufwManual for one added in Settings
}

// ufwManual is the RuleID of a ufw rule the admin added by hand: no forward
// owns it, so only the admin removes it.
const ufwManual = -1

// Manual reports whether the admin added the rule by hand.
func (r UfwRule) Manual() bool { return r.RuleID == ufwManual }

// SetUfwRuleManual hands a tracked rule over to the admin.
func (db *DB) SetUfwRuleManual(port int, proto string) error {
	_, err := db.sql.Exec(`UPDATE ufw_rules SET rule_id = ? WHERE port = ? AND proto = ?`, ufwManual, port, proto)
	return err
}

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())
}

func boolInt(b bool) int {
	if b {
		return 1
	}
	return 0
}