feat: initial commit

2026-10-13 00:25:29 +03:00
1 parent 353f62e257
commit c6852c6f1b
@@ -0,0 +1,452 @@
+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 '',
+    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`,
+	} {
+		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
+	LastSeen  time.Time
+}
+
+// Online reports whether the agent polled recently enough to be considered up.
+func (a Agent) Online() 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, enrolled_at, last_seen)
+		 VALUES(?, ?, ?, ?, ?, ?, ?, ?, ?)`,
+		a.ID, a.Name, a.Hostname, a.OS, a.Arch, a.TokenHash, a.Version, 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, last_seen FROM agents WHERE id = ?`, id).
+		Scan(&a.ID, &a.Name, &a.Hostname, &a.OS, &a.Arch, &a.TokenHash, &a.Version, &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, 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, &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
+}
+
+// TouchAgentVersion updates last_seen and the reported version together.
+func (db *DB) TouchAgentVersion(id, version string) error {
+	_, err := db.sql.Exec(`UPDATE agents SET last_seen = ?, version = ? WHERE id = ?`,
+		time.Now().Unix(), version, 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
+}
+
+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, created_at)
+		 VALUES(?, ?, ?, ?, ?, ?)`,
+		r.AgentID, r.ClientPort, r.HostPort, r.ListenAddr, r.Proto, 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 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); 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 FROM forward_rules WHERE id = ?`, id).
+		Scan(&r.ID, &r.AgentID, &r.ClientPort, &r.HostPort, &r.ListenAddr, &r.Proto)
+	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())
+}