Merge pull request #1462 from mozilla/fixup_duplicates_geomodel

Add test examples for geomodel alerts

Author: pwnbus

alerts/generic_alert_loader.py

@@ -115,9 +115,10 @@ def main(self):
         for cfg in self.configs:
             try:
                 self.process_alert(cfg)
-            except Exception:
+            except Exception as err:
+                self.error_thrown = err
                 traceback.print_exc(file=sys.stdout)
-                logger.error("Processing rule file {} failed".format(cfg.__str__()))
+                logger.exception("Processing rule file {} failed".format(cfg.__str__()))
 
     def onAggregation(self, aggreg):
         # aggreg['count']: number of items in the aggregation, ex: number of failed login attempts

alerts/geomodel/alert.py

@@ -1,13 +1,11 @@
 import math
 from operator import attrgetter
-from typing import NamedTuple, Optional
+from typing import List, NamedTuple, Optional
 
-from .locality import State, Locality
+from .locality import Locality, distance as geo_distance
 
 
-_AIR_TRAVEL_SPEED = 1000.0  # km/h
-
-_EARTH_RADIUS = 6373.0  # km # approximate
+_AIR_TRAVEL_SPEED = 277.778  # m/s
 
 # TODO: Switch to dataclasses when we move to Python3.7+
 
@@ -16,20 +14,30 @@ class Origin(NamedTuple):
     '''A description of a location.
     '''
 
+    ip: str
     city: str
     country: str
     latitude: float
     longitude: float
     geopoint: str
 
 
+class Hop(NamedTuple):
+    '''Describes a hop from one location to another that would be
+    physically impossible in the time between a user's activity in each
+    location.
+    '''
+
+    origin: Origin
+    destination: Origin
+
+
 class Alert(NamedTuple):
     '''A container for the data the alerts output by GeoModel contain.
     '''
 
     username: str
-    sourceipaddress: str
-    origin: Origin
+    hops: List[Hop]
 
 
 def _travel_possible(loc1: Locality, loc2: Locality) -> bool:
@@ -38,53 +46,61 @@ def _travel_possible(loc1: Locality, loc2: Locality) -> bool:
     actions took place.
     '''
 
-    lat1 = math.radians(loc1.latitude)
-    lat2 = math.radians(loc2.latitude)
-    lon1 = math.radians(loc1.longitude)
-    lon2 = math.radians(loc2.longitude)
+    dist_traveled = 1000 * geo_distance(loc1, loc2)  # Convert to metres
 
-    dlat = lat2 - lat1
-    dlon = lon2 - lon1
-
-    a = math.sin(dlat / 2.0) ** 2 +\
-        math.cos(lat1) * math.cos(lat2) * math.sin(dlon / 2.0) ** 2
-    c = 2 * math.atan2(math.sqrt(a), math.sqrt(1 - a))
-
-    distance = c * _EARTH_RADIUS
-
-    seconds_between = (loc2.lastaction - loc1.lastaction).total_seconds()
-    hours_between = math.ceil(seconds_between / 60.0 / 60.0)
+    seconds_between = abs((loc2.lastaction - loc1.lastaction).total_seconds())
 
     # We pad the time with an hour to account for things like planes being
     # slowed, network delays, etc.
-    return (distance / _AIR_TRAVEL_SPEED) <= (hours_between - 1)
+    ttt = (dist_traveled / _AIR_TRAVEL_SPEED)  # Time to travel the distance.
+    pad = math.ceil((1000 * min(loc1.radius, loc2.radius)) / _AIR_TRAVEL_SPEED)
+
+    return (ttt - pad) <= seconds_between
 
 
-def alert(user_state: State) -> Optional[Alert]:
+def alert(
+        username: str,
+        from_evts: List[Locality],
+        from_es: List[Locality]
+) -> Optional[Alert]:
     '''Determine whether an alert should fire given a particular user's
     locality state.  If an alert should fire, an `Alert` is returned, otherwise
     this function returns `None`.
     '''
 
-    locs_to_consider = sorted(user_state.localities, key=attrgetter('lastaction'))
+    relevant_es = sorted(from_es, key=attrgetter('lastaction'), reverse=True)[0:1]
+    all_evts = sorted(from_evts, key=attrgetter('lastaction'))
+    locs_to_consider = relevant_es + all_evts
 
     if len(locs_to_consider) < 2:
         return None
 
-    locations = locs_to_consider[-2:]
-
-    if _travel_possible(*locations):
+    pairs = [
+        (locs_to_consider[i], locs_to_consider[i + 1])
+        for i in range(len(locs_to_consider) - 1)
+    ]
+
+    hops = [
+        Hop(
+            Origin(
+                o.sourceipaddress,
+                o.city,
+                o.country,
+                o.latitude,
+                o.longitude,
+                '{},{}'.format(o.latitude, o.longitude)),
+            Origin(
+                d.sourceipaddress,
+                d.city,
+                d.country,
+                d.latitude,
+                d.longitude,
+                '{},{}'.format(d.latitude, d.longitude)))
+        for (o, d) in pairs
+        if not _travel_possible(o, d)
+    ]
+
+    if len(hops) == 0:
         return None
 
-    (ip, city, country, lat, lon) = (
-        locations[1].sourceipaddress,
-        locations[1].city,
-        locations[1].country,
-        locations[1].latitude,
-        locations[1].longitude
-    )
-
-    geo = '{0},{1}'.format(lat, lon)
-    origin = Origin(city, country, lat, lon, geo)
-
-    return Alert(user_state.username, ip, origin)
+    return Alert(username, hops)

alerts/geomodel/config.py

@@ -18,7 +18,6 @@ class Events(NamedTuple):
 
     search_window: dict
     lucene_query: str
-    username_path: str
 
 
 class Whitelist(NamedTuple):
@@ -38,5 +37,5 @@ class Config(NamedTuple):
     '''
 
     localities: Localities
