Build v2.3 - Multiple Enhancements including mailMessage processing

.github/workflows/pytest.yml

@@ -19,7 +19,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.10","3.12"]
+        python-version: ["3.12"]
 
     steps:
     - uses: actions/checkout@v4

.github/workflows/stat_function_build.yml

@@ -37,7 +37,7 @@ jobs:
       run: |
         pushd './${{ env.AZURE_FUNCTIONAPP_PACKAGE_PATH }}'
         python -m pip install --upgrade pip
-        pip install -r requirements.txt --target=".python_packages/lib/site-packages"
+        pip install -r requirements.txt --platform manylinux_2_17_x86_64 --only-binary=:all: --target=".python_packages/lib/site-packages"
         popd
 
     - name: 'ZIP Function App'

.vscode/launch.json

@@ -3,10 +3,14 @@
     "configurations": [
         {
             "name": "Attach to Python Functions",
-            "type": "python",
+            "type": "debugpy",
             "request": "attach",
-            "port": 9091,
-            "preLaunchTask": "func: host start"
+            "connect": {
+                "host": "localhost",
+                "port": 9091,
+            },
+            "preLaunchTask": "func: host start",
+            "postDebugTask": "Release Blocked Port"
         }
     ]
 }

.vscode/tasks.json

@@ -1,27 +1,33 @@
 {
-	"version": "2.0.0",
-	"tasks": [
-		{
-			"type": "func",
-			"label": "func: host start",
-			"command": "host start",
-			"problemMatcher": "$func-python-watch",
-			"isBackground": true,
-			"dependsOn": "pip install (functions)"
-		},
-		{
-			"label": "pip install (functions)",
-			"type": "shell",
-			"osx": {
-				"command": "${config:azureFunctions.pythonVenv}/bin/python -m pip install -r requirements.txt"
-			},
-			"windows": {
-				"command": "${config:azureFunctions.pythonVenv}\\Scripts\\python -m pip install -r requirements.txt"
-			},
-			"linux": {
-				"command": "${config:azureFunctions.pythonVenv}/bin/python -m pip install -r requirements.txt"
-			},
-			"problemMatcher": []
-		}
-	]
+    "version": "2.0.0",
+    "tasks": [
+        {
+            "type": "func",
+            "label": "func: host start",
+            "command": "host start",
+            "problemMatcher": "$func-python-watch",
+            "isBackground": true,
+            "dependsOn": "pip install (functions)"
+        },
+        {
+            "label": "pip install (functions)",
+            "type": "shell",
+            "osx": {
+                "command": "${config:azureFunctions.pythonVenv}/bin/python -m pip install -r requirements.txt"
+            },
+            "windows": {
+                "command": "${config:azureFunctions.pythonVenv}\\Scripts\\python -m pip install -r requirements.txt"
+            },
+            "linux": {
+                "command": "${config:azureFunctions.pythonVenv}/bin/python -m pip install -r requirements.txt"
+            },
+            "problemMatcher": []
+        },
+        {
+            "label": "Release Blocked Port",
+            "type": "shell",
+            "command": "powershell.exe -ExecutionPolicy Bypass -Command \"Stop-Process -Id (Get-NetTCPConnection -LocalPort 7071).OwningProcess -Force\"",
+            "problemMatcher": []
+        }
+    ]
 }

README.md

@@ -15,7 +15,7 @@ To debug in VS Code create a local.settings.json file in the root of the project
     "AZURE_TENANT_ID": "<TENANTID>",
     "AZURE_CLIENT_ID": "<CLIENTID>",
     "AZURE_CLIENT_SECRET": "<SECRET>",
-    "AZURE_AUTHORITY_HOST": "login.microsoftonline.com",
+    "AZURE_AUTHORITY_HOST": "https://login.microsoftonline.com",
     "ARM_ENDPOINT": "management.azure.com",
     "GRAPH_ENDPOINT": "graph.microsoft.com",
     "LOGANALYTICS_ENDPOINT": "api.loganalytics.io",

classes/__init__.py

@@ -46,6 +46,14 @@ def __init__(self, error:str, source_error:dict={}, status_code:int=400):
         self.source_error = source_error
         self.status_code = status_code
 
