Trying to avoid query_params attribute error

Author: ianhelle

msticpy/nbtools/entityschema.py

@@ -522,7 +522,7 @@ def description_str(self) -> str:
 @export
 class Algorithm(Enum):
     """FileHash Algorithm Enumeration."""
-    
+
     Unknown = 0
     MD5 = 1
     SHA1 = 2

msticpy/nbtools/security_alert_graph.py

@@ -28,7 +28,8 @@ def create_alert_graph(alert: SecurityAlert):
     alertentity_graph.add_node(alert['AlertType'],
                                name=alert['AlertType'],
                                time=str(alert['StartTimeUtc']),
-                               description='Alert: ' + alert['AlertDisplayName'],
+                               description='Alert: ' +
+                               alert['AlertDisplayName'],
                                color='red',
                                node_type='alert')
 
@@ -71,7 +72,8 @@ def create_alert_graph(alert: SecurityAlert):
 
             # if we have a previously created an edge to the alert, remove it
             if alertentity_graph.has_edge(alert['AlertType'], related_entity):
-                alertentity_graph.remove_edge(alert['AlertType'], related_entity)
+                alertentity_graph.remove_edge(
+                    alert['AlertType'], related_entity)
 
         # if we haven't added an edge to this entity from anything else,
         # add one to the alert
@@ -92,10 +94,13 @@ def add_related_alerts(related_alerts: pd.DataFrame, alertgraph: nx.Graph) ->nx.
 
     alert_host_node = _find_graph_node(related_alerts_graph, 'host', '')
 
-    related_alerts.apply(lambda x: _add_alert_node(related_alerts_graph, x), axis=1)
-    related_alerts.apply(lambda x: _add_related_alert_edges(related_alerts_graph,
-                                                            x,
-                                                            alert_host_node), axis=1)
+    related_alerts.apply(lambda x: _add_alert_node(
+        related_alerts_graph, x), axis=1)
+    if alert_host_node:
+        related_alerts.apply(lambda x:
+                             _add_related_alert_edges(related_alerts_graph,
+                                                      x,
+                                                      alert_host_node), axis=1)
     return related_alerts_graph
 
 
@@ -105,25 +110,29 @@ def _add_related_alert_edges(related_alerts_graph, alert_row, default_node):
         acct_node = _find_graph_node(related_alerts_graph, 'account',
                                      related_alert.primary_account.qualified_name)
         if acct_node is not None:
-            _add_related_alert_edge(related_alerts_graph, acct_node, related_alert)
+            _add_related_alert_edge(
+                related_alerts_graph, acct_node, related_alert)
 
     if related_alert.primary_process is not None:
         proc_node = _find_graph_node(related_alerts_graph,
                                      'process',
                                      related_alert.primary_process.ProcessFilePath)
         if proc_node is not None:
-            _add_related_alert_edge(related_alerts_graph, proc_node, related_alert)
+            _add_related_alert_edge(
+                related_alerts_graph, proc_node, related_alert)
 
     if related_alert.primary_host is not None:
         host_node = _find_graph_node(related_alerts_graph,
                                      'host', related_alert.primary_host['HostName'])
         if host_node is not None:
-            _add_related_alert_edge(related_alerts_graph, host_node, related_alert)
+            _add_related_alert_edge(
+                related_alerts_graph, host_node, related_alert)
 
     # if we haven't added an edge to this entity from anything else,
     # add one to the alert
     if not related_alerts_graph[related_alert['AlertType'] + '(R)']:
-        _add_related_alert_edge(related_alerts_graph, default_node, related_alert)
+        _add_related_alert_edge(related_alerts_graph,
+                                default_node, related_alert)
 
 
 def _add_alert_node(nx_graph, alert):
@@ -159,9 +168,11 @@ def _add_related_alert_edge(nx_graph, source, target):
 
     description = 'Related alert: {}  Count:{}'.format(target['AlertType'],
                                                        current_count)
