185 lines
6.7 KiB
Go
185 lines
6.7 KiB
Go
package postgres
|
|
|
|
import (
|
|
"context"
|
|
"crypto/sha256"
|
|
"crypto/subtle"
|
|
"encoding/hex"
|
|
"errors"
|
|
"fmt"
|
|
"net"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/jackc/pgx/v5"
|
|
"github.com/jackc/pgx/v5/pgxpool"
|
|
)
|
|
|
|
var ErrUnauthorized = errors.New("unauthorized")
|
|
|
|
type Store struct{ Pool *pgxpool.Pool }
|
|
|
|
type Agent struct {
|
|
ID string `json:"id"`
|
|
Hostname string `json:"hostname"`
|
|
LastIP string `json:"last_ip"`
|
|
Enabled bool `json:"enabled"`
|
|
FirstSeen time.Time `json:"first_seen"`
|
|
LastSeen time.Time `json:"last_seen"`
|
|
}
|
|
type Detection struct {
|
|
ID int64 `json:"id"`
|
|
Fingerprint string `json:"fingerprint"`
|
|
RuleName string `json:"rule_name"`
|
|
Severity string `json:"severity"`
|
|
Status string `json:"status"`
|
|
Hostname string `json:"hostname"`
|
|
UserName string `json:"user_name"`
|
|
SourceIP string `json:"source_ip"`
|
|
Workstation string `json:"workstation"`
|
|
Summary string `json:"summary"`
|
|
EventCode uint32 `json:"event_code"`
|
|
Score float64 `json:"score"`
|
|
WindowStart time.Time `json:"window_start"`
|
|
WindowEnd time.Time `json:"window_end"`
|
|
FirstSeen time.Time `json:"first_seen"`
|
|
LastSeen time.Time `json:"last_seen"`
|
|
Count int64 `json:"count"`
|
|
}
|
|
|
|
func Open(ctx context.Context, url string) (*Store, error) {
|
|
p, e := pgxpool.New(ctx, url)
|
|
if e != nil {
|
|
return nil, e
|
|
}
|
|
cctx, cancel := context.WithTimeout(ctx, 5*time.Second)
|
|
defer cancel()
|
|
if e = p.Ping(cctx); e != nil {
|
|
p.Close()
|
|
return nil, e
|
|
}
|
|
return &Store{Pool: p}, nil
|
|
}
|
|
func (s *Store) Close() { s.Pool.Close() }
|
|
|
|
func (s *Store) AuthenticateOrEnroll(ctx context.Context, tenant, hostname, apiKey, enrollmentKey, expectedEnrollment, remoteIP string) (string, error) {
|
|
hostname = strings.TrimSpace(hostname)
|
|
apiKey = strings.TrimSpace(apiKey)
|
|
var id, hash string
|
|
var enabled bool
|
|
err := s.Pool.QueryRow(ctx, `SELECT id::text, api_key_hash, enabled FROM agents WHERE tenant_id=$1 AND hostname=$2`, tenant, hostname).Scan(&id, &hash, &enabled)
|
|
if err != nil {
|
|
if !errors.Is(err, pgx.ErrNoRows) {
|
|
return "", err
|
|
}
|
|
if enrollmentKey == "" || !secureEqual(hashHex(enrollmentKey), hashHex(expectedEnrollment)) {
|
|
return "", ErrUnauthorized
|
|
}
|
|
err = s.Pool.QueryRow(ctx, `INSERT INTO agents(tenant_id,hostname,api_key_hash,last_ip) VALUES($1,$2,$3,$4) ON CONFLICT(tenant_id,hostname) DO NOTHING RETURNING id::text`, tenant, hostname, hashHex(apiKey), remoteIP).Scan(&id)
|
|
if err == nil {
|
|
return id, nil
|
|
}
|
|
if !errors.Is(err, pgx.ErrNoRows) {
|
|
return "", err
|
|
}
|
|
// A concurrent first request may have enrolled the host between SELECT and INSERT.
|
|
err = s.Pool.QueryRow(ctx, `SELECT id::text, api_key_hash, enabled FROM agents WHERE tenant_id=$1 AND hostname=$2`, tenant, hostname).Scan(&id, &hash, &enabled)
|
|
if err != nil || !enabled || !secureEqual(strings.ToLower(hash), hashHex(apiKey)) {
|
|
return "", ErrUnauthorized
|
|
}
|
|
return id, nil
|
|
}
|
|
if !enabled || !secureEqual(strings.ToLower(hash), hashHex(apiKey)) {
|
|
return "", ErrUnauthorized
|
|
}
|
|
_, err = s.Pool.Exec(ctx, `UPDATE agents SET last_seen=now(), last_ip=$3 WHERE tenant_id=$1 AND hostname=$2`, tenant, hostname, remoteIP)
|
|
return id, err
|
|
}
|
|
|
|
func (s *Store) ListAgents(ctx context.Context, tenant string) ([]Agent, error) {
|
|
rows, e := s.Pool.Query(ctx, `SELECT id::text,hostname,enabled,COALESCE(last_ip,''),first_seen,last_seen FROM agents WHERE tenant_id=$1 ORDER BY last_seen DESC`, tenant)
|
|
if e != nil {
|
|
return nil, e
|
|
}
|
|
defer rows.Close()
|
|
var out []Agent
|
|
for rows.Next() {
|
|
var a Agent
|
|
if e := rows.Scan(&a.ID, &a.Hostname, &a.Enabled, &a.LastIP, &a.FirstSeen, &a.LastSeen); e != nil {
|
|
return nil, e
|
|
}
|
|
out = append(out, a)
|
|
}
|
|
return out, rows.Err()
|
|
}
|
|
|
|
func (s *Store) UpsertDetection(ctx context.Context, d Detection, tenant string) error {
|
|
_, e := s.Pool.Exec(ctx, `INSERT INTO detections(tenant_id,fingerprint,rule_name,severity,status,hostname,user_name,source_ip,workstation,event_code,score,window_start,window_end,summary,hit_count,first_seen,last_seen)
|
|
VALUES($1,$2,$3,$4,'open',$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$11,$12)
|
|
ON CONFLICT(tenant_id,fingerprint) DO UPDATE SET
|
|
first_seen=LEAST(detections.first_seen,EXCLUDED.first_seen),
|
|
last_seen=GREATEST(detections.last_seen,EXCLUDED.last_seen),
|
|
window_start=LEAST(detections.window_start,EXCLUDED.window_start),
|
|
window_end=GREATEST(detections.window_end,EXCLUDED.window_end),
|
|
hit_count=GREATEST(detections.hit_count,EXCLUDED.hit_count),
|
|
score=GREATEST(detections.score,EXCLUDED.score),
|
|
summary=EXCLUDED.summary, workstation=EXCLUDED.workstation, updated_at=now()`, tenant, d.Fingerprint, d.RuleName, d.Severity, d.Hostname, d.UserName, d.SourceIP, d.Workstation, d.EventCode, d.Score, d.WindowStart, d.WindowEnd, d.Summary, d.Count)
|
|
return e
|
|
}
|
|
|
|
func (s *Store) ListDetections(ctx context.Context, tenant string, limit int) ([]Detection, error) {
|
|
rows, e := s.Pool.Query(ctx, `SELECT id,fingerprint,rule_name,severity,status,hostname,user_name,source_ip,workstation,event_code,score,window_start,window_end,summary,hit_count,first_seen,last_seen FROM detections WHERE tenant_id=$1 ORDER BY last_seen DESC LIMIT $2`, tenant, limit)
|
|
if e != nil {
|
|
return nil, e
|
|
}
|
|
defer rows.Close()
|
|
var out []Detection
|
|
for rows.Next() {
|
|
var d Detection
|
|
if e := rows.Scan(&d.ID, &d.Fingerprint, &d.RuleName, &d.Severity, &d.Status, &d.Hostname, &d.UserName, &d.SourceIP, &d.Workstation, &d.EventCode, &d.Score, &d.WindowStart, &d.WindowEnd, &d.Summary, &d.Count, &d.FirstSeen, &d.LastSeen); e != nil {
|
|
return nil, e
|
|
}
|
|
out = append(out, d)
|
|
}
|
|
return out, rows.Err()
|
|
}
|
|
func (s *Store) UpdateDetectionStatus(ctx context.Context, tenant string, id int64, status string) error {
|
|
if status != "open" && status != "investigating" && status != "closed" && status != "false_positive" {
|
|
return fmt.Errorf("invalid status")
|
|
}
|
|
_, e := s.Pool.Exec(ctx, `UPDATE detections SET status=$3, updated_at=now() WHERE tenant_id=$1 AND id=$2`, tenant, id, status)
|
|
return e
|
|
}
|
|
func (s *Store) DetectionCounts(ctx context.Context, tenant string) (map[string]int64, error) {
|
|
rows, e := s.Pool.Query(ctx, `SELECT severity,count(*) FROM detections WHERE tenant_id=$1 AND status IN ('open','investigating') GROUP BY severity`, tenant)
|
|
if e != nil {
|
|
return nil, e
|
|
}
|
|
defer rows.Close()
|
|
out := map[string]int64{}
|
|
for rows.Next() {
|
|
var k string
|
|
var n int64
|
|
if e := rows.Scan(&k, &n); e != nil {
|
|
return nil, e
|
|
}
|
|
out[k] = n
|
|
}
|
|
return out, rows.Err()
|
|
}
|
|
|
|
func hashHex(v string) string { h := sha256.Sum256([]byte(v)); return hex.EncodeToString(h[:]) }
|
|
func secureEqual(a, b string) bool {
|
|
if len(a) != len(b) {
|
|
return false
|
|
}
|
|
return subtle.ConstantTimeCompare([]byte(a), []byte(b)) == 1
|
|
}
|
|
func ClientIP(remote string) string {
|
|
h, _, e := net.SplitHostPort(remote)
|
|
if e == nil {
|
|
return h
|
|
}
|
|
return remote
|
|
}
|