{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "import os\r\n",
    "from datetime import datetime\r\n",
    "from pathlib import Path\r\n",
    "import logging\r\n",
    "\r\n",
    "import nbformat\r\n",
    "from azure.mgmt.resource import ResourceManagementClient\r\n",
    "from azure.storage.fileshare import ShareFileClient, ShareServiceClient\r\n",
    "from azure.storage.queue import QueueServiceClient\r\n",
    "\r\n",
    "import papermill as pm\r\n",
    "from msticpy.common.azure_auth import az_connect\r\n",
    "from msticpy.common.keyvault_client import BHKeyVaultClient\r\n",
    "from msticpy.data.azure_sentinel import AzureSentinel\r\n",
    "from azure.common.exceptions import CloudError"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "# Open a log file that tracks notebook executions\r\n",
    "path = Path.cwd().joinpath(\"notebook_execution.log\")\r\n",
    "logging.basicConfig(filename=path)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "# Populate with details relating to your environment\r\n",
    "# Tenant ID\r\n",
    "ten_id = \"YOUR TENANT ID\"\r\n",
    "# The name of the Key Vault containing AFS key\r\n",
    "vault_name = \"YOUR KV NAME\"\r\n",
    "# The secret name that the AFS key is stored in\r\n",
    "kv_sec_name = \"YOUR SECRET NAME\"\r\n",
    "# Subscription ID of the Azure Sentinel Workspace to get incidents from\r\n",
    "subscriptionId = \"YOUR SUBSCRIPTION ID\"\r\n",
    "# The name of the Resource Group of the Azure Sentinel Workspace to get incidents from\r\n",
    "resourceGroupName = \"YOUR RG NAME\"\r\n",
    "# The name of the Azure Sentinel Workspace to get incidents from\r\n",
    "workspaceName = \"YOUR WORKSPACE NAME\"\r\n",
    "# The name of the Azure Sentinel Workspace ID to get incidents from\r\n",
    "ws_id = \"YOUR WORKSPACE ID\"\r\n",
    "# The name of the Azure Storage Queue account used (if used)\r\n",
    "q_account = \"YOUR QUEUE ACCOUNT\"\r\n",
    "# The name of the Azure Storage Queue account used (if used)\r\n",
    "q_name = \"YOUR QUEUE NAME\"\r\n",
    "# Details of the Azure Machine Learning workspace to be used (sub_id = Subscription ID, RG = Resource Group name, AMLWorkspace = AML Workspace Name)\r\n",
    "AML_details = {\r\n",
    "    \"sub_id\": \"YOUR SUB ID\",\r\n",
    "    \"RG\": \"YOUR RG NAME\",\r\n",
    "    \"AMLWorkspace\": \"YOUR WORKSPACE NAME\",\r\n",
    "    \"ten_id\": ten_id,\r\n",
    "}"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "# Authenticate to Azure using Azure CLI or Managed Identity\r\n",
    "creds = az_connect([\"cli\", \"msi\"])\r\n",
    "token = creds.modern.get_token(\"https://management.azure.com/.default\")\r\n",
    "\r\n",
    "\r\n",
    "def get_api_headers():\r\n",
    "    token = creds.modern.get_token(\"https://management.azure.com/.default\")\r\n",
    "    \"\"\"Return authorization header with current token.\"\"\"\r\n",
    "    return {\r\n",
    "        \"Authorization\": f\"Bearer {token.token}\",\r\n",
    "        \"Content-Type\": \"application/json\",\r\n",
    "    }\r\n",
    "\r\n",
    "\r\n",
    "# Access key vault and get Azure Storage access key\r\n",
    "kv_c = BHKeyVaultClient(tenant_id=ten_id, vault_name=vault_name)\r\n",
    "afs_cred = kv_c.get_secret(kv_sec_name)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "# Connect to Azure Sentinel\r\n",
    "azs = AzureSentinel()\r\n",
    "azs.connect()\r\n",
    "logging.info(f\"{datetime.now()} - Successfully connected to Azure Sentinel\")"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "# Get recent Incidents from API\r\n",
    "try:\r\n",
    "    incidents = azs.get_incidents(\r\n",
    "        sub_id=subscriptionId, res_grp=resourceGroupName, ws_name=workspaceName\r\n",
    "    )\r\n",
    "    incident_ids = incidents[\"name\"].tolist()\r\n",
    "except CloudError:\r\n",
    "    logging.info(f\"{datetime.now()} - Unable to retreive incidents\")\r\n",
    "    incident_ids = []"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "# If using Queue method get incidents from queue - uncomment following cells to use this method\r\n",
    "# queue_service_client = QueueServiceClient(\r\n",
    "# account_url=q_account, credential=creds.modern, api_version=\"2019-07-07\"\r\n",
    "# )\r\n",
    "# q_client = queue_service_client.get_queue_client(q_name)\r\n",
    "# messages = q_client.receive_messages()\r\n",
    "# logging.info(f\"{datetime.now()} - Getting incidents from queue\")\r\n",
    "# incident_ids = [message[\"content\"] for message in messages]\r\n",
    "# q_client.clear_messages()"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "# Use a local archive to avoid processing of incidents more than once.\r\n",
    "try:\r\n",
    "    with open(\"incident_archive\", \"r\") as input_file:\r\n",
    "        incident_archive = input_file.read().splitlines()\r\n",
    "except FileNotFoundError:\r\n",
    "    incident_archive = []"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "incident_file = open(\"incident_archive\", \"w\")\r\n",
    "out_files = []\r\n",
    "\r\n",
    "# For each incident, if it has not already been processed then run the incident triage notebook with that ID\r\n",
    "if incident_ids:\r\n",
    "    for incident_id in incident_ids:\r\n",
    "        if incident_id not in incident_archive:\r\n",
    "            logging.info(f\"{datetime.now()} - Processing incident {incident_id}\")\r\n",
    "            (Path.cwd() / 'out').mkdir(parents=True, exist_ok=True)\r\n",
    "            out_path = Path.cwd().joinpath(f\"out/{incident_id}.ipynb\")\r\n",
    "            # If execution error occurs continue but record this in the log\r\n",
    "            try:\r\n",
    "                pm.execute_notebook(\r\n",
    "                    \"AutomatedNotebooks-IncidentTriage.ipynb\",\r\n",
    "                    str(out_path),\r\n",
    "                    parameters={\r\n",
    "                        \"incident_id\": incident_id,\r\n",
    "                        \"ten_id\": ten_id,\r\n",
    "                        \"ws_id\": ws_id,\r\n",
    "                    }\r\n",
    "                )\r\n",
    "                out_files.append(out_path)\r\n",
    "            except pm.PapermillExecutionError:\r\n",
    "                logging.info(\r\n",
    "                    f\"{datetime.now()} - Unable to process incident {incident_id} - skipping\"\r\n",
    "                )\r\n",
    "            # Once processed add incident to archive\r\n",
    "            incident_file.write(incident_id + \"\\n\")\r\n",
    "        else:\r\n",
    "            logging.info(\r\n",
    "                f\"{datetime.now()} - Incident {incident_id} has already been processed - skipping\"\r\n",
    "            )\r\n",
    "\r\n",
    "incident_file.close()"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "# Function to move a notebook from local path to Azure File Storage\r\n",
    "def move_to_afs(path, incident_id):\r\n",
    "    with open(path) as notebook:\r\n",
    "        notebook = notebook.read()\r\n",
    "        account = get_storage_acct()\r\n",
    "        notebook_name = path.name\r\n",
    "        share_name = get_share(account)\r\n",
    "        file_client = ShareFileClient(\r\n",
    "            account_url=f\"{account}.file.core.windows.net\",\r\n",
    "            share_name=share_name,\r\n",
    "            file_path=f\"Users/TriageNbs/{notebook_name}\",\r\n",
    "            credential=afs_cred,\r\n",
    "        )\r\n",
    "        file_client.upload_file(notebook)\r\n",
    "        path = (f\"https://ml.azure.com/fileexplorerAzNB?wsid=/subscriptions/{AML_details['sub_id']}\"\r\n",
    "                f\"/resourcegroups/{AML_details['RG']}/workspaces/AzureMLWorkspace&tid={AML_details['ten_id']}\"\r\n",
    "                f\"&activeFilePath=Users/{notebook_name}\")\r\n",
    "        write_to_incident(incident_id, path)\r\n",
    "        update_incident(incident_id)\r\n",
    "\r\n",
    "\r\n",
    "# Function to find the Azure Storage Account used by Azure ML\r\n",
    "def get_storage_acct():\r\n",
    "    res_client = ResourceManagementClient(creds.legacy, AML_details[\"sub_id\"])\r\n",
    "    res = res_client.resources.get(\r\n",
    "        AML_details[\"RG\"],\r\n",
    "        \"\",\r\n",
    "        \"Microsoft.MachineLearningServices/workspaces\",\r\n",
    "        \"\",\r\n",
    "        AML_details[\"AMLWorkspace\"],\r\n",
    "        \"2021-01-01\",\r\n",
    "    )\r\n",
    "    account = res.properties[\"storageAccount\"].split(\"/\")[-1]\r\n",
    "    return account\r\n",
    "\r\n",
    "\r\n",
    "# Function to get the correct file share to store notebook in\r\n",
    "def get_share(account):\r\n",
    "    ssc = ShareServiceClient(f\"{account}.file.core.windows.net\", afs_cred)\r\n",
    "    for share in list(ssc.list_shares()):\r\n",
    "        if share[\"name\"].startswith(\"code-\"):\r\n",
    "            return share[\"name\"]\r\n",
    "\r\n",
    "\r\n",
    "# Function to write a comment to the Azure Sentinel Incident that contains a link to the notebook\r\n",
    "def write_to_incident(incidentId, path):\r\n",
    "    html = f\"<a href='{path}'>View incident triage notebook in AML</a>\"\r\n",
    "    logging.info(\r\n",
    "        f\"{datetime.now()} - Adding link to notebook for incident: {incidentId}\"\r\n",
    "    )\r\n",
    "    azs.post_comment(\r\n",
    "        incident_id=incidentId,\r\n",
    "        comment=html,\r\n",
    "        sub_id=subscriptionId,\r\n",
    "        res_grp=resourceGroupName,\r\n",
    "        ws_name=workspaceName,\r\n",
    "    )\r\n",
    "\r\n",
    "\r\n",
    "# Function to update incident severity to High if triage determines a risk\r\n",
    "def update_incident(incidentId):\r\n",
    "    logging.info(f\"{datetime.now()} - Updating severity for {incidentId}\")\r\n",
    "    azs.update_incident(\r\n",
    "        incident_id=incidentId,\r\n",
    "        update_items={\"severity\": \"High\", \"status\": \"New\"},\r\n",
    "        sub_id=subscriptionId,\r\n",
    "        res_grp=resourceGroupName,\r\n",
    "        ws_name=workspaceName,\r\n",
    "    )\r\n",
    "\r\n",
    "\r\n",
    "# For each processed incident check if there was valuable output in the notebook and process it\r\n",
    "if out_files:\r\n",
    "    for path in out_files:\r\n",
    "        incident_id = str(path.name).split(\".\")[0]\r\n",
    "        try:\r\n",
    "            nb = nbformat.read(path, as_version=2)\r\n",
    "            for cell in nb[\"worksheets\"][0][\"cells\"]:\r\n",
    "                if \"output\" in cell[\"metadata\"][\"tags\"] and cell[\"outputs\"]:\r\n",
    "                    logging.info(\r\n",
    "                        f\"{datetime.now()} - Storing notebook for {incident_id} in AML\"\r\n",
    "                    )\r\n",
    "                    move_to_afs(path, incident_id)\r\n",
    "                    break\r\n",
    "            os.remove(str(path))\r\n",
    "        except FileNotFoundError:\r\n",
    "            continue"
   ],
   "outputs": [],
   "metadata": {}
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "papermill",
   "language": "python",
   "name": "papermill"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.12"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}