-    node_attrs = {target_node: {'count': current_count, 'description': description}}
+    node_attrs = {target_node: {
+        'count': current_count, 'description': description}}
     nx.set_node_attributes(nx_graph, node_attrs)
-    nx_graph.add_edge(source, target_node, weight=0.7, description='Related Alert')
+    nx_graph.add_edge(source, target_node, weight=0.7,
+                      description='Related Alert')
 
 
 def _get_account_qualified_name(account):
@@ -261,7 +272,8 @@ def _get_process_name_desc(entity):
 
 
 def _get_account_name_desc(entity):
-    e_name = (entity['NTDomain'] + '\\' if 'NTDomain' in entity else '') + entity['Name']
+    e_name = (entity['NTDomain'] +
+              '\\' if 'NTDomain' in entity else '') + entity['Name']
     e_name = '{}: {}'.format(entity['Type'], e_name)
     if 'IsDomainJoined' in entity:
         domain_joined = entity['IsDomainJoined']
@@ -291,6 +303,7 @@ def _get_host_name_desc(entity, os_family):
         domain_joined = 'false'
     if 'OSFamily' in entity:
         os_family = entity['OSFamily']
-    e_description = '{}\n({}, Domain-joined: {})'.format(e_name, os_family, domain_joined)
+    e_description = '{}\n({}, Domain-joined: {})'.format(e_name,
+                                                         os_family, domain_joined)
 
     return e_name, e_description

msticpy/nbtools/security_base.py

@@ -186,32 +186,38 @@ def query_params(self):
             dict(str, str) -- Dictionary of parameter names
 
         """
-        host_name = (self.primary_host.fqdn if self.primary_host else None)
-        proc_name = (self.primary_process.ImageFile.FullPath if
-                     self.primary_process and self.primary_process.ImageFile
-                     else None)
-        acct_name = self.primary_account.Name if self.primary_account else None
-        path_separator = self.path_separator
-        if self.data_family == DataFamily.WindowsSecurity:
-            proc_name = escape_windows_path(proc_name)
-            path_separator = escape_windows_path(self.path_separator)
-
-        dyn_query_params = {
-            'subscription_filter': self.subscription_filter(),
-            'host_filter_eq': self.host_filter(operator='=='),
-            'host_filter_neq': self.host_filter(operator='!='),
-            'host_name': host_name,
-            'account_name': acct_name,
-            'process_name': proc_name,
-            'logon_session_id': self.get_logon_id(),
-            'process_id': (self.primary_process.ProcessId if self.primary_process else None),
-            'path_separator': path_separator,
-            'data_family': self.data_family,
-            'data_environment': self.data_environment
-        }
-
-        dyn_query_params.update(self._custom_query_params)
-        return dyn_query_params
+        try:
+            if self.primary_host:
+                host_name = self.primary_host.fqdn
+            else:
+                host_name = None
+            proc_name = (self.primary_process.ImageFile.FullPath if
+                         self.primary_process and self.primary_process.ImageFile
+                         else None)
+            acct_name = self.primary_account.Name if self.primary_account else None
+            path_separator = self.path_separator
+            if self.data_family == DataFamily.WindowsSecurity:
+                proc_name = escape_windows_path(proc_name)
+                path_separator = escape_windows_path(self.path_separator)
+
+            dyn_query_params = {
+                'subscription_filter': self.subscription_filter(),
+                'host_filter_eq': self.host_filter(operator='=='),
+                'host_filter_neq': self.host_filter(operator='!='),
+                'host_name': host_name,
+                'account_name': acct_name,
+                'process_name': proc_name,
+                'logon_session_id': self.get_logon_id(),
+                'process_id': (self.primary_process.ProcessId if self.primary_process else None),
+                'path_separator': path_separator,
+                'data_family': self.data_family,
+                'data_environment': self.data_environment,
+            }
+
+            dyn_query_params.update(self._custom_query_params)
+            return dyn_query_params
+        except AttributeError:
+            return {}
 
     @property
     def data_family(self) -> DataFamily: