EvoBGP/internal/repository/postgres_firewall.go

package repository

import (
	"context"
	"encoding/hex"
	"errors"
	"strings"
	"time"

	"evobgp/internal/store"

	"github.com/google/uuid"
	"github.com/jackc/pgx/v5"
)

const firewallClientSelectCols = `
	id, name, COALESCE(hostname, ''), token_prefix, status,
	last_seen_at, COALESCE(last_seen_at_source, ''), COALESCE(last_seen_ip, ''),
	last_apply_at, COALESCE(last_apply_status, ''), COALESCE(last_apply_error, ''),
	COALESCE(last_apply_prefix_count, 0), COALESCE(last_apply_ip_count, 0),
	COALESCE(last_apply_packets_dropped, 0), COALESCE(last_apply_packets_accepted, 0),
	COALESCE(last_apply_source, ''),
	COALESCE(client_version, ''), created_at, approved_at, approved_by_api_key_id, revoked_at, created_by_user_id`

func (p *Postgres) ListFirewallClients(tenantID string) ([]*store.FirewallClient, error) {
	ctx := context.Background()
	rows, err := p.pool.Query(ctx, `
		SELECT `+firewallClientSelectCols+`
		FROM firewall_client WHERE tenant_id=$1 ORDER BY created_at DESC`, tenantID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []*store.FirewallClient
	for rows.Next() {
		c, err := scanFirewallClientRow(rows.Scan, tenantID)
		if err != nil {
			return nil, err
		}
		out = append(out, c)
	}
	return out, rows.Err()
}

func (p *Postgres) GetFirewallClient(tenantID, id string) (*store.FirewallClient, error) {
	ctx := context.Background()
	row := p.pool.QueryRow(ctx, `
		SELECT `+firewallClientSelectCols+`
		FROM firewall_client WHERE id=$1 AND tenant_id=$2`, id, tenantID)
	c, err := scanFirewallClientRow(row.Scan, tenantID)
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return nil, store.ErrNotFound
		}
		return nil, err
	}
	return c, nil
}

func (p *Postgres) CreateFirewallClient(tenantID string, in *store.FirewallClientCreate) (*store.FirewallClient, error) {
	if in == nil || strings.TrimSpace(in.Name) == "" || len(in.TokenHash) != 32 {
		return nil, store.ErrInvalidInput
	}
	id := uuid.NewString()
	ctx := context.Background()
	var createdBy any
	if v := strings.TrimSpace(in.CreatedByUserID); v != "" {
		createdBy = v
	}
	_, err := p.pool.Exec(ctx, `
		INSERT INTO firewall_client (id, tenant_id, name, hostname, token_prefix, token_hash, client_version, created_by_user_id)
		VALUES ($1,$2,$3,$4,$5,$6,$7,$8)`,
		id, tenantID, strings.TrimSpace(in.Name), strings.TrimSpace(in.Hostname),
		in.TokenPrefix, in.TokenHash, strings.TrimSpace(in.ClientVersion), createdBy)
	if err != nil {
		return nil, err
	}
	return p.GetFirewallClient(tenantID, id)
}

func (p *Postgres) UpdateFirewallClient(tenantID, id string, patch *store.FirewallClientPatch) (*store.FirewallClient, error) {
	cur, err := p.GetFirewallClient(tenantID, id)
	if err != nil {
		return nil, err
	}
	if patch == nil {
		return nil, store.ErrInvalidInput
	}
	name := cur.Name
	hostname := cur.Hostname
	if patch.Name != nil {
		name = strings.TrimSpace(*patch.Name)
		if name == "" {
			return nil, store.ErrInvalidInput
		}
	}
	if patch.Hostname != nil {
		hostname = strings.TrimSpace(*patch.Hostname)
	}
	ctx := context.Background()
	_, err = p.pool.Exec(ctx, `UPDATE firewall_client SET name=$3, hostname=$4 WHERE id=$1 AND tenant_id=$2`,
		id, tenantID, name, hostname)
	if err != nil {
		return nil, err
	}
	return p.GetFirewallClient(tenantID, id)
}