+class STATServerError(STATError):
+    """STAT exception raised when an API call returns a 5xx series error.
+
+    This exception is a specialized version of STATError for cases where
+    a server side error was encountered and a retry may be needed.
+    """
+    pass
+
 class STATNotFound(STATError):
     """STAT exception raised when an API call returns a 404 Not Found error.
 
@@ -54,6 +62,14 @@ class STATNotFound(STATError):
     """
     pass
 
+class STATFailedToDecodeToken(STATError):
+    """STAT exception raised when the JWT can't be decoded to check for app roles."""
+    pass
+
+class STATInsufficientPermissions(STATError):
+    """STAT exception raised when the STAT Function identity does not have sufficient permissions."""
+    pass
+
 class STATTooManyRequests(STATError):
     """STAT exception raised when an API call returns a 429 Too Many Requests error.
 
@@ -85,6 +101,7 @@ def __init__(self):
         self.AccountsCount = 0
         self.AccountsOnPrem = []
         self.Alerts = []
+        self.CreatedTime = ''
         self.Domains = []
         self.DomainsCount = 0
         self.EntitiesCount = 0
@@ -99,6 +116,8 @@ def __init__(self):
         self.IncidentARMId = ""
         self.IncidentTriggered = False
         self.IncidentAvailable = False
+        self.MailMessages = []
+        self.MailMessagesCount = 0
         self.ModuleVersions = {}
         self.MultiTenantConfig = {}
         self.OtherEntities = []
@@ -117,6 +136,7 @@ def __init__(self):
     def load_incident_trigger(self, req_body):
 
         self.IncidentARMId = req_body['object']['id']
+        self.CreatedTime = req_body['object']['properties']['createdTimeUtc']
         self.IncidentTriggered = True
         self.IncidentAvailable = True
         self.SentinelRGARMId = "/subscriptions/" + req_body['workspaceInfo']['SubscriptionId'] + "/resourceGroups/" + req_body['workspaceInfo']['ResourceGroupName']
@@ -127,6 +147,7 @@ def load_incident_trigger(self, req_body):
 
     def load_alert_trigger(self, req_body):
         self.IncidentTriggered = False
+        self.CreatedTime = req_body['EndTimeUtc']
         self.SentinelRGARMId = "/subscriptions/" + req_body['WorkspaceSubscriptionId'] + "/resourceGroups/" + req_body['WorkspaceResourceGroup']
         self.WorkspaceId = req_body['WorkspaceId']
 
@@ -135,6 +156,7 @@ def load_from_input(self, basebody):
         self.AccountsCount = basebody['AccountsCount']
         self.AccountsOnPrem = basebody.get('AccountsOnPrem', [])
         self.Alerts = basebody.get('Alerts', [])
+        self.CreatedTime = basebody.get('CreatedTime', '')
         self.Domains = basebody['Domains']
         self.DomainsCount = basebody['DomainsCount']
         self.EntitiesCount = basebody['EntitiesCount']
@@ -149,6 +171,8 @@ def load_from_input(self, basebody):
         self.IncidentTriggered = basebody['IncidentTriggered']
         self.IncidentAvailable = basebody['IncidentAvailable']
         self.IncidentARMId = basebody['IncidentARMId']
+        self.MailMessages = basebody.get('MailMessages', [])
+        self.MailMessagesCount = basebody.get('MailMessagesCount', 0)
         self.ModuleVersions = basebody['ModuleVersions']
         self.MultiTenantConfig = basebody.get('MultiTenantConfig', {})
         self.OtherEntities = basebody['OtherEntities']
@@ -189,11 +213,16 @@ def add_account_entity(self, data):
     def add_onprem_account_entity(self, data):
         self.AccountsOnPrem.append(data)
 
-    def get_ip_list(self):
+    def get_ip_list(self, include_mail_ips:bool=True):
         ip_list = []
         for ip in self.IPs:
             ip_list.append(ip['Address'])
 
+        if include_mail_ips:
+            for message in self.MailMessages:
+                if message.get('senderDetail', {}).get('ipv4'):
+                    ip_list.append(message.get('senderDetail', {}).get('ipv4')) 
+
         return ip_list
 
     def get_domain_list(self):
@@ -203,27 +232,43 @@ def get_domain_list(self):
 
         return domain_list
 
-    def get_url_list(self):
+    def get_url_list(self, include_mail_urls:bool=True):
         url_list = []
         for url in self.URLs:
             url_list.append(url['Url'])
 
+        if include_mail_urls:
+            for message in self.MailMessages:
+                for url in message.get('urls', []):
+                    url_list.append(url.get('url'))
+
         return url_list
 
-    def get_filehash_list(self):
+    def get_filehash_list(self, include_mail_hashes:bool=True):
         hash_list = []
         for hash in self.FileHashes:
             hash_list.append(hash['FileHash'])
+
+        if include_mail_hashes:
+            for message in self.MailMessages:
+                for attachment in message.get('attachments', []):
+                    if attachment.get('sha256'):
+                        hash_list.append(attachment.get('sha256'))
 
         return hash_list
 
-    def get_ip_kql_table(self):
+    def get_ip_kql_table(self, include_mail_ips:bool=True):
 
         ip_data = []
 
         for ip in self.IPs:
             ip_data.append({'Address': ip.get('Address'), 'Latitude': ip.get('GeoData').get('latitude'), 'Longitude': ip.get('GeoData').get('longitude'), \
                             'Country': ip.get('GeoData').get('country'), 'State': ip.get('GeoData').get('state')})
+
+        if include_mail_ips:
+            for message in self.MailMessages:
+                if message.get('senderDetail', {}).get('ipv4'):
+                    ip_data.append({'Address': message.get('senderDetail', {}).get('ipv4')})          
 
         encoded = urllib.parse.quote(json.dumps(ip_data))
 
@@ -268,12 +313,17 @@ def get_host_kql_table(self):
 '''
         return kql
 
-    def get_url_kql_table(self):
+    def get_url_kql_table(self, include_mail_urls:bool=True):
         url_data = []
 
         for url in self.URLs:
             url_data.append({'Url': url.get('Url')})
 
+        if include_mail_urls:
+            for message in self.MailMessages:
+                for url in message.get('urls', []):
+                    url_data.append({'Url': url.get('url')})
+
         encoded = urllib.parse.quote(json.dumps(url_data))
 
         kql = f'''let urlEntities = print t = todynamic(url_decode('{encoded}'))
@@ -282,12 +332,18 @@ def get_url_kql_table(self):
 '''
         return kql
 
-    def get_filehash_kql_table(self):
+    def get_filehash_kql_table(self, include_mail_hashes:bool=True):
         hash_data = []
 
         for hash in self.FileHashes:
             hash_data.append({'FileHash': hash.get('FileHash'), 'Algorithm': hash.get('Algorithm')})
 
+        if include_mail_hashes:
+            for message in self.MailMessages:
+                for attachment in message.get('attachments', []):
+                    if attachment.get('sha256'):
+                        hash_data.append({'FileHash': attachment.get('sha256'), 'Algorithm': 'SHA256'})
+
         encoded = urllib.parse.quote(json.dumps(hash_data))
 
         kql = f'''let hashEntities = print t = todynamic(url_decode('{encoded}'))
@@ -308,6 +364,21 @@ def get_domain_kql_table(self):
         kql = f'''let domainEntities = print t = todynamic(url_decode('{encoded}'))
 | mv-expand t
 | project Domain=tostring(t.Domain);
+'''
+        return kql
+
+    def get_mail_kql_table(self):
+
+        mail_data = []
+
+        for mail in self.MailMessages:
+            mail_data.append({'rec': mail.get('recipientEmailAddress'), 'nid': mail.get('networkMessageId'), 'send': mail.get('senderDetail', {}).get('fromAddress'), 'sendfrom': mail.get('senderDetail', {}).get('mailFromAddress')})
+
+        encoded = urllib.parse.quote(json.dumps(mail_data))
+
+        kql = f'''let mailEntities = print t = todynamic(url_decode('{encoded}'))
+| mv-expand t
+| project RecipientEmailAddress=tostring(t.rec), NetworkMessageId=tostring(t.nid), SenderMailFromAddress=tostring(t.send), SenderFromAddress=tostring(t.sendfrom);
 '''
         return kql
 

debug/debug.py

@@ -26,6 +26,12 @@ def debug_module (req_body):
         case 'comment':
             default_debug(debug_out)
             comment_debug(debug_out)
+        case 'tag':
+            default_debug(debug_out)
+            tag_debug(debug_out)
+        case 'task':
+            default_debug(debug_out)
+            task_debug(debug_out)
         case 'exception':
             exception_debug(debug_out)
         case _:
@@ -145,6 +151,42 @@ def comment_debug(debug_out:DebugModule):
     base_module.MultiTenantConfig = debug_out.Params.get('MultiTenantConfig', {})
     base_module.IncidentARMId = debug_out.Params['IncidentARMId']
     base_module.IncidentAvailable = True
-    response = rest.add_incident_comment(base_module, comment)
-    debug_out.CommentStatus = response.status_code
-    debug_out.CommentResponse = response.json()
+    try: 
+        response = rest.add_incident_comment(base_module, comment, raise_on_error=True)
+        debug_out.CommentStatus = response.status_code
+        debug_out.CommentResponse = response.json()
+    except STATError as e:
+        debug_out.CommentStatus = 'Comment failed'
+        debug_out.CommentSourceError = e.source_error
+        debug_out.CommentStatusCode = e.status_code
+
+def task_debug(debug_out:DebugModule):
+    title = debug_out.Params.get('TaskTitle')
+    description = debug_out.Params.get('TaskDescription', '')
+    base_module = BaseModule()
+    base_module.MultiTenantConfig = debug_out.Params.get('MultiTenantConfig', {})
+    base_module.IncidentARMId = debug_out.Params['IncidentARMId']
+    base_module.IncidentAvailable = True
+    try: 
+        response = rest.add_incident_task(base_module, title, description, raise_on_error=True)
+        debug_out.TaskStatus = response.status_code
+        debug_out.TaskResponse = response.json()
+    except STATError as e:
+        debug_out.TaskStatus = 'Task failed'
+        debug_out.TaskSourceError = e.source_error
+        debug_out.TaskStatusCode = e.status_code
+
+def tag_debug(debug_out:DebugModule):
+    tag = [debug_out.Params.get('Tag')]
+    base_module = BaseModule()
+    base_module.MultiTenantConfig = debug_out.Params.get('MultiTenantConfig', {})
+    base_module.IncidentARMId = debug_out.Params['IncidentARMId']
+    base_module.IncidentAvailable = True
+    try: 
+        response = rest.add_incident_tags(base_module, tag, raise_on_error=True)
+        debug_out.TagStatus = response.status_code
+        debug_out.TagResponse = response.json()
+    except STATError as e:
+        debug_out.TagStatus = 'Tag failed'
+        debug_out.TagSourceError = e.source_error
+        debug_out.TagStatusCode = e.status_code

host.json

@@ -10,6 +10,6 @@
   },
   "extensionBundle": {
     "id": "Microsoft.Azure.Functions.ExtensionBundle",
-    "version": "[3.*, 4.0.0)"
+    "version": "[4.0.0, 5.0.0)"
   }
 }

modules/aadrisks.py

@@ -1,11 +1,15 @@
 from classes import BaseModule, Response, AADModule, STATError, STATNotFound
 from shared import rest, data
 import json, datetime
+import logging
 
 def execute_aadrisks_module (req_body):
 
     #Inputs AddIncidentComments, AddIncidentTask, Entities, IncidentTaskInstructions, LookbackInDays, MFAFailureLookup, MFAFraudLookup, SuspiciousActivityReportLookup
-
+    # Log module invocation with parameters (excluding BaseModuleBody)
+    log_params = {k: v for k, v in req_body.items() if k != 'BaseModuleBody'}
+    logging.info(f'AAD Risks Module invoked with parameters: {log_params}')
+
     base_object = BaseModule()
     base_object.load_from_input(req_body['BaseModuleBody'])