From c36541305b0e80ce87b98ef61bdc1ff719b5b972 Mon Sep 17 00:00:00 2001
From: Tom Windell <thomas@thinkst.com>
Date: Wed, 24 Jul 2024 20:11:15 +0200
Subject: [PATCH] Change get_src_ip_continent and fix tests

---
 canarytokens/models.py     |  5 +++--
 canarytokens/utils.py      |  6 +++---
 tests/units/test_models.py | 25 ++++++++++++++++++++-----
 tests/units/test_utils.py  | 26 +++++++++++++-------------
 4 files changed, 39 insertions(+), 23 deletions(-)

diff --git a/canarytokens/models.py b/canarytokens/models.py
index 283f88f66..797e833df 100644
--- a/canarytokens/models.py
+++ b/canarytokens/models.py
@@ -1326,8 +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]
 
-            continent = get_src_ip_continent(additional_data)
-            additional_data["geo_info"]["continent"] = continent
+            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")
diff --git a/canarytokens/utils.py b/canarytokens/utils.py
index d0c53c808..6acddbfdf 100644
--- a/canarytokens/utils.py
+++ b/canarytokens/utils.py
@@ -78,16 +78,16 @@ def get_deployed_commit_sha(commit_sha_file: Path = Path("/COMMIT_SHA")):
 #     return inner
 
 
-def get_src_ip_continent(additional_data: dict) -> 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:
-        additional_data (dict): The "country" key contains an 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 = additional_data.get("geo_info", {}).get("country")
+    country = geo_data.get("country")
     if country is not None:
         # AQ is the ISO 3166-2 code for Antarctica, and is returned from IPinfo,
         # but it's not included in pycountry_convert.
diff --git a/tests/units/test_models.py b/tests/units/test_models.py
index 6b823b47e..c99136c17 100644
--- a/tests/units/test_models.py
+++ b/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",
+                },
             },
         ),
         (
diff --git a/tests/units/test_utils.py b/tests/units/test_utils.py
index 8c1ff2f67..d934a42c1 100644
--- a/tests/units/test_utils.py
+++ b/tests/units/test_utils.py
@@ -25,19 +25,19 @@ def test_coerce_to_float():
 
 
 @pytest.mark.parametrize(
-    "additional_data, continent",
+    "geo_info, continent",
     [
-        ({"geo_info": {"country": "ZA"}}, "AF"),
-        ({"geo_info": {"country": "AQ"}}, "AN"),
-        ({"geo_info": {"country": "CN"}}, "AS"),
-        ({"geo_info": {"country": "GB"}}, "EU"),
-        ({"geo_info": {"country": "US"}}, "NA"),
-        ({"geo_info": {"country": "AU"}}, "OC"),
-        ({"geo_info": {"country": "AR"}}, "SA"),
-        ({"geo_info": {"country": "Mordor"}}, "NO_CONTINENT"),
-        ({"geo_info": {"bogon": True}}, "NO_CONTINENT"),
-        ({"geo_info": {}}, "NO_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(additional_data, continent):
-    assert continent == get_src_ip_continent(additional_data)
+def test_get_src_ip_continent(geo_info, continent):
+    assert continent == get_src_ip_continent(geo_info)