<commit_before>package conntrack

import (
	"net"
)

type TypeFilter uint8

const (
	SNATFilter TypeFilter = 1 << iota
	DNATFilter
	RoutedFilter
	LocalFilter
)

var localIPs = make([]*net.IPNet, 0)

func isLocalIP(ip net.IP) bool {
	for _, localIP := range localIPs {
		if localIP.IP.Equal(ip) {
			return true
		}
	}

	return false
}

func init() {
	addresses, err := net.InterfaceAddrs()
	if err != nil {
		panic(err)
	}

	for _, address := range addresses {
		localIPs = append(localIPs, address.(*net.IPNet))
	}
}

func (flows FlowSlice) FilterByType(which TypeFilter) FlowSlice {
	natFlows := make(FlowSlice, 0, len(flows))

	snat := (which & SNATFilter) > 0
	dnat := (which & DNATFilter) > 0
	local := (which & LocalFilter) > 0
	routed := (which & RoutedFilter) > 0

	for _, flow := range flows {
		if (snat && isSNAT(flow)) ||
			(dnat && isDNAT(flow)) ||
			(local && isLocal(flow)) ||
			(routed && isRouted(flow)) {

			natFlows = append(natFlows, flow)
		}
	}

	return natFlows
}

func (flows FlowSlice) FilterByProtocol(protocol string) FlowSlice {
	filtered := make(FlowSlice, 0, len(flows))

	for _, flow := range flows {
		if flow.Protocol == protocol {
			filtered = append(filtered, flow)
		}
	}

	return filtered
}

func isSNAT(flow Flow) bool {
	// SNATed flows should reply to our WAN IP, not a LAN IP.
	if flow.Original.Source.Equal(flow.Reply.Destination) {
		return false
	}

	if !flow.Original.Destination.Equal(flow.Reply.Source) {
		return false
	}

	return true
}
<commit_msg>conntrack: add filtering by state and a generic func based filter<commit_after>package conntrack

import (
	"net"
)

type TypeFilter uint8

const (
	SNATFilter TypeFilter = 1 << iota
	DNATFilter
	RoutedFilter
	LocalFilter
)

var localIPs = make([]*net.IPNet, 0)

func isLocalIP(ip net.IP) bool {
	for _, localIP := range localIPs {
		if localIP.IP.Equal(ip) {
			return true
		}
	}

	return false
}

func init() {
	addresses, err := net.InterfaceAddrs()
	if err != nil {
		panic(err)
	}

	for _, address := range addresses {
		localIPs = append(localIPs, address.(*net.IPNet))
	}
}

func (flows FlowSlice) Filter(filter func(flow Flow) bool) FlowSlice {
	filtered := make(FlowSlice, 0, len(flows))

	for _, flow := range flows {
		if filter(flow) {
			filtered = append(filtered, flow)
		}
	}

	return filtered
}

func (flows FlowSlice) FilterByType(which TypeFilter) FlowSlice {
	snat := (which & SNATFilter) > 0
	dnat := (which & DNATFilter) > 0
	local := (which & LocalFilter) > 0
	routed := (which & RoutedFilter) > 0

	return flows.Filter(func(flow Flow) bool {
		return ((snat && isSNAT(flow)) ||
			(dnat && isDNAT(flow)) ||
			(local && isLocal(flow)) ||
			(routed && isRouted(flow)))
	})
}

func (flows FlowSlice) FilterByProtocol(protocol string) FlowSlice {
	return flows.Filter(func(flow Flow) bool {
		return flow.Protocol == protocol
	})
}

func (flows FlowSlice) FilterByState(state string) FlowSlice {
	return flows.Filter(func(flow Flow) bool {
		return flow.State == state
	})
}

func isSNAT(flow Flow) bool {
	// SNATed flows should reply to our WAN IP, not a LAN IP.
	if flow.Original.Source.Equal(flow.Reply.Destination) {
		return false
	}

	if !flow.Original.Destination.Equal(flow.Reply.Source) {
		return false
	}

	return true
}

{'fraction_non_alphanumeric': 0.13010125805461797, 'fraction_numerical': 0.005216324025774778, 'mean_word_length': 3.780058651026393, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1133580', 'n_tokens_mistral': 1260, 'n_tokens_neox': 1167, 'n_words': 310}