Refactor entity list handling

intel/entity.go

@@ -16,6 +16,43 @@ import (
 	"golang.org/x/net/publicsuffix"
 )
 
+// ListMatch represents an entity that has been
+// matched against filterlists.
+type ListMatch struct {
+	Entity        string
+	ActiveLists   []string
+	InactiveLists []string
+}
+
+func (lm *ListMatch) String() string {
+	inactive := ""
+	if len(lm.InactiveLists) > 0 {
+		inactive = " and in deactivated lists " + strings.Join(lm.InactiveLists, ", ")
+	}
+	return fmt.Sprintf(
+		"%s in activated lists %s%s",
+		lm.Entity,
+		strings.Join(lm.ActiveLists, ","),
+		inactive,
+	)
+}
+
+// ListBlockReason is a list of list matches.
+type ListBlockReason []ListMatch
+
+func (br ListBlockReason) String() string {
+	if len(br) == 0 {
+		return ""
+	}
+
+	matches := make([]string, len(br))
+	for idx, lm := range br {
+		matches[idx] = lm.String()
+	}
+
+	return strings.Join(matches, " and ")
+}
+
 // Entity describes a remote endpoint in many different ways.
 // It embeddes a sync.Mutex but none of the endpoints own
 // functions performs locking. The caller MUST ENSURE
@@ -60,8 +97,18 @@ type Entity struct {
 
 	location *geoip.Location
 
-	Lists    []string
-	ListsMap filterlists.LookupMap
+	// BlockedByLists holds list source IDs that
+	// are used to block the entity.
+	BlockedByLists []string
+
+	// BlockedEntities holds a list of entities that
+	// have been blocked. Values can be used as a key
+	// for the ListOccurences map.
+	BlockedEntities []string
+
+	// ListOccurences is a map that matches an entity (Domain, IPs, ASN, Country, Sub-domain)
+	// to a list of sources where the entity has been observed in.
+	ListOccurences map[string][]string
 
 	// we only load each data above at most once
 	fetchLocationOnce  sync.Once
@@ -90,8 +137,11 @@ func (e *Entity) ResetLists() {
 	// TODO(ppacher): our actual goal is to reset the domain
 	// list right now so we could be more efficient by keeping
 	// the other lists around.
-	e.Lists = nil
-	e.ListsMap = nil
+
+	// FIXME
+	//e.Lists = nil
+	//e.ListsMap = nil
+	e.ListOccurences = nil
 	e.domainListLoaded = false
 	e.ipListLoaded = false
 	e.countryListLoaded = false
@@ -238,9 +288,19 @@ func (e *Entity) getLists() {
 	e.getCountryLists()
 }
 
-func (e *Entity) mergeList(list []string) {
-	e.Lists = mergeStringList(e.Lists, list)
-	e.ListsMap = buildLookupMap(e.Lists)
+func (e *Entity) mergeList(key string, list []string) {
+	if len(list) == 0 {
+		return
+	}
+
+	if e.ListOccurences == nil {
+		e.ListOccurences = make(map[string][]string)
+	}
+
+	e.ListOccurences[key] = mergeStringList(e.ListOccurences[key], list)
+
+	//e.Lists = mergeStringList(e.Lists, list)
+	//e.ListsMap = buildLookupMap(e.Lists)
 }
 
 func (e *Entity) getDomainLists() {
@@ -284,7 +344,7 @@ func (e *Entity) getDomainLists() {
 				return
 			}
 
-			e.mergeList(list)
+			e.mergeList(d, list)
 		}
 		e.domainListLoaded = true
 	})
@@ -328,7 +388,8 @@ func (e *Entity) getASNLists() {
 
 	log.Tracef("intel: loading ASN list for %d", asn)
 	e.loadAsnListOnce.Do(func() {
-		list, err := filterlists.LookupASNString(fmt.Sprintf("%d", asn))
+		asnStr := fmt.Sprintf("%d", asn)
+		list, err := filterlists.LookupASNString(asnStr)
 		if err != nil {
 			log.Errorf("intel: failed to get ASN blocklist for %d: %s", asn, err)
 			e.loadAsnListOnce = sync.Once{}
@@ -336,7 +397,7 @@ func (e *Entity) getASNLists() {
 		}
 
 		e.asnListLoaded = true
-		e.mergeList(list)
+		e.mergeList(asnStr, list)
 	})
 }
 
@@ -360,7 +421,7 @@ func (e *Entity) getCountryLists() {
 		}
 
 		e.countryListLoaded = true
-		e.mergeList(list)
+		e.mergeList(country, list)
 	})
 }
 