func (p *Postgres) ApproveFirewallClient(tenantID, id, approverAPIKeyID string) (*store.FirewallClient, error) {
	ctx := context.Background()
	tag, err := p.pool.Exec(ctx, `
		UPDATE firewall_client
		SET status='approved', approved_at=now(), approved_by_api_key_id=$3, revoked_at=NULL
		WHERE id=$1 AND tenant_id=$2 AND status != 'revoked'`, id, tenantID, nullIfEmpty(approverAPIKeyID))
	if err != nil {
		return nil, err
	}
	if tag.RowsAffected() == 0 {
		return nil, store.ErrNotFound
	}
	return p.GetFirewallClient(tenantID, id)
}

func (p *Postgres) RevokeFirewallClient(tenantID, id string) error {
	ctx := context.Background()
	tag, err := p.pool.Exec(ctx, `
		UPDATE firewall_client SET status='revoked', revoked_at=now() WHERE id=$1 AND tenant_id=$2`, id, tenantID)
	if err != nil {
		return err
	}
	if tag.RowsAffected() == 0 {
		return store.ErrNotFound
	}
	return nil
}

func (p *Postgres) DeleteFirewallClient(tenantID, id string) error {
	ctx := context.Background()
	tag, err := p.pool.Exec(ctx, `DELETE FROM firewall_client WHERE id=$1 AND tenant_id=$2`, id, tenantID)
	if err != nil {
		return err
	}
	if tag.RowsAffected() == 0 {
		return store.ErrNotFound
	}
	return nil
}

func (p *Postgres) LookupFirewallClientByTokenHash(hash []byte) (*store.FirewallClient, error) {
	if len(hash) != 32 {
		return nil, store.ErrNotFound
	}
	ctx := context.Background()
	row := p.pool.QueryRow(ctx, `
		SELECT tenant_id, `+firewallClientSelectCols+`
		FROM firewall_client WHERE token_hash=$1`, hash)
	c, err := scanFirewallClientLookupRow(row.Scan)
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return nil, store.ErrNotFound
		}
		return nil, err
	}
	return c, nil
}

func (p *Postgres) TouchFirewallClientLastSeen(id, source, clientIP, clientVersion string) error {
	ctx := context.Background()
	_, err := p.pool.Exec(ctx, `
		UPDATE firewall_client SET last_seen_at=now(), last_seen_at_source=$2, last_seen_ip=$3,
			client_version=COALESCE(NULLIF($4,''), client_version)
		WHERE id=$1`, id, strings.TrimSpace(source), strings.TrimSpace(clientIP), strings.TrimSpace(clientVersion))
	return err
}

func (p *Postgres) TouchFirewallClientLastApply(id, source, status, errMsg string, prefixCount, ipCount int, packetsDropped, packetsAccepted int64) error {
	ctx := context.Background()
	_, err := p.pool.Exec(ctx, `
		UPDATE firewall_client SET last_apply_at=now(), last_apply_source=$2, last_apply_status=$3,
			last_apply_error=$4, last_apply_prefix_count=$5, last_apply_ip_count=$6,
			last_apply_packets_dropped=$7, last_apply_packets_accepted=$8
		WHERE id=$1`, id, strings.TrimSpace(source), strings.TrimSpace(status), strings.TrimSpace(errMsg), prefixCount, ipCount, packetsDropped, packetsAccepted)
	return err
}

