<commit_before>package flows

import (
	"bytes"
	"net"
	"sort"
	"sync"
	"time"

	"github.com/PreetamJinka/cistern/state/metrics"
)

const (
	FlowProtocolTCP = "TCP"
	FlowProtocolUDP = "UDP"

	alpha float32 = 1
)

type Flow struct {
	Source      net.IP `json:"source"`
	Destination net.IP `json:"destination"`
	Protocol    string `json:"protocol"`

	bytesDerivative   metrics.MetricState
	packetsDerivative metrics.MetricState

	bytes   uint64
	packets uint64

	BytesPerSecond   float32 `json:"bytesPerSecond"`
	PacketsPerSecond float32 `json:"packetsPerSecond`

	lastUpdated time.Time
}

type flowKey struct {
	protocol string
	addr1    [16]byte
	port1    int
	addr2    [16]byte
	port2    int
}

func toFlowKey(protocol string, source, destination net.IP, sourcePort, destPort int) flowKey {
	source = source.To16()
	destination = destination.To16()

	key := flowKey{
		protocol: protocol,
	}

	if bytes.Compare([]byte(source), []byte(destination)) < 0 {
		copy(key.addr1[:], source)
		copy(key.addr2[:], destination)
		key.port1 = sourcePort
		key.port2 = destPort
		return key
	}

	copy(key.addr1[:], destination)
	copy(key.addr2[:], source)
	key.port1 = destPort
	key.port2 = sourcePort
	return key
}

type TopTalkers struct {
	flows map[flowKey]Flow

	lock sync.Mutex
}

func NewTopTalkers(window time.Duration) *TopTalkers {
	toptalkers := &TopTalkers{
		flows: map[flowKey]Flow{},
	}

	go func() {
		for _ = range time.Tick(window) {
			toptalkers.compact(window)
		}
	}()

	return toptalkers
}

func (t *TopTalkers) compact(window time.Duration) {
	t.lock.Lock()
	defer t.lock.Unlock()

	timeout := time.Now().Add(-window)

	for key, flow := range t.flows {
		if flow.lastUpdated.Before(timeout) {
			delete(t.flows, key)
		}
	}
}

func (t *TopTalkers) Update(protocol string,
	source, destination net.IP, sourcePort, destPort int, bytes int) {
	t.lock.Lock()
	defer t.lock.Unlock()

	key := toFlowKey(protocol, source, destination, sourcePort, destPort)
	var flow Flow

	var present bool
	if flow, present = t.flows[key]; !present {
		flow.Source = source
		flow.Destination = destination
		flow.bytesDerivative = metrics.DerivativeState{}
		flow.packetsDerivative = metrics.DerivativeState{}
		flow.Protocol = protocol
	}<commit_msg>add ports to flow data<commit_after>package flows

import (
	"bytes"
	"net"
	"sort"
	"sync"
	"time"

	"github.com/PreetamJinka/cistern/state/metrics"
)

const (
	FlowProtocolTCP = "TCP"
	FlowProtocolUDP = "UDP"

	alpha float32 = 1
)

type Flow struct {
	Source          net.IP `json:"source"`
	Destination     net.IP `json:"destination"`
	SourcePort      int    `json:"sourcePort"`
	DestinationPort int    `json:"destinationPort"`

	Protocol string `json:"protocol"`

	bytesDerivative   metrics.MetricState
	packetsDerivative metrics.MetricState

	bytes   uint64
	packets uint64

	BytesPerSecond   float32 `json:"bytesPerSecond"`
	PacketsPerSecond float32 `json:"packetsPerSecond`

	lastUpdated time.Time
}

type flowKey struct {
	protocol string
	addr1    [16]byte
	port1    int
	addr2    [16]byte
	port2    int
}

func toFlowKey(protocol string, source, destination net.IP, sourcePort, destPort int) flowKey {
	source = source.To16()
	destination = destination.To16()

	key := flowKey{
		protocol: protocol,
	}

	if bytes.Compare([]byte(source), []byte(destination)) < 0 {
		copy(key.addr1[:], source)
		copy(key.addr2[:], destination)
		key.port1 = sourcePort
		key.port2 = destPort
		return key
	}

	copy(key.addr1[:], destination)
	copy(key.addr2[:], source)
	key.port1 = destPort
	key.port2 = sourcePort
	return key
}

type TopTalkers struct {
	flows map[flowKey]Flow

	lock sync.Mutex
}

func NewTopTalkers(window time.Duration) *TopTalkers {
	toptalkers := &TopTalkers{
		flows: map[flowKey]Flow{},
	}

	go func() {
		for _ = range time.Tick(window) {
			toptalkers.compact(window)
		}
	}()

	return toptalkers
}

func (t *TopTalkers) compact(window time.Duration) {
	t.lock.Lock()
	defer t.lock.Unlock()

	timeout := time.Now().Add(-window)

	for key, flow := range t.flows {
		if flow.lastUpdated.Before(timeout) {
			delete(t.flows, key)
		}
	}
}

func (t *TopTalkers) Update(protocol string,
	source, destination net.IP, sourcePort, destPort int, bytes int) {
	t.lock.Lock()
	defer t.lock.Unlock()

	key := toFlowKey(protocol, source, destination, sourcePort, destPort)
	var flow Flow

	var present bool
	if flow, present = t.flows[key]; !present {
		flow.Source = source
		flow.Destination = destination
		flow.SourcePort = sourcePort
		flow.DestinationPort = destPort
		flow.bytesDerivative = metrics.DerivativeState{}
		flow.packetsDerivative = metrics.DerivativeState{}
		flow.Protocol = protocol
	}
{'fraction_non_alphanumeric': 0.12024908739531888, 'fraction_numerical': 0.0137427528451793, 'mean_word_length': 4.129955947136564, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 3, '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': '471703', 'n_tokens_mistral': 1747, 'n_tokens_neox': 1640, 'n_words': 384}