Containarr/server/internal/store/store.go

package store

import (
	"database/sql"
	"time"

	_ "github.com/mattn/go-sqlite3"
)

type Agent struct {
	ID         string
	Token      string
	Hostname   string
	Alias      string
	IPAddress  string
	Arch       string
	OS         string
	LastSeenAt time.Time
	Online     bool
}

type Store struct {
	db *sql.DB
}

func New(path string) (*Store, error) {
	db, err := sql.Open("sqlite3", path+"?_journal_mode=WAL&_foreign_keys=on")
	if err != nil {
		return nil, err
	}
	s := &Store{db: db}
	return s, s.migrate()
}

func (s *Store) migrate() error {
	_, err := s.db.Exec(`
		CREATE TABLE IF NOT EXISTS users (
			username      TEXT PRIMARY KEY,
			password_hash TEXT NOT NULL
		);
		CREATE TABLE IF NOT EXISTS agents (
			id           TEXT PRIMARY KEY,
			token        TEXT UNIQUE NOT NULL,
			hostname     TEXT NOT NULL,
			alias        TEXT NOT NULL DEFAULT '',
			ip_address   TEXT NOT NULL DEFAULT '',
			arch         TEXT NOT NULL DEFAULT '',
			os           TEXT NOT NULL DEFAULT '',
			last_seen_at DATETIME,
			online       INTEGER NOT NULL DEFAULT 0
		);
	`)
	if err != nil {
		return err
	}
	// Idempotent — ignore error if column already exists.
	for _, col := range []string{
		`ALTER TABLE agents ADD COLUMN alias TEXT NOT NULL DEFAULT ''`,
		`ALTER TABLE agents ADD COLUMN ip_address TEXT NOT NULL DEFAULT ''`,
	} {
		s.db.Exec(col)
	}
	return nil
}

func (s *Store) Close() error { return s.db.Close() }

func (s *Store) UpsertAgent(a *Agent) error {
	_, err := s.db.Exec(`
		INSERT INTO agents (id, token, hostname, alias, ip_address, arch, os, last_seen_at, online)
		VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
		ON CONFLICT(token) DO UPDATE SET
			hostname     = excluded.hostname,
			ip_address   = excluded.ip_address,
			arch         = excluded.arch,
			os           = excluded.os,
			last_seen_at = excluded.last_seen_at,
			online       = excluded.online
	`, a.ID, a.Token, a.Hostname, a.Alias, a.IPAddress, a.Arch, a.OS, a.LastSeenAt, boolToInt(a.Online))
	return err
}

func (s *Store) AgentByToken(token string) (*Agent, error) {
	row := s.db.QueryRow(`
		SELECT id, token, hostname, alias, ip_address, arch, os, last_seen_at, online
		FROM agents WHERE token = ?`, token)
	return scanAgent(row)
}

func (s *Store) GetAgent(id string) (*Agent, error) {
	row := s.db.QueryRow(`
		SELECT id, token, hostname, alias, ip_address, arch, os, last_seen_at, online
		FROM agents WHERE id = ?`, id)
	return scanAgent(row)
}

func (s *Store) ListAgents() ([]*Agent, error) {
	rows, err := s.db.Query(`
		SELECT id, token, hostname, alias, ip_address, arch, os, last_seen_at, online
		FROM agents ORDER BY hostname`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var agents []*Agent
	for rows.Next() {
		a := &Agent{}
		var online int
		var lastSeen sql.NullTime
		if err := rows.Scan(&a.ID, &a.Token, &a.Hostname, &a.Alias, &a.IPAddress, &a.Arch, &a.OS, &lastSeen, &online); err != nil {
			return nil, err
		}
		if lastSeen.Valid {
			a.LastSeenAt = lastSeen.Time
		}
		a.Online = online == 1
		agents = append(agents, a)
	}
	return agents, rows.Err()
}

func (s *Store) SetAgentOffline(id string) error {
	_, err := s.db.Exec(`UPDATE agents SET online = 0 WHERE id = ?`, id)
	return err
}

func (s *Store) CreateAgentToken(id, token, hostname string) error {
	_, err := s.db.Exec(`
		INSERT OR IGNORE INTO agents (id, token, hostname, arch, os, online)
		VALUES (?, ?, ?, '', '', 0)
	`, id, token, hostname)
	return err
}

func (s *Store) UpdateAgentAlias(id, alias string) error {
	_, err := s.db.Exec(`UPDATE agents SET alias = ? WHERE id = ?`, alias, id)
	return err
}

// ── Users ─────────────────────────────────────────────────────────────────────

func (s *Store) GetUserHash(username string) (string, error) {
	var hash string
	err := s.db.QueryRow(`SELECT password_hash FROM users WHERE username = ?`, username).Scan(&hash)
	return hash, err
}

func (s *Store) UpsertUser(username, hash string) error {
	_, err := s.db.Exec(`
		INSERT INTO users (username, password_hash) VALUES (?, ?)
		ON CONFLICT(username) DO UPDATE SET password_hash = excluded.password_hash
	`, username, hash)
	return err
}

func (s *Store) UserExists(username string) (bool, error) {
	var n int
	err := s.db.QueryRow(`SELECT COUNT(*) FROM users WHERE username = ?`, username).Scan(&n)
	return n > 0, err
}

func scanAgent(row *sql.Row) (*Agent, error) {
	a := &Agent{}
	var online int
	var lastSeen sql.NullTime
	err := row.Scan(&a.ID, &a.Token, &a.Hostname, &a.Alias, &a.IPAddress, &a.Arch, &a.OS, &lastSeen, &online)
	if err != nil {
		return nil, err
	}
	if lastSeen.Valid {
		a.LastSeenAt = lastSeen.Time
	}
	a.Online = online == 1
	return a, nil
}

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