EvoBGP/internal/store/memory_firewall.go

package store

import (
	"encoding/hex"
	"sort"
	"strings"
	"time"

	"github.com/google/uuid"
)

func (m *Memory) ListFirewallClients(tenantID string) ([]*FirewallClient, error) {
	m.mu.RLock()
	defer m.mu.RUnlock()
	var out []*FirewallClient
	for _, rec := range m.firewallClients {
		if rec.TenantID == tenantID {
			out = append(out, firewallClientCopy(&rec.FirewallClient))
		}
	}
	sort.Slice(out, func(i, j int) bool { return out[i].CreatedAt.After(out[j].CreatedAt) })
	return out, nil
}

func (m *Memory) GetFirewallClient(tenantID, id string) (*FirewallClient, error) {
	m.mu.RLock()
	defer m.mu.RUnlock()
	rec, ok := m.firewallClients[id]
	if !ok || rec.TenantID != tenantID {
		return nil, ErrNotFound
	}
	return firewallClientCopy(&rec.FirewallClient), nil
}

func (m *Memory) CreateFirewallClient(tenantID string, in *FirewallClientCreate) (*FirewallClient, error) {
	if in == nil || strings.TrimSpace(in.Name) == "" || len(in.TokenHash) != 32 {
		return nil, ErrInvalidInput
	}
	m.mu.Lock()
	defer m.mu.Unlock()
	if _, ok := m.tenants[tenantID]; !ok {
		return nil, ErrTenantScope
	}
	hashKey := hex.EncodeToString(in.TokenHash)
	if _, dup := m.firewallHashIndex[hashKey]; dup {
		return nil, ErrInvalidInput
	}
	now := time.Now().UTC()
	id := uuid.NewString()
	rec := &firewallClientRec{
		FirewallClient: FirewallClient{
			ID:            id,
			TenantID:      tenantID,
			Name:          strings.TrimSpace(in.Name),
			Hostname:      strings.TrimSpace(in.Hostname),
			TokenPrefix:   in.TokenPrefix,
			Status:        "pending",
			ClientVersion: strings.TrimSpace(in.ClientVersion),
			CreatedAt:     now,
		},
		TokenHash: append([]byte(nil), in.TokenHash...),
	}
	m.firewallClients[id] = rec
	m.firewallHashIndex[hashKey] = id
	return firewallClientCopy(&rec.FirewallClient), nil
}

func (m *Memory) UpdateFirewallClient(tenantID, id string, patch *FirewallClientPatch) (*FirewallClient, error) {
	if patch == nil {
		return nil, ErrInvalidInput
	}
	m.mu.Lock()
	defer m.mu.Unlock()
	rec, ok := m.firewallClients[id]
	if !ok || rec.TenantID != tenantID {
		return nil, ErrNotFound
	}
	if patch.Name != nil {
		n := strings.TrimSpace(*patch.Name)
		if n == "" {
			return nil, ErrInvalidInput
		}
		rec.Name = n
	}
	if patch.Hostname != nil {
		rec.Hostname = strings.TrimSpace(*patch.Hostname)
	}
	return firewallClientCopy(&rec.FirewallClient), nil
}

func (m *Memory) ApproveFirewallClient(tenantID, id, approverAPIKeyID string) (*FirewallClient, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	rec, ok := m.firewallClients[id]
	if !ok || rec.TenantID != tenantID {
		return nil, ErrNotFound
	}
	if rec.Status == "revoked" {
		return nil, ErrInvalidInput
	}
	now := time.Now().UTC()
	rec.Status = "approved"
	rec.ApprovedAt = &now
	rec.ApprovedByAPIKeyID = strings.TrimSpace(approverAPIKeyID)
	rec.RevokedAt = nil
	return firewallClientCopy(&rec.FirewallClient), nil
}

func (m *Memory) RevokeFirewallClient(tenantID, id string) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	rec, ok := m.firewallClients[id]
	if !ok || rec.TenantID != tenantID {
		return ErrNotFound
	}
	now := time.Now().UTC()
	rec.Status = "revoked"
	rec.RevokedAt = &now
	return nil
}

func (m *Memory) DeleteFirewallClient(tenantID, id string) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	rec, ok := m.firewallClients[id]
	if !ok || rec.TenantID != tenantID {
		return ErrNotFound
	}
	hashKey := hex.EncodeToString(rec.TokenHash)
	delete(m.firewallHashIndex, hashKey)
	delete(m.firewallClients, id)
	for rid, rule := range m.firewallRules {
		if rule.ClientID != nil && *rule.ClientID == id {
			delete(m.firewallRules, rid)
		}
	}
	return nil
}

func (m *Memory) LookupFirewallClientByTokenHash(hash []byte) (*FirewallClient, error) {
	if len(hash) != 32 {
		return nil, ErrNotFound
	}
	m.mu.RLock()
	defer m.mu.RUnlock()
	id, ok := m.firewallHashIndex[hex.EncodeToString(hash)]
	if !ok {
		return nil, ErrNotFound
	}
	rec, ok := m.firewallClients[id]
	if !ok {
		return nil, ErrNotFound
	}
	return firewallClientCopy(&rec.FirewallClient), nil
}

