Updating CountryName to CountryOrRegionName (#796)

* Updating CountryName to CountryOrRegionName Fixing test data and tests. Updating version * Remove unneeded importlib-resources from setup.py * Suppressing flake8 objections to method capitalization
microsoft · Sep 12, 2024 · b5e052a · b5e052a
1 parent 2a81450
commit b5e052a
Show file tree

Hide file tree

Showing 10 changed files with 28 additions and 17 deletions.
diff --git a/msticpy/_version.py b/msticpy/_version.py
@@ -1,3 +1,3 @@
 """Version file."""
 
-VERSION = "2.12.0"
+VERSION = "2.13.0"
diff --git a/msticpy/analysis/syslog_utils.py b/msticpy/analysis/syslog_utils.py
@@ -97,7 +97,7 @@ def create_host_record(
         ip_entity = IpAddress()
         ip_entity.Address = host_hb["ComputerIP"]
         geoloc_entity = GeoLocation()
-        geoloc_entity.CountryName = host_hb["RemoteIPCountry"]  # type: ignore
+        geoloc_entity.CountryOrRegionName = host_hb["RemoteIPCountry"]  # type: ignore
         geoloc_entity.Longitude = host_hb["RemoteIPLongitude"]  # type: ignore
         geoloc_entity.Latitude = host_hb["RemoteIPLatitude"]  # type: ignore
         ip_entity.Location = geoloc_entity  # type: ignore

diff --git a/msticpy/context/geoip.py b/msticpy/context/geoip.py
@@ -301,7 +301,7 @@ def _create_ip_entity(ip_loc: dict, ip_entity) -> IpAddress:
         geo_entity = GeoLocation()
         geo_entity.CountryCode = ip_loc["country_code"]
 
-        geo_entity.CountryName = ip_loc["country_name"]
+        geo_entity.CountryOrRegionName = ip_loc["country_name"]
         geo_entity.State = ip_loc["region_name"]
         geo_entity.City = ip_loc["city"]
         geo_entity.Longitude = ip_loc["longitude"]
@@ -554,7 +554,7 @@ def _create_ip_entity(
             ip_entity.Address = ip_address
         geo_entity = GeoLocation()
         geo_entity.CountryCode = geo_match.get("country", {}).get("iso_code", None)
-        geo_entity.CountryName = (
+        geo_entity.CountryOrRegionName = (
             geo_match.get("country", {}).get("names", {}).get("en", None)
         )
         subdivs = geo_match.get("subdivisions", [])

diff --git a/msticpy/context/ip_utils.py b/msticpy/context/ip_utils.py
@@ -192,7 +192,7 @@ def create_ip_record(
     ip_entity.VMUUID = ip_hb["VMUUID"]  # type: ignore
     ip_entity.SubscriptionId = ip_hb["SubscriptionId"]  # type: ignore
     geoloc_entity = GeoLocation()  # type: ignore
-    geoloc_entity.CountryName = ip_hb["RemoteIPCountry"]  # type: ignore
+    geoloc_entity.CountryOrRegionName = ip_hb["RemoteIPCountry"]  # type: ignore
     geoloc_entity.Longitude = ip_hb["RemoteIPLongitude"]  # type: ignore
     geoloc_entity.Latitude = ip_hb["RemoteIPLatitude"]  # type: ignore
     ip_entity.Location = geoloc_entity  # type: ignore

diff --git a/msticpy/context/tiproviders/ip_quality_score.py b/msticpy/context/tiproviders/ip_quality_score.py
@@ -70,7 +70,7 @@ def parse_results(self: Self, response: dict) -> tuple[bool, ResultSeverity, Any
                 "FraudScore": response["RawResult"].get("fraud_score"),
                 "ISP": response["RawResult"].get("ISP"),
                 "ASN": response["RawResult"].get("ASN"),
-                "Country": response["RawResult"].get("country_code"),
+                "CountryOrRegion": response["RawResult"].get("country_code"),
                 "Region": response["RawResult"].get("city"),
                 "City": response["RawResult"].get("region"),
                 "Organization": response["RawResult"].get("organization"),

diff --git a/msticpy/datamodel/entities/geo_location.py b/msticpy/datamodel/entities/geo_location.py
@@ -26,7 +26,7 @@ class GeoLocation(Entity, ContextObject):
     ----------
     CountryCode : str
         GeoLocation CountryCode
-    CountryName : str
+    CountryOrRegionName : str
         GeoLocation CountryName
     State : str
         GeoLocation State
@@ -62,7 +62,7 @@ def __init__(self, src_entity: Mapping[str, Any] = None, **kwargs):
 
         """
         self.CountryCode: Optional[str] = None
-        self.CountryName: Optional[str] = None
+        self.CountryOrRegionName: Optional[str] = None
         self.State: Optional[str] = None
         self.City: Optional[str] = None
         self.Longitude: Optional[float] = None
@@ -80,6 +80,16 @@ def name_str(self) -> str:
         """Return Entity Name."""
         return self.CountryCode or self.__class__.__name__
 
+    @property
+    def CountryName(self) -> Optional[str]:  # noqa: N802
+        """Return CountryName."""
+        return self.CountryOrRegionName
+
+    @CountryName.setter
+    def CountryName(self, value: str):  # noqa: N802
+        """Set CountryName."""
+        self.CountryOrRegionName = value
+
     @property
     def coordinates(self) -> Tuple[float, float]:
         """Return Latitude/Longitude as a tuple of floats."""
@@ -91,7 +101,7 @@ def coordinates(self) -> Tuple[float, float]:
         # str
         "CountryCode": None,
         # str
-        "CountryName": None,
+        "CountryOrRegionName": None,
         # str
         "State": None,
         # str

diff --git a/msticpy/vis/foliummap.py b/msticpy/vis/foliummap.py
@@ -744,7 +744,8 @@ def icon_mapper(icon_key):
             tooltip_columns = _default_columns(data, [ip_column, "CountryCode", "City"])
         if not popup_columns:
             popup_columns = _default_columns(
-                data, [ip_column, "CountryName", "City", lat_column, long_column]
+                data,
+                [ip_column, "CountryOrRegionName", "City", lat_column, long_column],
             )
     else:
         if not tooltip_columns:

diff --git a/tests/init/pivot/test_pivot_register.py b/tests/init/pivot/test_pivot_register.py
@@ -227,7 +227,7 @@ def data_providers():
             pivot_func="geoloc",
             func_param="value",
             src_col="ip",
-            exp_col="CountryName",
+            exp_col="CountryOrRegionName",
             exp_val="United States",
         ),
         id="IpAddress-geoip_maxmind",
@@ -243,7 +243,7 @@ def data_providers():
     #         pivot_func="geoloc_ips",
     #         func_param="value",
     #         src_col="ip",
-    #         exp_col="CountryName",
+    #         exp_col="CountryOrRegionName",
     #         exp_val="United States",
     #     ),
     #     id="IpAddress-geoip_ipstack",

diff --git a/tests/testdata/ip_locs.csv b/tests/testdata/ip_locs.csv
@@ -1,4 +1,4 @@
-,AllExtIPs,AdditionalData,Type,CountryCode,CountryName,State,City,Longitude,Latitude,IpAddress
+,AllExtIPs,AdditionalData,Type,CountryCode,CountryOrRegionName,State,City,Longitude,Latitude,IpAddress
 0,65.55.44.109,{},geolocation,US,United States,Virginia,Boydton,-78.375,36.6534,65.55.44.109
 1,13.71.172.128,{},geolocation,CA,Canada,Ontario,Toronto,-79.4195,43.6644,13.71.172.128
 2,13.71.172.130,{},geolocation,CA,Canada,Ontario,Toronto,-79.4195,43.6644,13.71.172.130

diff --git a/tests/vis/test_folium.py b/tests/vis/test_folium.py
@@ -222,7 +222,7 @@ def create_geo_entity(row):
     """Return GeoLocation entity - test helper."""
     # get subset of fields for GeoLocation
     loc_props = row[
-        ["CountryCode", "CountryName", "State", "City", "Longitude", "Latitude"]
+        ["CountryCode", "CountryOrRegionName", "State", "City", "Longitude", "Latitude"]
     ]
     return GeoLocation(**loc_props.to_dict())
 
@@ -242,7 +242,7 @@ def icon_map_func(key):
 _IP_COL = "AllExtIPs"
 _LAT_COL = "Latitude"
 _LONG_COL = "Longitude"
-_COUNTRY_COL = "CountryName"
+_COUNTRY_COL = "CountryOrRegionName"
 
 
 @attr.s(auto_attribs=True)
@@ -272,8 +272,8 @@ class PlotMapTest:
     (PlotMapTest(name="layer", **_STD_ARGS, layer_column=_COUNTRY_COL)),  # type: ignore
     (PlotMapTest(name="icon_dict", **_STD_ARGS, icon_column=_COUNTRY_COL, icon_map=icon_map)),  # type: ignore
     (PlotMapTest(name="icon_map", **_STD_ARGS, icon_column=_COUNTRY_COL, icon_map=icon_map_func)),  # type: ignore
-    (PlotMapTest(name="popup", **_STD_ARGS, popup_columns=["CountryName", "State", "City"])),  # type: ignore
-    (PlotMapTest(name="tooltip", **_STD_ARGS, tooltip_columns=["CountryName", "State", "City"])),  # type: ignore
+    (PlotMapTest(name="popup", **_STD_ARGS, popup_columns=["CountryOrRegionName", "State", "City"])),  # type: ignore
+    (PlotMapTest(name="tooltip", **_STD_ARGS, tooltip_columns=["CountryOrRegionName", "State", "City"])),  # type: ignore
 ]
 
 _PM_IDS = [pmt.name for pmt in _PM_TEST_PARAMS]