elmprodvpn/selective-vpn-api/app/trafficcandidates/candidates.go

package trafficcandidates

import (
	"sort"
	"strconv"
	"strings"
	"time"
)

type Subnet struct {
	CIDR     string
	Dev      string
	Kind     string
	LinkDown bool
}

type Unit struct {
	Unit        string
	Description string
	Cgroup      string
}

type UID struct {
	UID      int
	User     string
	Examples []string
}

type Response struct {
	GeneratedAt string
	Subnets     []Subnet
	Units       []Unit
	UIDs        []UID
}

type Deps struct {
	RunCommand              func(name string, args ...string) (stdout, stderr string, exitCode int, err error)
	ParseRouteDevice        func(fields []string) string
	IsVPNLikeIface          func(iface string) bool
	IsContainerIface        func(iface string) bool
	IsAutoBypassDestination func(dst string) bool
}

func Collect(now time.Time, deps Deps) Response {
	if now.IsZero() {
		now = time.Now().UTC()
	}
	return Response{
		GeneratedAt: now.UTC().Format(time.RFC3339),
		Subnets:     candidateSubnets(deps),
		Units:       candidateUnits(deps),
		UIDs:        candidateUIDs(deps),
	}
}

func candidateSubnets(deps Deps) []Subnet {
	if deps.RunCommand == nil {
		return nil
	}
	out, _, code, _ := deps.RunCommand("ip", "-4", "route", "show", "table", "main")
	if code != 0 {
		return nil
	}

	seen := map[string]struct{}{}
	items := make([]Subnet, 0, 24)

	for _, raw := range strings.Split(out, "\n") {
		line := strings.TrimSpace(raw)
		if line == "" {
			continue
		}
		fields := strings.Fields(line)
		if len(fields) == 0 {
			continue
		}

		dst := strings.TrimSpace(fields[0])
		if dst == "" || dst == "default" {
			continue
		}
		dev := ""
		if deps.ParseRouteDevice != nil {
			dev = deps.ParseRouteDevice(fields)
		}
		if dev == "" || dev == "lo" {
			continue
		}
		if deps.IsVPNLikeIface != nil && deps.IsVPNLikeIface(dev) {
			continue
		}

		isDocker := deps.IsContainerIface != nil && deps.IsContainerIface(dev)
		isLocal := deps.IsAutoBypassDestination != nil && deps.IsAutoBypassDestination(dst)
		if !isDocker && !isLocal {
			continue
		}

		kind := "lan"
		if isDocker {
			kind = "docker"
		} else if strings.Contains(" "+strings.ToLower(line)+" ", " scope link ") {
			kind = "link"
		}

		key := kind + "|" + dst + "|" + dev
		if _, ok := seen[key]; ok {
			continue
		}
		seen[key] = struct{}{}

		items = append(items, Subnet{
			CIDR:     dst,
			Dev:      dev,
			Kind:     kind,
			LinkDown: strings.Contains(strings.ToLower(line), " linkdown"),
		})
	}

	sort.Slice(items, func(i, j int) bool {
		if items[i].Kind != items[j].Kind {
			return items[i].Kind < items[j].Kind
		}
		if items[i].Dev != items[j].Dev {
			return items[i].Dev < items[j].Dev
		}
		return items[i].CIDR < items[j].CIDR
	})
	return items
}

func candidateUnits(deps Deps) []Unit {
	if deps.RunCommand == nil {
		return nil
	}
	stdout, _, code, _ := deps.RunCommand(
		"systemctl",
		"list-units",
		"--type=service",
		"--state=running",
		"--no-legend",
		"--no-pager",
		"--plain",
	)
	if code != 0 {
		return nil
	}

	seen := map[string]struct{}{}
	items := make([]Unit, 0, 32)
	for _, raw := range strings.Split(stdout, "\n") {
		line := strings.TrimSpace(raw)
		if line == "" {
			continue
		}
		fields := strings.Fields(line)
		if len(fields) < 1 {
			continue
		}
		unit := strings.TrimSpace(fields[0])
		if unit == "" {
			continue
		}
		if _, ok := seen[unit]; ok {
			continue
		}
		seen[unit] = struct{}{}

		desc := ""
		if len(fields) >= 5 {
			desc = strings.Join(fields[4:], " ")
		}

		items = append(items, Unit{
			Unit:        unit,
			Description: strings.TrimSpace(desc),
			Cgroup:      "system.slice/" + unit,
		})
	}

	sort.Slice(items, func(i, j int) bool {
		return items[i].Unit < items[j].Unit
	})
	return items
}

func candidateUIDs(deps Deps) []UID {
	if deps.RunCommand == nil {
		return nil
	}
	stdout, _, code, _ := deps.RunCommand("ps", "-eo", "uid,user,comm", "--no-headers")
	if code != 0 {
		return nil
	}

	type agg struct {
		uid   int
		user  string
		comms map[string]struct{}
	}

	aggs := map[int]*agg{}
	for _, raw := range strings.Split(stdout, "\n") {
		line := strings.TrimSpace(raw)
		if line == "" {
			continue
		}
		fields := strings.Fields(line)
		if len(fields) < 2 {
			continue
		}
		uidN, err := strconv.Atoi(strings.TrimSpace(fields[0]))
		if err != nil || uidN < 0 {
			continue
		}
		user := strings.TrimSpace(fields[1])
		comm := ""
		if len(fields) >= 3 {
			comm = strings.TrimSpace(fields[2])
		}

		a := aggs[uidN]
		if a == nil {
			a = &agg{uid: uidN, user: user, comms: map[string]struct{}{}}
			aggs[uidN] = a
		}
		if a.user == "" && user != "" {
			a.user = user
		}
		if comm != "" {
			a.comms[comm] = struct{}{}
		}
	}

	items := make([]UID, 0, len(aggs))
	for _, a := range aggs {
		examples := make([]string, 0, len(a.comms))
		for c := range a.comms {
			examples = append(examples, c)
		}
		sort.Strings(examples)
		if len(examples) > 3 {
			examples = examples[:3]
		}
		items = append(items, UID{
			UID:      a.uid,
			User:     a.user,
			Examples: examples,
		})
	}

	sort.Slice(items, func(i, j int) bool {
		return items[i].UID < items[j].UID
	})
	return items
}