Fix get_src_ip_continent and update unit tests (#549)

* Update get_scr_ip_continent to handle bogon responses from IPinfo API

* Change get_src_ip_continent and fix tests

* Change indentation

* Early return on get_src_ip_country

canarytokens/models.py

@@ -1326,10 +1326,9 @@ def get_additional_data_for_notification(self) -> Dict[str, Any]:
             if self.src_data and key in self.src_data:
                 self.src_data[replacement] = self.src_data[key]
 
-        if additional_data.get("geo_info", {}).get("country") is not None:
-            additional_data["geo_info"]["continent"] = get_src_ip_continent(
-                additional_data["geo_info"]["country"]
-            )
+        if additional_data.get("geo_info") is not None:
+            continent = get_src_ip_continent(additional_data["geo_info"])
+            additional_data["geo_info"]["continent"] = continent
 
         time = datetime.utcnow()
         additional_data["time_hm"] = time.strftime("%H:%M")

canarytokens/utils.py

@@ -78,15 +78,18 @@ def get_deployed_commit_sha(commit_sha_file: Path = Path("/COMMIT_SHA")):
 #     return inner
 
 
-def get_src_ip_continent(country: str) -> str:
+def get_src_ip_continent(geo_data: dict) -> str:
     """Helper function that returns the continent of country given it's ISO 3166-2 code.
 
     Args:
-        country (str): ISO 3166-2 code
+        geo_data (dict): The "country" key contains an ISO 3166-2 code
 
     Returns:
         str: A two character code representing a continent
     """
+    country = geo_data.get("country")
+    if country is None:
+        return "NO_CONTINENT"
     # AQ is the ISO 3166-2 code for Antarctica, and is returned from IPinfo,
     # but it's not included in pycountry_convert.
     if country == "AQ":

tests/units/test_models.py

@@ -19,7 +19,6 @@
     DNSTokenRequest,
     DownloadContentTypes,
     DownloadMSWordResponse,
-    GeoIPBogonInfo,
     GeoIPInfo,
     LegacyTokenHistory,
     LegacyTokenHit,
@@ -476,11 +475,19 @@ def test_all_requests_have_a_response():
             WebBugTokenHit,
             {
                 "useragent": "python 3.10",
-                "geo_info": GeoIPBogonInfo(ip="127.0.0.1", bogon=True),
+                "geo_info": {
+                    "ip": "127.0.0.1",
+                    "bogon": True,
+                    "continent": "NO_CONTINENT",
+                },
             },
             {
                 "useragent": "python 3.10",
-                "geo_info": GeoIPBogonInfo(ip="127.0.0.1", bogon=True),
+                "geo_info": {
+                    "ip": "127.0.0.1",
+                    "bogon": True,
+                    "continent": "NO_CONTINENT",
+                },
             },
         ),
         (
@@ -584,11 +591,19 @@ def test_get_additional_data_for_email(history_type, hit_type, seed_data):
         (
             {
                 "useragent": "python 3.10",
-                "geo_info": GeoIPBogonInfo(ip="127.0.0.1", bogon=True),
+                "geo_info": {
+                    "ip": "127.0.0.1",
+                    "bogon": True,
+                    "continent": "NO_CONTINENT",
+                },
             },
             {
                 "useragent": "python 3.10",
-                "geo_info": GeoIPBogonInfo(ip="127.0.0.1", bogon=True),
+                "geo_info": {
+                    "ip": "127.0.0.1",
+                    "bogon": True,
+                    "continent": "NO_CONTINENT",
+                },
             },
         ),
         (

tests/units/test_utils.py

@@ -1,3 +1,5 @@
+import pytest
+
 from canarytokens.utils import (
     coerce_to_float,
     get_deployed_commit_sha,
@@ -22,12 +24,20 @@ def test_coerce_to_float():
     assert not coerce_to_float("notafloat")
 
 
-def test_get_src_ip_continent():
-    assert "AF" == get_src_ip_continent("ZA")
-    assert "AN" == get_src_ip_continent("AQ")
-    assert "AS" == get_src_ip_continent("CN")
-    assert "EU" == get_src_ip_continent("GB")
-    assert "NA" == get_src_ip_continent("US")
-    assert "OC" == get_src_ip_continent("AU")
-    assert "SA" == get_src_ip_continent("AR")
-    assert "NO_CONTINENT" == get_src_ip_continent("1234")
+@pytest.mark.parametrize(
+    "geo_info, continent",
+    [
+        ({"country": "ZA"}, "AF"),
+        ({"country": "AQ"}, "AN"),
+        ({"country": "CN"}, "AS"),
+        ({"country": "GB"}, "EU"),
+        ({"country": "US"}, "NA"),
+        ({"country": "AU"}, "OC"),
+        ({"country": "AR"}, "SA"),
+        ({"country": "Mordor"}, "NO_CONTINENT"),
+        ({"bogon": True}, "NO_CONTINENT"),
+        ({}, "NO_CONTINENT"),
+    ],
+)
+def test_get_src_ip_continent(geo_info, continent):
+    assert continent == get_src_ip_continent(geo_info)