Merge pull request #3 from Microsoft/RSAFixes

Rsa fixes

Author: ianhelle

msticpy/_version.py

@@ -1,2 +1,2 @@
 """Version file."""
-VERSION = "0.0.73"
+VERSION = "0.1.03"

msticpy/nbtools/entityschema.py

@@ -46,23 +46,28 @@ def __init__(self, src_entity=None, **kwargs):
                     self[k] = src_entity[k]
 
                     if v is not None:
-                        if v == RegistryHive.__name__:
-                            self[k] = RegistryHive(src_entity[k])
-                        elif v == OSFamily.__name__:
-                            self[k] = OSFamily(src_entity[k])
-                        elif v == ElevationToken.__name__:
-                            self[k] = ElevationToken(src_entity[k])
-                        elif v == GeoLocation.__name__:
-                            self[k] = GeoLocation(src_entity[k])
-                        elif v == Algorithm.__name__:
-                            self[k] = Algorithm(src_entity[k])
-                        elif isinstance(v, tuple):
-                            entity_list = []
-                            for col_entity in src_entity[k]:
-                                entity_list.append(Entity.instantiate_entity(col_entity))
-                            self[k] = entity_list
-                        else:
-                            self[k] = Entity.instantiate_entity(src_entity[k])
+                        try:
+                            if v == RegistryHive.__name__:
+                                self[k] = RegistryHive[src_entity[k]]
+                            elif v == OSFamily.__name__:
+                                self[k] = OSFamily[src_entity[k]]
+                            elif v == ElevationToken.__name__:
+                                self[k] = ElevationToken[src_entity[k]]
+                            elif v == GeoLocation.__name__:
+                                self[k] = GeoLocation[src_entity[k]]
+                            elif v == Algorithm.__name__:
+                                self[k] = Algorithm[src_entity[k]]
+                            elif isinstance(v, tuple):
+                                entity_list = []
+                                for col_entity in src_entity[k]:
+                                    entity_list.append(Entity.instantiate_entity(col_entity))
+                                self[k] = entity_list
+                            else:
+                                self[k] = Entity.instantiate_entity(src_entity[k])
+                        except KeyError:
+                            # Catch key errors from invalid enum values
+                            self[k] = None
+
             # add AdditionalData dictionary if it's populated
             if 'AdditionalData' in src_entity:
                 self['AdditionalData'] = src_entity['AdditionalData']
@@ -517,6 +522,7 @@ def description_str(self) -> str:
 @export
 class Algorithm(Enum):
     """FileHash Algorithm Enumeration."""
+
     Unknown = 0
     MD5 = 1
     SHA1 = 2
@@ -862,7 +868,7 @@ def description_str(self) -> str:
     _entity_schema = {
         # Hive (type System.Nullable`1
         # [Microsoft.Azure.Security.Detection.AlertContracts.V3.Entities.RegistryHive])
-        'Hive': RegistryHive,
+        'Hive': 'RegistryHive',
         # Key (type System.String)
         'Key': None
     }

msticpy/nbtools/nbwidgets.py

@@ -372,8 +372,11 @@ def _get_alert(self, alert_id):
         if selected_alerts.shape[0] > 0:
             alert = pd.Series(selected_alerts.iloc[0])
             if isinstance(alert['ExtendedProperties'], str):
-                alert['ExtendedProperties'] = json.loads(
-                    (alert['ExtendedProperties']))
+                try:
+                    alert['ExtendedProperties'] = json.loads(
+                        (alert['ExtendedProperties']))
+                except JSONDecodeError:
+                    pass
             if isinstance(alert['Entities'], str):
                 try:
                     alert['Entities'] = json.loads((alert['Entities']))

msticpy/nbtools/security_alert.py

@@ -32,8 +32,8 @@ def __init__(self, src_row: pd.Series = None):
         super().__init__(src_row=src_row)
 
         # add entities to dictionary to remove dups
+        self._src_entities = dict()
         if 'Entities' in src_row:
-            self._src_entities = dict()
             self._extract_entities(src_row)
 
         if 'ExtendedProperties' in src_row:
@@ -44,6 +44,8 @@ def __init__(self, src_row: pd.Series = None):
                     self.extended_properties = json.loads(src_row.ExtendedProperties)
                 except JSONDecodeError:
                     pass
+        else:
+            self.extended_properties = []
         self._find_os_family()
 
     @property

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
-        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: