elmprodvpn/selective-vpn-api/app/trafficprofiles/store.go

package trafficprofiles

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync"
	"time"
)

type Profile struct {
	ID         string
	Name       string
	AppKey     string
	Command    string
	Target     string
	TTLSec     int
	VPNProfile string
	CreatedAt  string
	UpdatedAt  string
}

type UpsertRequest struct {
	ID         string
	Name       string
	AppKey     string
	Command    string
	Target     string
	TTLSec     int
	VPNProfile string
}

type Deps struct {
	CanonicalizeAppKey func(appKey, command string) string
	SanitizeID         func(string) string
	DefaultTTLSec      int
}

type Store struct {
	mu sync.Mutex

	statePath          string
	canonicalizeAppKey func(appKey, command string) string
	sanitizeID         func(string) string
	defaultTTLSec      int
}

type state struct {
	Version   int       `json:"version"`
	UpdatedAt string    `json:"updated_at"`
	Profiles  []Profile `json:"profiles,omitempty"`
}

func NewStore(statePath string, deps Deps) *Store {
	canon := deps.CanonicalizeAppKey
	if canon == nil {
		canon = func(appKey, _ string) string { return strings.TrimSpace(appKey) }
	}
	sanitize := deps.SanitizeID
	if sanitize == nil {
		sanitize = defaultSanitizeID
	}
	return &Store{
		statePath:          strings.TrimSpace(statePath),
		canonicalizeAppKey: canon,
		sanitizeID:         sanitize,
		defaultTTLSec:      deps.DefaultTTLSec,
	}
}

func (s *Store) List() []Profile {
	s.mu.Lock()
	defer s.mu.Unlock()

	st := s.loadStateLocked()
	out := append([]Profile(nil), st.Profiles...)
	sort.Slice(out, func(i, j int) bool {
		return out[i].UpdatedAt > out[j].UpdatedAt
	})
	return out
}

func (s *Store) Upsert(req UpsertRequest) (Profile, error) {
	s.mu.Lock()
	defer s.mu.Unlock()

	st := s.loadStateLocked()

	target := strings.ToLower(strings.TrimSpace(req.Target))
	if target == "" {
		target = "vpn"
	}
	if target != "vpn" && target != "direct" {
		return Profile{}, fmt.Errorf("target must be vpn|direct")
	}

	cmd := strings.TrimSpace(req.Command)
	if cmd == "" {
		return Profile{}, fmt.Errorf("missing command")
	}

	appKey := strings.TrimSpace(req.AppKey)
	if appKey == "" {
		fields := strings.Fields(cmd)
		if len(fields) > 0 {
			appKey = strings.TrimSpace(fields[0])
		}
	}
	appKey = s.canonicalizeAppKey(appKey, cmd)
	if appKey == "" {
		return Profile{}, fmt.Errorf("cannot infer app_key")
	}

	id := strings.TrimSpace(req.ID)
	if id == "" {
		for _, p := range st.Profiles {
			if strings.TrimSpace(p.AppKey) == appKey && strings.ToLower(strings.TrimSpace(p.Target)) == target {
				id = strings.TrimSpace(p.ID)
				break
			}
		}
	}
	if id == "" {
		id = deriveProfileID(appKey, target, st.Profiles, s.sanitizeID)
	}
	if id == "" {
		return Profile{}, fmt.Errorf("cannot derive profile id")
	}

	name := strings.TrimSpace(req.Name)
	if name == "" {
		name = filepath.Base(appKey)
		if name == "" || name == "/" || name == "." {
			name = id
		}
	}

	ttl := req.TTLSec
	if ttl <= 0 {
		ttl = s.defaultTTLSec
	}

	vpnProfile := strings.TrimSpace(req.VPNProfile)
	now := time.Now().UTC().Format(time.RFC3339)

	prof := Profile{
		ID:         id,
		Name:       name,
		AppKey:     appKey,
		Command:    cmd,
		Target:     target,
		TTLSec:     ttl,
		VPNProfile: vpnProfile,
		UpdatedAt:  now,
	}

	updated := false
	for i := range st.Profiles {
		if strings.TrimSpace(st.Profiles[i].ID) != id {
			continue
		}
		prof.CreatedAt = strings.TrimSpace(st.Profiles[i].CreatedAt)
		if prof.CreatedAt == "" {
			prof.CreatedAt = now
		}
		st.Profiles[i] = prof
		updated = true
		break
	}
	if !updated {
		prof.CreatedAt = now
		st.Profiles = append(st.Profiles, prof)
	}

	if err := s.saveStateLocked(st); err != nil {
		return Profile{}, err
	}
	return prof, nil
}

func (s *Store) Delete(id string) (bool, string) {
	s.mu.Lock()
	defer s.mu.Unlock()

	id = strings.TrimSpace(id)
	if id == "" {
		return false, "empty id"
	}

	st := s.loadStateLocked()
	kept := st.Profiles[:0]
	found := false
	for _, p := range st.Profiles {
		if strings.TrimSpace(p.ID) == id {
			found = true
			continue
		}
		kept = append(kept, p)
	}
	st.Profiles = kept

	if !found {
		return true, "not found"
	}
	if err := s.saveStateLocked(st); err != nil {
		return false, err.Error()
	}
	return true, "deleted"
}

func Dedupe(in []Profile, canonicalize func(appKey, command string) string) ([]Profile, bool) {
	if canonicalize == nil {
		canonicalize = func(appKey, _ string) string { return strings.TrimSpace(appKey) }
	}
	if len(in) <= 1 {
		return in, false
	}

	out := make([]Profile, 0, len(in))
	byID := map[string]int{}
	byAppTarget := map[string]int{}
	changed := false

	for _, raw := range in {
		p := raw
		p.ID = strings.TrimSpace(p.ID)
		p.Target = strings.ToLower(strings.TrimSpace(p.Target))
		p.AppKey = canonicalize(p.AppKey, p.Command)

		if p.ID == "" {
			changed = true
			continue
		}
		if p.Target != "vpn" && p.Target != "direct" {
			p.Target = "vpn"
			changed = true
		}

		if idx, ok := byID[p.ID]; ok {
			if preferProfile(p, out[idx]) {
				out[idx] = p
			}
			changed = true
			continue
		}

		if p.AppKey != "" {
			key := p.Target + "|" + p.AppKey
			if idx, ok := byAppTarget[key]; ok {
				if preferProfile(p, out[idx]) {
					byID[p.ID] = idx
					out[idx] = p
				}
				changed = true
				continue
			}
			byAppTarget[key] = len(out)
		}

		byID[p.ID] = len(out)
		out = append(out, p)
	}
	return out, changed
}

func (s *Store) loadStateLocked() state {
	st := state{Version: 1}
	if strings.TrimSpace(s.statePath) == "" {
		return st
	}
	data, err := os.ReadFile(s.statePath)
	if err != nil {
		return st
	}
	if err := json.Unmarshal(data, &st); err != nil {
		return state{Version: 1}
	}
	if st.Version == 0 {
		st.Version = 1
	}
	if st.Profiles == nil {
		st.Profiles = nil
	}

	changed := false
	for i := range st.Profiles {
		canon := s.canonicalizeAppKey(st.Profiles[i].AppKey, st.Profiles[i].Command)
		if canon != "" && strings.TrimSpace(st.Profiles[i].AppKey) != canon {
			st.Profiles[i].AppKey = canon
			changed = true
		}
		st.Profiles[i].Target = strings.ToLower(strings.TrimSpace(st.Profiles[i].Target))
	}
	if deduped, dedupChanged := Dedupe(st.Profiles, s.canonicalizeAppKey); dedupChanged {
		st.Profiles = deduped
		changed = true
	}
	if changed {
		_ = s.saveStateLocked(st)
	}
	return st
}

func (s *Store) saveStateLocked(st state) error {
	st.Version = 1
	st.UpdatedAt = time.Now().UTC().Format(time.RFC3339)

	data, err := json.MarshalIndent(st, "", "  ")
	if err != nil {
		return err
	}
	if strings.TrimSpace(s.statePath) == "" {
		return fmt.Errorf("state path is empty")
	}
	if err := os.MkdirAll(filepath.Dir(s.statePath), 0o755); err != nil {
		return err
	}
	tmp := s.statePath + ".tmp"
	if err := os.WriteFile(tmp, data, 0o644); err != nil {
		return err
	}
	return os.Rename(tmp, s.statePath)
}

func deriveProfileID(appKey string, target string, existing []Profile, sanitize func(string) string) string {
	if sanitize == nil {
		sanitize = defaultSanitizeID
	}
	base := filepath.Base(strings.TrimSpace(appKey))
	if base == "" || base == "/" || base == "." {
		base = "app"
	}
	base = sanitize(base)
	if base == "" {
		base = "app"
	}

	idBase := base + "-" + strings.ToLower(strings.TrimSpace(target))
	id := idBase

	used := map[string]struct{}{}
	for _, p := range existing {
		used[strings.TrimSpace(p.ID)] = struct{}{}
	}
	if _, ok := used[id]; !ok {
		return id
	}
	for i := 2; i < 1000; i++ {
		cand := fmt.Sprintf("%s-%d", idBase, i)
		if _, ok := used[cand]; !ok {
			return cand
		}
	}
	return ""
}

func preferProfile(cand, cur Profile) bool {
	cu := strings.TrimSpace(cand.UpdatedAt)
	ou := strings.TrimSpace(cur.UpdatedAt)
	if cu != ou {
		if cu == "" {
			return false
		}
		if ou == "" {
			return true
		}
		return cu > ou
	}

	cc := strings.TrimSpace(cand.CreatedAt)
	oc := strings.TrimSpace(cur.CreatedAt)
	if cc != oc {
		if cc == "" {
			return false
		}
		if oc == "" {
			return true
		}
		return cc > oc
	}

	if strings.TrimSpace(cand.Command) != "" && strings.TrimSpace(cur.Command) == "" {
		return true
	}
	return false
}

func defaultSanitizeID(s string) string {
	in := strings.ToLower(strings.TrimSpace(s))
	var b strings.Builder
	b.Grow(len(in))
	lastDash := false
	for i := 0; i < len(in); i++ {
		ch := in[i]
		isAZ := ch >= 'a' && ch <= 'z'
		is09 := ch >= '0' && ch <= '9'
		if isAZ || is09 {
			b.WriteByte(ch)
			lastDash = false
			continue
		}
		if !lastDash {
			b.WriteByte('-')
			lastDash = true
		}
	}
	return strings.Trim(b.String(), "-")
}