func (m *Memory) TouchFirewallClientLastSeen(id, source, clientIP, clientVersion string) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	rec, ok := m.firewallClients[id]
	if !ok {
		return ErrNotFound
	}
	now := time.Now().UTC()
	rec.LastSeenAt = &now
	rec.LastSeenAtSource = strings.TrimSpace(source)
	rec.LastSeenIP = strings.TrimSpace(clientIP)
	if v := strings.TrimSpace(clientVersion); v != "" {
		rec.ClientVersion = v
	}
	return nil
}

func (m *Memory) TouchFirewallClientLastApply(id, source, status, errMsg string, prefixCount, ipCount int, packetsDropped, packetsAccepted int64) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	rec, ok := m.firewallClients[id]
	if !ok {
		return ErrNotFound
	}
	now := time.Now().UTC()
	rec.LastApplyAt = &now
	rec.LastApplySource = strings.TrimSpace(source)
	rec.LastApplyStatus = strings.TrimSpace(status)
	rec.LastApplyError = strings.TrimSpace(errMsg)
	rec.LastApplyPrefixCount = prefixCount
	rec.LastApplyIPCount = ipCount
	rec.LastApplyPacketsDropped = packetsDropped
	rec.LastApplyPacketsAccepted = packetsAccepted
	return nil
}

func (m *Memory) ListActiveFirewallClientHashes() ([]FirewallClientAuthRow, error) {
	m.mu.RLock()
	defer m.mu.RUnlock()
	var out []FirewallClientAuthRow
	for _, rec := range m.firewallClients {
		if rec.Status != "approved" {
			continue
		}
		out = append(out, FirewallClientAuthRow{
			ID:        rec.ID,
			TenantID:  rec.TenantID,
			TokenHash: append([]byte(nil), rec.TokenHash...),
		})
	}
	return out, nil
}

func (m *Memory) ListApprovedFirewallClientsForReplication(tenantID string) ([]FirewallClientReplicationRow, error) {
	m.mu.RLock()
	defer m.mu.RUnlock()
	var out []FirewallClientReplicationRow
	for _, rec := range m.firewallClients {
		if rec.TenantID != tenantID || rec.Status != "approved" {
			continue
		}
		out = append(out, FirewallClientReplicationRow{
			ClientID:     rec.ID,
			Name:         rec.Name,
			TokenHashHex: hex.EncodeToString(rec.TokenHash),
		})
	}
	sort.Slice(out, func(i, j int) bool { return out[i].ClientID < out[j].ClientID })
	return out, nil
}

func (m *Memory) ListFirewallRules(tenantID string, clientID *string) ([]*FirewallRule, error) {
	m.mu.RLock()
	defer m.mu.RUnlock()
	var out []*FirewallRule
	for _, rule := range m.firewallRules {
		if rule.TenantID != tenantID {
			continue
		}
		if clientID == nil {
			if rule.ClientID != nil {
				continue
			}
		} else if rule.ClientID == nil || *rule.ClientID != *clientID {
			continue
		}
		out = append(out, firewallRuleCopy(rule))
	}
	sort.Slice(out, func(i, j int) bool { return out[i].Priority < out[j].Priority })
	return out, nil
}

func (m *Memory) ListAllFirewallRulesForClient(tenantID, clientID string) ([]*FirewallRule, error) {
	tenantRules, err := m.ListFirewallRules(tenantID, nil)
	if err != nil {
		return nil, err
	}
	cid := clientID
	clientRules, err := m.ListFirewallRules(tenantID, &cid)
	if err != nil {
		return nil, err
	}
	out := make([]*FirewallRule, 0, len(tenantRules)+len(clientRules))
	out = append(out, clientRules...)
	out = append(out, tenantRules...)
	return out, nil
}

func (m *Memory) ListAllFirewallRulesForReplication(tenantID string) ([]*FirewallRule, error) {
	m.mu.RLock()
	defer m.mu.RUnlock()
	var out []*FirewallRule
	for _, rule := range m.firewallRules {
		if rule.TenantID == tenantID {
			out = append(out, firewallRuleCopy(rule))
		}
	}
	sort.Slice(out, func(i, j int) bool {
		ac, bc := "", ""
		if out[i].ClientID != nil {
			ac = *out[i].ClientID
		}
		if out[j].ClientID != nil {
			bc = *out[j].ClientID
		}
		if ac != bc {
			return ac < bc
		}
		return out[i].Priority < out[j].Priority
	})
	return out, nil
}