func (p *Postgres) ListActiveFirewallClientHashes() ([]store.FirewallClientAuthRow, error) {
	ctx := context.Background()
	rows, err := p.pool.Query(ctx, `
		SELECT id, tenant_id, token_hash FROM firewall_client WHERE status='approved'`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []store.FirewallClientAuthRow
	for rows.Next() {
		var row store.FirewallClientAuthRow
		if err := rows.Scan(&row.ID, &row.TenantID, &row.TokenHash); err != nil {
			return nil, err
		}
		if len(row.TokenHash) != 32 {
			continue
		}
		out = append(out, row)
	}
	return out, rows.Err()
}

func (p *Postgres) ListApprovedFirewallClientsForReplication(tenantID string) ([]store.FirewallClientReplicationRow, error) {
	ctx := context.Background()
	rows, err := p.pool.Query(ctx, `
		SELECT id, name, token_hash FROM firewall_client
		WHERE tenant_id=$1 AND status='approved' ORDER BY id`, tenantID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []store.FirewallClientReplicationRow
	for rows.Next() {
		var id, name string
		var hash []byte
		if err := rows.Scan(&id, &name, &hash); err != nil {
			return nil, err
		}
		out = append(out, store.FirewallClientReplicationRow{
			ClientID:     id,
			Name:         name,
			TokenHashHex: hex.EncodeToString(hash),
		})
	}
	return out, rows.Err()
}

func (p *Postgres) ListFirewallRules(tenantID string, clientID *string) ([]*store.FirewallRule, error) {
	ctx := context.Background()
	var rows pgx.Rows
	var err error
	if clientID == nil {
		rows, err = p.pool.Query(ctx, `
			SELECT id, client_id, priority, action, community_id, comment, created_at, updated_at, created_by_user_id
			FROM firewall_rule WHERE tenant_id=$1 AND client_id IS NULL ORDER BY priority`, tenantID)
	} else {
		rows, err = p.pool.Query(ctx, `
			SELECT id, client_id, priority, action, community_id, comment, created_at, updated_at, created_by_user_id
			FROM firewall_rule WHERE tenant_id=$1 AND client_id=$2 ORDER BY priority`, tenantID, *clientID)
	}
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	return scanFirewallRules(rows, tenantID)
}

func (p *Postgres) ListAllFirewallRulesForClient(tenantID, clientID string) ([]*store.FirewallRule, error) {
	ctx := context.Background()
	rows, err := p.pool.Query(ctx, `
		SELECT id, client_id, priority, action, community_id, comment, created_at, updated_at, created_by_user_id
		FROM firewall_rule
		WHERE tenant_id=$1 AND (client_id IS NULL OR client_id=$2)
		ORDER BY CASE WHEN client_id IS NULL THEN 1 ELSE 0 END, priority`, tenantID, clientID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	return scanFirewallRules(rows, tenantID)
}

func (p *Postgres) ListAllFirewallRulesForReplication(tenantID string) ([]*store.FirewallRule, error) {
	ctx := context.Background()
	rows, err := p.pool.Query(ctx, `
		SELECT id, client_id, priority, action, community_id, comment, created_at, updated_at, created_by_user_id
		FROM firewall_rule WHERE tenant_id=$1
		ORDER BY CASE WHEN client_id IS NULL THEN 1 ELSE 0 END, client_id, priority`, tenantID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	return scanFirewallRules(rows, tenantID)
}

func (p *Postgres) CreateFirewallRule(tenantID string, clientID *string, in *store.FirewallRuleCreate) (*store.FirewallRule, error) {
	if in == nil || !store.ValidFirewallAction(in.Action) {
		return nil, store.ErrInvalidInput
	}
	if clientID != nil {
		if _, err := p.GetFirewallClient(tenantID, *clientID); err != nil {
			return nil, err
		}
	}
	priority := 1
	if in.Priority != nil && *in.Priority >= 1 {
		priority = *in.Priority
	} else {
		next, err := p.nextFirewallRulePriority(tenantID, clientID)
		if err != nil {
			return nil, err
		}
		priority = next
	}
	id := uuid.NewString()
	ctx := context.Background()
	var createdBy any
	if v := strings.TrimSpace(in.CreatedByUserID); v != "" {
		createdBy = v
	}
	_, err := p.pool.Exec(ctx, `
		INSERT INTO firewall_rule (id, tenant_id, client_id, priority, action, community_id, comment, created_by_user_id)
		VALUES ($1,$2,$3,$4,$5,$6,$7,$8)`,
		id, tenantID, clientID, priority, strings.ToLower(strings.TrimSpace(in.Action)), in.CommunityID, strings.TrimSpace(in.Comment), createdBy)
	if err != nil {
		return nil, err
	}
	return p.GetFirewallRule(tenantID, id)
}

func (p *Postgres) GetFirewallRule(tenantID, id string) (*store.FirewallRule, error) {
	ctx := context.Background()
	row := p.pool.QueryRow(ctx, `
		SELECT id, client_id, priority, action, community_id, comment, created_at, updated_at, created_by_user_id
		FROM firewall_rule WHERE id=$1 AND tenant_id=$2`, id, tenantID)
	r, err := scanFirewallRuleRow(row.Scan, tenantID)
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return nil, store.ErrNotFound
		}
		return nil, err
	}
	return r, nil
}

func (p *Postgres) UpdateFirewallRule(tenantID, ruleID string, patch *store.FirewallRulePatch) (*store.FirewallRule, error) {
	cur, err := p.GetFirewallRule(tenantID, ruleID)
	if err != nil {
		return nil, err
	}
	if patch == nil {
		return nil, store.ErrInvalidInput
	}
	action := cur.Action
	communityID := cur.CommunityID
	comment := cur.Comment
	if patch.Action != nil {
		if !store.ValidFirewallAction(*patch.Action) {
			return nil, store.ErrInvalidInput
		}
		action = strings.ToLower(strings.TrimSpace(*patch.Action))
	}
	if patch.ClearCommunity {
		communityID = nil
	} else if patch.CommunityID != nil {
		communityID = patch.CommunityID
	}
	if patch.Comment != nil {
		comment = strings.TrimSpace(*patch.Comment)
	}
	ctx := context.Background()
	_, err = p.pool.Exec(ctx, `
		UPDATE firewall_rule SET action=$3, community_id=$4, comment=$5, updated_at=now()
		WHERE id=$1 AND tenant_id=$2`, ruleID, tenantID, action, communityID, comment)
	if err != nil {
		return nil, err
	}
	return p.GetFirewallRule(tenantID, ruleID)
}

func (p *Postgres) DeleteFirewallRule(tenantID, ruleID string) error {
	ctx := context.Background()
	tag, err := p.pool.Exec(ctx, `DELETE FROM firewall_rule WHERE id=$1 AND tenant_id=$2`, ruleID, tenantID)
	if err != nil {
		return err
	}
	if tag.RowsAffected() == 0 {
		return store.ErrNotFound
	}
	return nil
}

func (p *Postgres) ReorderFirewallRules(tenantID string, clientID *string, orderedIDs []string) error {
	ctx := context.Background()
	tx, err := p.pool.Begin(ctx)
	if err != nil {
		return err
	}
	defer func() { _ = tx.Rollback(ctx) }()
	for i, id := range orderedIDs {
		var execErr error
		if clientID == nil {
			_, execErr = tx.Exec(ctx, `
				UPDATE firewall_rule SET priority=$3, updated_at=now()
				WHERE id=$1 AND tenant_id=$2 AND client_id IS NULL`, id, tenantID, i+1)
		} else {
			_, execErr = tx.Exec(ctx, `
				UPDATE firewall_rule SET priority=$4, updated_at=now()
				WHERE id=$1 AND tenant_id=$2 AND client_id=$3`, id, tenantID, *clientID, i+1)
		}
		if execErr != nil {
			return execErr
		}
	}
	return tx.Commit(ctx)
}

func (p *Postgres) nextFirewallRulePriority(tenantID string, clientID *string) (int, error) {
	ctx := context.Background()
	var max int
	var err error
	if clientID == nil {
		err = p.pool.QueryRow(ctx, `
			SELECT COALESCE(MAX(priority),0) FROM firewall_rule WHERE tenant_id=$1 AND client_id IS NULL`, tenantID).Scan(&max)
	} else {
		err = p.pool.QueryRow(ctx, `
			SELECT COALESCE(MAX(priority),0) FROM firewall_rule WHERE tenant_id=$1 AND client_id=$2`, tenantID, *clientID).Scan(&max)
	}
	if err != nil {
		return 0, err
	}
	return max + 1, nil
}

func scanFirewallRules(rows pgx.Rows, tenantID string) ([]*store.FirewallRule, error) {
	var out []*store.FirewallRule
	for rows.Next() {
		r, err := scanFirewallRuleRow(rows.Scan, tenantID)
		if err != nil {
			return nil, err
		}
		out = append(out, r)
	}
	return out, rows.Err()
}

func scanFirewallRuleRow(scan scanFn, tenantID string) (*store.FirewallRule, error) {
	var r store.FirewallRule
	r.TenantID = tenantID
	var clientID, communityID, createdBy *string
	if err := scan(&r.ID, &clientID, &r.Priority, &r.Action, &communityID, &r.Comment, &r.CreatedAt, &r.UpdatedAt, &createdBy); err != nil {
		return nil, err
	}
	r.ClientID = clientID
	r.CommunityID = communityID
	if createdBy != nil {
		r.CreatedByUserID = strings.TrimSpace(*createdBy)
	}
	return &r, nil
}

func scanFirewallClientRow(scan scanFn, tenantID string) (*store.FirewallClient, error) {
	var c store.FirewallClient
	c.TenantID = tenantID
	var approvedBy, createdBy *string
	var lastSeen, lastApply, approved, revoked *time.Time
	var prefixCount, ipCount *int
	var packetsDropped, packetsAccepted *int64
	if err := scan(
		&c.ID, &c.Name, &c.Hostname, &c.TokenPrefix, &c.Status,
		&lastSeen, &c.LastSeenAtSource, &c.LastSeenIP,
		&lastApply, &c.LastApplyStatus, &c.LastApplyError,
		&prefixCount, &ipCount, &packetsDropped, &packetsAccepted, &c.LastApplySource,
		&c.ClientVersion, &c.CreatedAt, &approved, &approvedBy, &revoked, &createdBy,
	); err != nil {
		return nil, err
	}
	return finishFirewallClientScan(&c, lastSeen, lastApply, approved, revoked, approvedBy, createdBy, prefixCount, ipCount, packetsDropped, packetsAccepted), nil
}

func scanFirewallClientLookupRow(scan scanFn) (*store.FirewallClient, error) {
	var c store.FirewallClient
	var approvedBy, createdBy *string
	var lastSeen, lastApply, approved, revoked *time.Time
	var prefixCount, ipCount *int
	var packetsDropped, packetsAccepted *int64
	if err := scan(
		&c.TenantID, &c.ID, &c.Name, &c.Hostname, &c.TokenPrefix, &c.Status,
		&lastSeen, &c.LastSeenAtSource, &c.LastSeenIP,
		&lastApply, &c.LastApplyStatus, &c.LastApplyError,
		&prefixCount, &ipCount, &packetsDropped, &packetsAccepted, &c.LastApplySource,
		&c.ClientVersion, &c.CreatedAt, &approved, &approvedBy, &revoked, &createdBy,
	); err != nil {
		return nil, err
	}
	return finishFirewallClientScan(&c, lastSeen, lastApply, approved, revoked, approvedBy, createdBy, prefixCount, ipCount, packetsDropped, packetsAccepted), nil
}

func finishFirewallClientScan(c *store.FirewallClient, lastSeen, lastApply, approved, revoked *time.Time, approvedBy, createdBy *string, prefixCount, ipCount *int, packetsDropped, packetsAccepted *int64) *store.FirewallClient {
	c.LastSeenAt = lastSeen
	c.LastApplyAt = lastApply
	c.ApprovedAt = approved
	c.RevokedAt = revoked
	if approvedBy != nil {
		c.ApprovedByAPIKeyID = *approvedBy
	}
	if createdBy != nil {
		c.CreatedByUserID = strings.TrimSpace(*createdBy)
	}
	if prefixCount != nil {
		c.LastApplyPrefixCount = *prefixCount
	}
	if ipCount != nil {
		c.LastApplyIPCount = *ipCount
	}
	if packetsDropped != nil {
		c.LastApplyPacketsDropped = *packetsDropped
	}
	if packetsAccepted != nil {
		c.LastApplyPacketsAccepted = *packetsAccepted
	}
	return c
}

func nullIfEmpty(s string) any {
	if strings.TrimSpace(s) == "" {
		return nil
	}
	return s
}