@@ -393,28 +454,71 @@ func (e *Entity) getIPLists() {
 			return
 		}
 		e.ipListLoaded = true
-		e.mergeList(list)
+		e.mergeList(ip.String(), list)
 	})
 }
 
-// GetLists returns the filter list identifiers the entity matched and whether this data is set.
-func (e *Entity) GetLists() ([]string, bool) {
+// LoadLists searches all filterlists for all occurences of
+// this entity.
+func (e *Entity) LoadLists() bool {
 	e.getLists()
 
-	if e.Lists == nil {
-		return nil, false
+	if e.ListOccurences == nil {
+		return false
 	}
-	return e.Lists, true
+	return true
 }
 
-// GetListsMap is like GetLists but returns a lookup map for list IDs.
-func (e *Entity) GetListsMap() (filterlists.LookupMap, bool) {
-	e.getLists()
+// MatchLists matches the entities lists against a slice
+// of source IDs and  updates various entity properties
+// like BlockedByLists, ListOccurences and BlockedEntitites.
+func (e *Entity) MatchLists(lists []string) bool {
+	e.BlockedByLists = nil
+	e.BlockedEntities = nil
 
-	if e.ListsMap == nil {
-		return nil, false
+	lm := makeMap(lists)
+	for key, keyLists := range e.ListOccurences {
+		for _, keyListID := range keyLists {
+			if _, ok := lm[keyListID]; ok {
+				e.BlockedByLists = append(e.BlockedByLists, keyListID)
+				e.BlockedEntities = append(e.BlockedEntities, key)
+			}
+		}
 	}
-	return e.ListsMap, true
+
+	makeDistinct(e.BlockedByLists)
+
+	return len(e.BlockedByLists) > 0
+}
+
+// ListBlockReason returns the block reason for this entity.
+func (e *Entity) ListBlockReason() ListBlockReason {
+	blockedBy := make([]ListMatch, len(e.BlockedEntities))
+
+	lm := makeMap(e.BlockedByLists)
+
+	for idx, blockedEntity := range e.BlockedEntities {
+		if entityLists, ok := e.ListOccurences[blockedEntity]; ok {
+			var activeLists []string
+			var inactiveLists []string
+
+			for _, l := range entityLists {
+				if _, ok := lm[l]; ok {
+					activeLists = append(activeLists, l)
+				} else {
+					inactiveLists = append(inactiveLists, l)
+				}
+			}
+
+			blockedBy[idx] = ListMatch{
+				Entity:        blockedEntity,
+				ActiveLists:   activeLists,
+				InactiveLists: inactiveLists,
+			}
+		}
+	}
+
+	return blockedBy
 }
 
 func mergeStringList(a, b []string) []string {
@@ -434,21 +538,26 @@ func mergeStringList(a, b []string) []string {
 	return res
 }
 
-func buildLookupMap(l []string) filterlists.LookupMap {
-	m := make(filterlists.LookupMap, len(l))
-
-	for _, s := range l {
-		m[s] = struct{}{}
-	}
-
-	return m
-}
-
 func makeDistinct(slice []string) []string {
-	lm := buildLookupMap(slice)
-	result := make([]string, 0, len(lm))
-	for key := range lm {
-		result = append(result, key)
+	m := make(map[string]struct{}, len(slice))
+	var result []string
+
+	for _, v := range slice {
+		if _, ok := m[v]; ok {
+			continue
+		}
+
+		m[v] = struct{}{}
+		result = append(result, v)
 	}
+
 	return result
 }
+
+func makeMap(slice []string) map[string]struct{} {
+	lm := make(map[string]struct{})
+	for _, v := range slice {
+		lm[v] = struct{}{}
+	}
+	return lm
+}