func (m *Memory) CreateFirewallRule(tenantID string, clientID *string, in *FirewallRuleCreate) (*FirewallRule, error) {
	if in == nil || !ValidFirewallAction(in.Action) {
		return nil, ErrInvalidInput
	}
	m.mu.Lock()
	defer m.mu.Unlock()
	if _, ok := m.tenants[tenantID]; !ok {
		return nil, ErrTenantScope
	}
	if clientID != nil {
		if rec, ok := m.firewallClients[*clientID]; !ok || rec.TenantID != tenantID {
			return nil, ErrNotFound
		}
	}
	priority := m.nextFirewallRulePriorityLocked(tenantID, clientID)
	if in.Priority != nil && *in.Priority >= 1 {
		priority = *in.Priority
	}
	if m.firewallRulePriorityTakenLocked(tenantID, clientID, priority, "") {
		return nil, ErrInvalidInput
	}
	now := time.Now().UTC()
	id := uuid.NewString()
	rule := &FirewallRule{
		ID:          id,
		TenantID:    tenantID,
		ClientID:    clientID,
		Priority:    priority,
		Action:      strings.ToLower(strings.TrimSpace(in.Action)),
		CommunityID: in.CommunityID,
		Comment:     strings.TrimSpace(in.Comment),
		CreatedAt:   now,
		UpdatedAt:   now,
	}
	m.firewallRules[id] = rule
	return firewallRuleCopy(rule), nil
}

func (m *Memory) UpdateFirewallRule(tenantID, ruleID string, patch *FirewallRulePatch) (*FirewallRule, error) {
	if patch == nil {
		return nil, ErrInvalidInput
	}
	m.mu.Lock()
	defer m.mu.Unlock()
	rule, ok := m.firewallRules[ruleID]
	if !ok || rule.TenantID != tenantID {
		return nil, ErrNotFound
	}
	if patch.Action != nil {
		if !ValidFirewallAction(*patch.Action) {
			return nil, ErrInvalidInput
		}
		rule.Action = strings.ToLower(strings.TrimSpace(*patch.Action))
	}
	if patch.ClearCommunity {
		rule.CommunityID = nil
	} else if patch.CommunityID != nil {
		rule.CommunityID = patch.CommunityID
	}
	if patch.Comment != nil {
		rule.Comment = strings.TrimSpace(*patch.Comment)
	}
	rule.UpdatedAt = time.Now().UTC()
	return firewallRuleCopy(rule), nil
}

func (m *Memory) DeleteFirewallRule(tenantID, ruleID string) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	rule, ok := m.firewallRules[ruleID]
	if !ok || rule.TenantID != tenantID {
		return ErrNotFound
	}
	delete(m.firewallRules, ruleID)
	return nil
}

func (m *Memory) ReorderFirewallRules(tenantID string, clientID *string, orderedIDs []string) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	scope := make(map[string]*FirewallRule)
	for id, rule := range m.firewallRules {
		if rule.TenantID != tenantID {
			continue
		}
		if clientID == nil {
			if rule.ClientID != nil {
				continue
			}
		} else if rule.ClientID == nil || *rule.ClientID != *clientID {
			continue
		}
		scope[id] = rule
	}
	if len(orderedIDs) != len(scope) {
		return ErrInvalidInput
	}
	now := time.Now().UTC()
	for i, id := range orderedIDs {
		rule, ok := scope[id]
		if !ok {
			return ErrInvalidInput
		}
		rule.Priority = i + 1
		rule.UpdatedAt = now
	}
	return nil
}

func (m *Memory) nextFirewallRulePriorityLocked(tenantID string, clientID *string) int {
	max := 0
	for _, rule := range m.firewallRules {
		if rule.TenantID != tenantID {
			continue
		}
		if clientID == nil {
			if rule.ClientID != nil {
				continue
			}
		} else if rule.ClientID == nil || *rule.ClientID != *clientID {
			continue
		}
		if rule.Priority > max {
			max = rule.Priority
		}
	}
	return max + 1
}

func (m *Memory) firewallRulePriorityTakenLocked(tenantID string, clientID *string, priority int, exceptID string) bool {
	for id, rule := range m.firewallRules {
		if id == exceptID || rule.TenantID != tenantID || rule.Priority != priority {
			continue
		}
		if clientID == nil {
			if rule.ClientID == nil {
				return true
			}
			continue
		}
		if rule.ClientID != nil && *rule.ClientID == *clientID {
			return true
		}
	}
	return false
}

func firewallClientCopy(c *FirewallClient) *FirewallClient {
	cp := *c
	cp.LastSeenAt = cloneTime(c.LastSeenAt)
	cp.LastApplyAt = cloneTime(c.LastApplyAt)
	cp.ApprovedAt = cloneTime(c.ApprovedAt)
	cp.RevokedAt = cloneTime(c.RevokedAt)
	return &cp
}

func firewallRuleCopy(r *FirewallRule) *FirewallRule {
	cp := *r
	if r.ClientID != nil {
		v := *r.ClientID
		cp.ClientID = &v
	}
	if r.CommunityID != nil {
		v := *r.CommunityID
		cp.CommunityID = &v
	}
	return &cp
}

func cloneTime(t *time.Time) *time.Time {
	if t == nil {
		return nil
	}
	v := *t
	return &v
}