Skip to content
Snippets Groups Projects
iprep.go 1.94 KiB
Newer Older
ale's avatar
ale committed
package proto
ale's avatar
ale committed

ale's avatar
ale committed
// Map is an in-memory aggregate representation: type/ip/count.
ale's avatar
ale committed
type Map map[string]map[string]int64

func (m Map) Update(a *Aggregate) {
	for _, bt := range a.ByType {
		for _, e := range bt.ByIp {
			m.Incr(bt.Type, e.Ip, e.Count)
		}
	}
}

func (m Map) Incr(evType, ip string, n int64) {
	byIP, ok := m[evType]
	if !ok {
		byIP = make(map[string]int64)
		m[evType] = byIP
	}
	byIP[ip] += n
}

func (m Map) ToAggregate() *Aggregate {
	aggr := &Aggregate{
		ByType: make([]*AggregateTypeEntry, 0, len(m)),
	}
	for evType, byType := range m {
		typeEntry := &AggregateTypeEntry{
			Type: evType,
			ByIp: make([]*AggregateIPEntry, 0, len(byType)),
		}
		for ip, count := range byType {
			typeEntry.ByIp = append(typeEntry.ByIp, &AggregateIPEntry{
				Ip:    ip,
				Count: count,
			})
		}
		aggr.ByType = append(aggr.ByType, typeEntry)
	}
	return aggr
}

func (m Map) ByIP(ip string) map[string]int64 {
	out := make(map[string]int64)
	for t, bt := range m {
		if ev, ok := bt[ip]; ok {
			out[t] = ev
		}
	}
	return out
}

// MakeAggregate builds an Aggregate out of Events. Event details are
// lost in the aggregate.
func MakeAggregate(events []*Event) *Aggregate {
	tmp := make(Map)
	for _, ev := range events {
		tmp.Incr(ev.Type, ev.Ip, 1)
	}
	return tmp.ToAggregate()

}

func (a *Aggregate) GetCount(evType, ip string) (int64, bool) {
	for _, bt := range a.ByType {
		if bt.Type != evType {
			continue
		}
		for _, bi := range bt.ByIp {
			if bi.Ip == ip {
				return bi.Count, true
			}
		}
	}
	return 0, false
}

func (a *Aggregate) Merge(b *Aggregate) *Aggregate {
	tmp := make(Map)
	tmp.Update(a)
	tmp.Update(b)
	return tmp.ToAggregate()
}

func (a *Aggregate) AddEvent(e *Event) {
	a.ByType = append(a.ByType, &AggregateTypeEntry{
		Type: e.Type,
		ByIp: []*AggregateIPEntry{
			&AggregateIPEntry{
				Ip:    e.Ip,
				Count: 1,
			},
		},
	})
}

func (a *Aggregate) Recalc() {
	tmp := make(Map)
	tmp.Update(a)
	b := tmp.ToAggregate()
	a.ByType = b.ByType
}