-    events: List[Events]
+    events: Events
     whitelist: Whitelist

alerts/geomodel/execution.py

@@ -0,0 +1,97 @@
+from datetime import datetime
+from typing import Callable, NamedTuple, Optional
+
+from mozdef_util.elasticsearch_client import ElasticsearchClient as ESClient
+from mozdef_util.query_models import SearchQuery, TermMatch
+from mozdef_util.utilities.toUTC import toUTC
+
+
+_TYPE_NAME = 'execution_state'
+
+
+class ExecutionState(NamedTuple):
+    '''A record of an alert's execution at a particular time, used to create a
+    sliding window through which an alert can query for relevant events and
+    not run the risk of missing any due to relying only on searching some
+    configured amount of time in the past.
+    '''
+
+    type_: str
+    #  alert_name: str
+    execution_time: datetime
+
+    def new(executed_at: Optional[datetime]=None) -> 'ExecutionState':
+        '''Construct a new `ExecutionState` representing the execution of an
+        alert at a specific time.
+        By default, the execution time will be set to when this function is
+        called if not explicitly provided.
+        '''
+
+        if executed_at is None:
+            executed_at = toUTC(datetime.now())
+
+        return ExecutionState(_TYPE_NAME, executed_at)
+
+
+class Record(NamedTuple):
+    '''A container for data identifying an `ExecutionState` in ElasticSearch.
+    '''
+
+    identifier: Optional[str]
+    state: ExecutionState
+
+    def new(state: ExecutionState) -> 'Record':
+        '''Construct a new `Record` that, when stored, will result in a new
+        document being inserted into ElasticSearch.
+        '''
+
+        return Record('', state)
+
+
+Index = str
+StoreInterface = Callable[[Record, Index], None]
+LoadInterface = Callable[[Index], Optional[Record]]
+
+
+def _dict_take(dictionary, keys):
+    return {key: dictionary[key] for key in keys}
+
+
+def store(client: ESClient) -> StoreInterface:
+    '''Wrap an `ElasticsearchClient` in a `StoreInterface` closure to be
+    invoked without requiring direct access to the client in order to
+    persist an `ExecutionState`.
+    '''
+
+    def wrapper(record: Record, esindex: Index):
+        doc = dict(record.state._asdict())
+
+        client.save_object(index=esindex, body=doc, doc_id=record.identifier)
+
+    return wrapper
+
+
+def load(client: ESClient) -> LoadInterface:
+    '''Wrap an `ElasticsearchClient` in a `LoadInterface` closure to be
+    invoed without requiring direct access to the client in order to retrieve
+    an `ExecutionState`.
+    '''
+
+    def wrapper(esindex: Index=None) -> Optional[Record]:
+        query = SearchQuery()
+        query.add_must(TermMatch('type_', _TYPE_NAME))
+
+        results = query.execute(client, indices=[esindex])
+
+        if len(results['hits']) == 0:
+            return None
+
+        eid = results['hits'][0]['_id']
+
+        state = ExecutionState(**_dict_take(
+            results['hits'][0].get('_source', {}),
+            ExecutionState._fields))
+
+        return Record(eid, state)
+
+    return wrapper

