Merge pull request #20 from microsoft/pebryan/2021-7-2MinorFixes

Minor fixes to account and ip nblts
microsoft · Oct 27, 2021 · bb908cd · bb908cd
2 parents cc391c9 + 9b2b549
commit bb908cd
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git a/msticnb/nb/azsent/account/account_summary.py b/msticnb/nb/azsent/account/account_summary.py
@@ -1180,6 +1180,9 @@ def _create_ip_summary(data, ip_col, geoip):
         )  # get geoip and whois
         .merge(data, left_on="IpAddress", right_on=ip_col)
     )
+    for col in group_cols:
+        if col not in all_data.columns:
+            group_cols.pop(group_cols.index(col))
     ip_summary = all_data.groupby(group_cols).agg(
         Count=pd.NamedAgg("Operation", "count"),
         FirstOperation=pd.NamedAgg("TimeGenerated", "min"),

diff --git a/msticnb/nb/azsent/network/ip_summary.py b/msticnb/nb/azsent/network/ip_summary.py
@@ -605,7 +605,11 @@ def _get_heartbeat(qry_prov, src_ip, result):
     nb_data_wait("Heartbeat")
     if result.ip_type == "Public":
         result.heartbeat = qry_prov.Network.get_heartbeat_for_ip(ip_address=src_ip)
-    elif result.host_entity.HostName and result.host_entity.HostName != "unknown":
+    elif (
+        result.host_entity
+        and result.host_entity.HostName
+        and result.host_entity.HostName != "unknown"
+    ):
         result.heartbeat = qry_prov.Network.get_heartbeat_for_host(
             host_name=result.host_entity.HostName
         )