package proto import "net" // Map is an in-memory aggregate representation: type/ip/count. 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: e.Count, }, }, }) } func (a *Aggregate) Recalc() { tmp := make(Map) tmp.Update(a) b := tmp.ToAggregate() a.ByType = b.ByType } // IPMask is a utility type to mask v4/v6 addresses retaining // different number of bits for each protocol. type IPMask struct { v4Mask net.IPMask v6Mask net.IPMask } var DefaultIPMask = NewIPMask(32, 64) func NewIPMask(v4bits, v6bits int) IPMask { return IPMask{ v4Mask: net.CIDRMask(v4bits, 32), v6Mask: net.CIDRMask(v6bits, 128), } } func (m IPMask) MaskIP(ip net.IP) net.IP { if ip.To4() == nil { return ip.Mask(m.v6Mask) } return ip.Mask(m.v4Mask) } func (m IPMask) MaskString(ipstr string) string { return m.MaskIP(net.ParseIP(ipstr)).String() }