alerts/geomodel/locality.py

@@ -1,3 +1,4 @@
+import math
 from datetime import datetime, timedelta
 from typing import Any, Callable, Dict, List, NamedTuple, Optional
 
@@ -9,6 +10,8 @@
 # Default radius (in Kilometres) that a locality should have.
 _DEFAULT_RADIUS_KM = 50.0
 
+_EARTH_RADIUS = 6373.0  # km # approximate
+
 # TODO: Switch to dataclasses when we move to Python3.7+
 
 
@@ -62,17 +65,6 @@ class Update(NamedTuple):
     state: State
     did_update: bool
 
-    def flat_map(fn: Callable[[State], 'Update'], u: 'Update') -> 'Update':
-        '''Apply a function to a `State` that produces an `Update` against the
-        state contained within an established `Update`.  The resulting `Update`
-        will have its `did_update` field set to `True` if either the original
-        or the new `Update` are `True`.
-        '''
-
-        new = fn(u.state)
-
-        return Update(new.state, u.did_update or new.did_update)
-
 
 JournalInterface = Callable[[Entry, str], None]
 QueryInterface = Callable[[SearchQuery, str], Optional[Entry]]
@@ -191,7 +183,7 @@ def update(state: State, from_evt: State) -> Update:
             # If we find that the new state's locality has been recorded
             # for the user in question, we only want to update it if either
             # their IP changed or the new time of activity is more recent.
-            if loc1.city == loc2.city and loc1.country == loc2.country:
+            if distance(loc1, loc2) <= min(loc1.radius, loc2.radius):
                 did_find = True
 
                 new_more_recent = loc1.lastaction > loc2.lastaction
@@ -229,3 +221,23 @@ def remove_outdated(state: State, days_valid: int) -> Update:
     return Update(
         state=State(state.type_, state.username, new_localities),
         did_update=len(new_localities) != len(state.localities))
+
+
+def distance(loc1: Locality, loc2: Locality) -> float:
+    '''Compute the distance between two localities, returning the geographical
+    distance in kilometres (KM).
+    '''
+
+    lat1 = math.radians(loc1.latitude)
+    lat2 = math.radians(loc2.latitude)
+    lon1 = math.radians(loc1.longitude)
+    lon2 = math.radians(loc2.longitude)
+
+    dlat = lat2 - lat1
+    dlon = lon2 - lon1
+
+    a = math.sin(dlat / 2.0) ** 2 +\
+        math.cos(lat1) * math.cos(lat2) * math.sin(dlon / 2.0) ** 2
+    c = 2 * math.atan2(math.sqrt(a), math.sqrt(1 - a))
+
+    return c * _EARTH_RADIUS

alerts/geomodel_location.json

@@ -4,15 +4,12 @@
     "valid_duration_days": 1,
     "radius_kilometres": 50
   },
-  "events": [
-    {
+  "events": {
       "search_window": {
         "minutes": 5
       },
-      "lucene_query": "tags:auth0",
-      "username_path": "details.username"
-    }
-  ],
+      "lucene_query": "tags:auth0"
+  },
   "whitelist": {
     "users": [],
     "cidrs": []