Guided Hunting - Detect potential network beaconing using Apache Spark via Azure Synapse.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "source": [
        "## Guided Hunting - Detect potential network beaconing using Apache Spark via Azure Synapse\n",
        "\n",
        "__Notebook Version:__ 1.0<br>\n",
        "__Python Version:__ Python 3.10 - AzureML<br>\n",
        "__Required Packages:__ azureml-synapse, Msticpy, azure-storage-file-datalake <br>\n",
        "__Platforms Supported:__  Azure Machine Learning Notebooks connected to Azure Synapse Workspace\n",
        "     \n",
        "__Data Source Required:__ Yes\n",
        "\n",
        "__Data Source:__ CommonSecurityLogs\n",
        "\n",
        "__Spark Version:__ 3.1 or above\n",
        "    \n",
        "### Description\n",
        "In this sample guided scenario notebook, we will demonstrate how to set up continuous data pipeline to store data into azure data lake storage (ADLS) and \n",
        "then hunt on that data at scale using distributed processing via Azure Synapse workspace connected to serverless Spark pool. \n",
        "Once historical dataset is available in ADLS , we can start performing common hunt operations, create a baseline of normal behavior using PySpark API and also apply data transformations \n",
        "to find anomalous behaviors such as periodic network beaconing as explained in the blog - [Detect Network beaconing via Intra-Request time delta patterns in Microsoft Sentinel - Microsoft Tech Community](https://techcommunity.microsoft.com/t5/azure-sentinel/detect-network-beaconing-via-intra-request-time-delta-patterns/ba-p/779586). \n",
        "You can use various other spark API to perform other data transformation to understand the data better. \n",
        "The output generated can also be further enriched to populate Geolocation information and also visualize using Msticpy capabilities to identify any anomalies. \n",
        ".<br>\n",
        "*** Python modules download may be needed. ***<br>\n",
        "*** Please run the cells sequentially to avoid errors.  Please do not use \"run all cells\". *** <br>\n",
        "\n",
        "## Table of Contents\n",
        "1. Warm-up\n",
        "2. Authentication to Azure Resources\n",
        "3. Configure Azure ML and Azure Synapse Analytics\n",
        "4. Load the Historical and current data\n",
        "5. Data Wrangling using Spark\n",
        "6. Enrich the results\n",
        "7. Conclusion\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "source": [
        "### Warm-up"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "source": [
        "> **Note**: Install below packages only for the first time and restart the kernel once done."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "gather": {
          "logged": 1632406406186
        },
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "outputs": [],
      "source": [
        "# Install AzureML Synapse package to use spark magics\n",
        "import sys\n",
        "!{sys.executable} -m pip install azureml-synapse"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "outputs": [],
      "source": [
        "# Install Azure storage datalake library to manipulate file systems\n",
        "import sys\n",
        "!{sys.executable} -m pip install azure-storage-file-datalake --pre"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "outputs": [],
      "source": [
        "# Install Azure storage datalake library to manipulate file systems\n",
        "import sys\n",
        "!{sys.executable} -m pip install msticpy"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "source": [
        "*** $\\color{red}{Note:~After~installing~the~packages,~please~restart~the~kernel.}$ ***"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "gather": {
          "logged": 1632982861159
        },
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "outputs": [],
      "source": [
        "# Load Python libraries that will be used in this notebook\n",
        "from azure.common.client_factory import get_client_from_cli_profile\n",
        "from azure.common.credentials import get_azure_cli_credentials\n",
        "from azure.mgmt.resource import ResourceManagementClient\n",
        "# from azureml.core import Workspace, LinkedService, SynapseWorkspaceLinkedServiceConfiguration, Datastore\n",
        "# from azureml.core.compute import SynapseCompute, ComputeTarget\n",
        "from datetime import timedelta, datetime\n",
        "# from azure.storage.filedatalake import DataLakeServiceClient\n",
        "# from azure.core._match_conditions import MatchConditions\n",
        "# from azure.storage.filedatalake._models import ContentSettings\n",
        "\n",
        "import json\n",
        "import os, uuid, sys\n",
        "import IPython\n",
        "import pandas as pd\n",
        "from ipywidgets import widgets, Layout\n",
        "from IPython.display import display, HTML\n",
        "from pathlib import Path\n",
        "\n",
        "REQ_PYTHON_VER = \"3.10\"\n",
        "REQ_MSTICPY_VER = \"2.12.0\"\n",
        "\n",
        "\n",
        "# If not using Azure Notebooks, install msticpy with\n",
        "# !pip install msticpy\n",
        "\n",
        "from msticpy import init_notebook\n",
        "extra_imports = [\n",
        "    \"msticpy.nbtools.nbdisplay, draw_alert_entity_graph\",\n",
        "    \"msticpy.sectools.ip_utils, convert_to_ip_entities\",\n",
        "    \"msticpy.nbtools.ti_browser, browse_results\",\n",
        "    \"IPython.display, Image\",\n",
        "    \"msticpy.sectools.ip_utils, get_whois_info\",\n",
        "    \"msticpy.sectools.ip_utils, get_ip_type\"\n",
        "]\n",
        "init_notebook(\n",
        "    namespace=globals(),\n",
        "    # additional_packages=[\"azureml-synapse\", \"azure-cli\", \"azure-storage-file-datalake\"],\n",
        "    extra_imports=extra_imports,\n",
        ");\n",
        "\n",
        "\n",
        "WIDGET_DEFAULTS = {\n",
        "    \"layout\": Layout(width=\"95%\"),\n",
        "    \"style\": {\"description_width\": \"initial\"},\n",
        "}\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "source": [
        "## Configure Azure ML and Azure Synapse Analytics\n",
        "\n",
        "Please use notebook [Configurate Azure ML and Azure Synapse Analytics](https://github.com/Azure/Azure-Sentinel-Notebooks/blob/master/Configurate%20Azure%20ML%20and%20Azure%20Synapse%20Analytics.ipynb) to configure environment. \n",
        "\n",
        "The notebook will configure existing Azure synapse workspace to create and connect to Spark pool. You can then create linked service and connect AML workspace to Azure Synapse workspaces.\n",
        "<br>It will also configure data export rules to export data from Log analytics workspace CommonSecurityLog table to Azure Data lake storage Gen 2.\n",
        "\n",
        "> **Note**: Specify the input parameters in below step in order to connect AML workspace to synapse workspace using linked service."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "outputs": [],
      "source": [
        "amlworkspace = \"<aml workspace name>\"  # fill in your AML workspace name\n",
        "subscription_id = \"<subscription id>\" # fill in your subscription id\n",
        "resource_group = '<resource group of AML workspace>' # fill in your resource groups for AML workspace\n",
        "linkedservice = '<linked service name>' # fill in your linked service created to connect to synapse workspace"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "source": [
        "### Authentication to Azure Resources\n",
        "\n",
        "In this step we will connect aml workspace to linked service connected to Azure Synapse workspace"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "gather": {
          "logged": 1632982882257
        },
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "outputs": [],
      "source": [
        "# Get the aml workspace\n",
        "aml_workspace = Workspace.get(name=amlworkspace, subscription_id=subscription_id, resource_group=resource_group)\n",
        "\n",
        "# Retrieve a known linked service\n",
        "linked_service = LinkedService.get(aml_workspace, linkedservice)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "source": [
        "## Start Spark Session\n",
        "Enter your Synapse Spark compute below. To find the Spark compute, please follow these steps:  </br>\n",
        "1. On the AML Studio left menu, navigate to **Linked Services**  </br>\n",
        "2. Click on the name of the Link Service you want to use </br>\n",
        "3. Select **Spark pools** tab </br>\n",
        "4. Get the Name of the Spark pool you want to use. </br>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "outputs": [],
      "source": [
        "synapse_spark_compute = input('Synapse Spark compute:')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "gather": {
          "logged": 1632959931295
        },
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "outputs": [],
      "source": [
        "# Start spark session\n",
        "%synapse start -s $subscription_id -w $amlworkspace -r $resource_group -c $synapse_spark_compute"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "source": [
        "## Data Preparation\n",
        "In this step, we will define several details associated with ADLS account and specify input date and lookback period to calculate baseline. Based on the input dates and lookback period , we will load the data.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "outputs": [],
      "source": [
        "%%synapse\n",
        "\n",
        "# Primary storage info\n",
        "account_name = '<storage account name>' # fill in your primary account name\n",
        "container_name = '<container name>' # fill in your container name\n",
        "subscription_id = '<subscription if>' # fill in your subscription id\n",
        "resource_group = '<resource group>' # fill in your resource groups for ADLS\n",
        "workspace_name = '<Microsoft sentinel/log analytics workspace name>' # fill in your workspace name\n",
        "device_vendor = \"Fortinet\"  # Replace your desired network vendor from commonsecuritylogs\n",
        "\n",
        "# Datetime and lookback parameters\n",
        "end_date = \"<enter date in the format yyyy-MM-dd e.g.2021-09-17>\"  # fill in your input date\n",
        "lookback_days = 21 # fill in lookback days if you want to run it on historical data. make sure you have historical data available in ADLS"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "gather": {
          "logged": 1632245154374
        },
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "outputs": [],
      "source": [
        "%%synapse\n",
        "from pyspark.sql.types import *\n",
        "from pyspark.sql.window import Window\n",
        "from pyspark.sql.functions import lag, col\n",
        "from pyspark.sql.functions import *\n",
        "from pyspark.sql import functions as F\n",
        "from datetime import timedelta, datetime, date\n",
        "\n",
        "# Compiling ADLS paths from date string\n",
        "end_date_str = end_date.split(\"-\")\n",
        "current_path = f\"/y={end_date_str[0]}/m={end_date_str[1]}/d={end_date_str[2]}\"\n",
        "\n",
        "def generate_adls_paths(end_date, lookback_days, adls_path):\n",
        "    endDate = datetime.strptime(end_date, '%Y-%m-%d')\n",
        "    endDate = endDate - timedelta(days=1)\n",
        "    startDate = endDate - timedelta(days=lookback_days)\n",
        "    daterange = [startDate + timedelta(days=x) for x in range((endDate-startDate).days + 1)]\n",
        "\n",
        "    pathlist = []\n",
        "    for day in daterange:\n",
        "        date_str = day.strftime('%Y-%m-%d').split(\"-\")\n",
        "        day_path = adls_path + f\"/y={date_str[0]}/m={date_str[1]}/d={date_str[2]}\"\n",
        "        pathlist.append(day_path)\n",
        "\n",
        "    return pathlist\n",
        "\n",
        "adls_path = f'abfss://{container_name}@{account_name}.dfs.core.windows.net/WorkspaceResourceId=/subscriptions/{subscription_id}/resourcegroups/{resource_group}/providers/microsoft.operationalinsights/workspaces/{workspace_name}'\n",
        "current_day_path = adls_path + current_path\n",
        "historical_path = generate_adls_paths(end_date, lookback_days, adls_path)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "source": [
        "### Load Current day\n",
        "\n",
        "In this step, you will load the data based on the input date specified."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "outputs": [],
      "source": [
        "%%synapse\n",
        "\n",
        "try:\n",
        "    current_df = (\n",
        "        spark.read.option(\"recursiveFileLook\", \"true\")\n",
        "        .option(\"header\", \"true\")\n",
        "        .json(current_day_path)\n",
        "    )\n",
        "\n",
        "    current_df = (\n",
        "            current_df\n",
        "            .select(\n",
        "                \"TimeGenerated\",\n",
        "                \"SourceIP\",\n",
        "                \"SourcePort\",\n",
        "                \"DestinationIP\",\n",
        "                \"DestinationPort\",\n",
        "                \"Protocol\",\n",
        "                \"ReceivedBytes\",\n",
        "                \"SentBytes\",\n",
        "                \"DeviceVendor\",\n",
        "                    )\n",
        "            .filter(F.col(\"DeviceVendor\") == device_vendor)\n",
        "                )\n",
        "except Exception as e:\n",
        "    print(f\"Could note load the data due to error:{e}\")\n",
        "\n",
        "#Display the count of records\n",
        "print(f\"No of records loaded from the current date specified: {current_df.count()}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "source": [
        "### Load Historical data\n",
        "\n",
        "You can also perform the analysis on all historical data available in your ADLS account. The notebook is currently configured to run only on current date specified in input. \n",
        "\n",
        "If you need to perform the same analysis on historical data, run the cell below and under Data Wrangling using Spark -> Filtering Data code cell, replace `current_df` with `historical_df` variable. \n",
        "<br>**Otherwise SKIP running below cell as it will result in an error if you do not have historical data**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "outputs": [],
      "source": [
        "%%synapse\n",
        "\n",
        "try:\n",
        "    #Read Previous days data\n",
        "    historical_df = (\n",
        "        spark.read.option(\"recursiveFileLook\", \"true\")\n",
        "        .option(\"header\", \"true\")\n",
        "        .json(historical_path[-1])\n",
        "    )\n",
        "    historical_df = historical_df.select(\n",
        "        \"TimeGenerated\",\n",
        "        \"SourceIP\",\n",
        "        \"SourcePort\",\n",
        "        \"DestinationIP\",\n",
        "        \"DestinationPort\",\n",
        "        \"Protocol\",\n",
        "        \"ReceivedBytes\",\n",
        "        \"SentBytes\",\n",
        "        \"DeviceVendor\",\n",
        "    ).filter(F.col(\"DeviceVendor\") == device_vendor)\n",
        "\n",
        "    #Read all historical days data per day and union them together\n",
        "    for path in historical_path[:-1]:\n",
        "        daily_table = (\n",
        "            spark.read.option(\"recursiveFileLook\", \"true\")\n",
        "            .option(\"header\", \"true\")\n",
        "            .json(path)\n",
        "        )\n",
        "        daily_table = daily_table.select(\n",
        "            \"TimeGenerated\",\n",
        "            \"SourceIP\",\n",
        "            \"SourcePort\",\n",
        "            \"DestinationIP\",\n",
        "            \"DestinationPort\",\n",
        "            \"Protocol\",\n",
        "            \"ReceivedBytes\",\n",
        "            \"SentBytes\",\n",
        "            \"DeviceVendor\",\n",
        "        ).filter(F.col(\"DeviceVendor\") == device_vendor)\n",
        "        historical_df = historical_df.union(daily_table)\n",
        "\n",
        "except Exception as e:\n",
        "    print(f\"Could not load the data due to error:\\n {e}\")\n",
        "\n",
        "#Display the count of records\n",
        "print(f\"\\n No of records loaded from the lookback days specified: {historical_df.count()}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "source": [
        "## Data Wrangling using Spark"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "source": [
        "### Filtering data\n",
        "In this step, we will prepare dataset by filtering logs to only destination as Public/external IPs. For this, we are using regex and rlike spark API to filter internal to external network traffic."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "outputs": [],
      "source": [
        "%%synapse\n",
        "\n",
        "PrivateIPregex = (\"^127\\.|^10\\.|^172\\.1[6-9]\\.|^172\\.2[0-9]\\.|^172\\.3[0-1]\\.|^192\\.168\\.\")\n",
        "cooked_df = (current_df   # replace historical_df if you want to use historical data\n",
        "            # use below filter if you have Palo Alto logs\n",
        "            # .filter(\n",
        "            #     (F.col(\"Activity\") == \"TRAFFIC\")\n",
        "            #     )\n",
        "            .withColumn(\n",
        "            \"DestinationIsPrivate\", F.col(\"DestinationIP\").rlike(PrivateIPregex)\n",
        "                        )\n",
        "            .filter(F.col(\"DestinationIsPrivate\") == \"false\")\n",
        "            .withColumn(\"TimeGenerated\", F.col(\"TimeGenerated\").cast(\"timestamp\"))\n",
        "            )\n",
        "\n",
        "cooked_df.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "source": [
        "### Baseline data to filter known Source IP and Destination IPs\n",
        "\n",
        "In this step, you can either analyze Historical data or current data to filter source IP and destination IP per defined criteria. \n",
        "\n",
        "In below example, we are filtering the Source IP which has daily event count more than the specified threshold.\n",
        "<br> Also, you can filter the destination IPs whom very less source IPs are connecting. This will reduce false positives be filtering destination IPs which are commonly seen from internal systems which are likely benign."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "outputs": [],
      "source": [
        "%%synapse\n",
        "daily_event_count_threshold = 100  # Replace the threshold based on your environment or use default values\n",
        "degree_of_srcip_threshold = 25    # Replace the threshold based on your environment or use default values\n",
        "\n",
        "# Filtering IP list per TotalEventsThreshold\n",
        "csl_srcip = (\n",
        "        cooked_df.groupBy(\"SourceIP\")\n",
        "        .count()\n",
        "        .filter(F.col(\"count\") > daily_event_count_threshold)\n",
        "        .orderBy(F.col(\"count\"), ascending=False)\n",
        "    )\n",
        "\n",
        "# Filtering Destination IP list per Degree of Source IPs threshold\n",
        "csl_dstip = (\n",
        "        cooked_df.groupBy(\"DestinationIP\")\n",
        "        .agg(F.countDistinct(\"SourceIP\").alias(\"DegreeofSourceIps\"))\n",
        "        .filter(F.col(\"DegreeofSourceIps\") < degree_of_srcip_threshold)\n",
        "    )\n",
        "\n",
        "# Filtering IP list per Daily event threshold\n",
        "baseline_df = (\n",
        "        cooked_df.join(csl_srcip, [\"SourceIP\"])\n",
        "        .join(csl_dstip, [\"DestinationIP\"])\n",
        "        .select(\n",
        "            \"TimeGenerated\",\n",
        "            \"SourceIP\",\n",
        "            \"SourcePort\",\n",
        "            \"DestinationIP\",\n",
        "            \"DestinationPort\",\n",
        "            \"Protocol\",\n",
        "            \"ReceivedBytes\",\n",
        "            \"SentBytes\",\n",
        "            \"DeviceVendor\",\n",
        "        )\n",
        "    )\n",
        "\n",
        "baseline_df.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "source": [
        "### Rank the datasets and Calculate PercentageBeaconing\n",
        "\n",
        "In this step, we will use spark to wrangle the data by applying below transformations.\n",
        "- Sort the dataset per SourceIP\n",
        "- Calculate the time difference between First event and next event.\n",
        "- Partition dataset per Source IP, Destination IP, Destination Port\n",
        "- Window dataset into consecutive 3 to Calculate the Timedeltalistcount based on cluster of 1-3 timedelta events.\n",
        "- Calculate percentagebeacon between TotalEventscount and Timedeltalistcount\n",
        "- Apply thresholds to further reduce false positives\n",
        "\n",
        "** SPARK References:**\n",
        "- https://spark.apache.org/docs/latest/api/python/getting_started/quickstart.html\n",
        "- https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql.html#window"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "outputs": [],
      "source": [
        "%%synapse\n",
        "time_delta_threshold = 15   # Replace thresholds in seconds interval between 2 successive events. min 10 to anything maximum.\n",
        "percent_beacon_threshold = 75  # Replace thresholds in percentage. Max value can be 100.\n",
        "\n",
        "# Serialize the dataset by sorting timegenerated and partition by SourceIP and WorkspaceId\n",
        "w = (\n",
        "        Window()\n",
        "        .partitionBy(F.col(\"SourceIP\"))\n",
        "        .orderBy(F.col(\"TimeGenerated\"))\n",
        "    )\n",
        "\n",
        "# Calculate new timestamp column of next event\n",
        "csl_beacon_df = baseline_df.select(\n",
        "        \"*\", lag(\"TimeGenerated\").over(w).alias(\"prev_TimeStamp\")\n",
        "    ).na.drop()\n",
        "\n",
        "# Calculate timedelta difference between previoud and next timestamp\n",
        "timeDiff = F.unix_timestamp(\"TimeGenerated\") - F.unix_timestamp(\"prev_TimeStamp\")\n",
        "\n",
        "# Add new column as timedelta\n",
        "csl_beacon_df = csl_beacon_df.withColumn(\"Timedelta\", timeDiff).filter(\n",
        "        F.col(\"Timedelta\") > time_delta_threshold\n",
        "    )\n",
        "\n",
        "csl_beacon_ranked = csl_beacon_df.groupBy(\n",
        "        \"DeviceVendor\",\n",
        "        \"SourceIP\",\n",
        "        \"DestinationIP\",\n",
        "        \"DestinationPort\",\n",
        "        \"Protocol\",\n",
        "        \"Timedelta\",\n",
        "    ).agg(\n",
        "        F.count(\"Timedelta\").alias(\"Timedeltacount\"),\n",
        "        F.sum(\"SentBytes\").alias(\"TotalSentBytes\"),\n",
        "        F.sum(\"ReceivedBytes\").alias(\"TotalReceivedBytes\"),\n",
        "        F.count(\"*\").alias(\"Totalevents\"),\n",
        "    )\n",
        "\n",
        "w1 = (\n",
        "        Window.partitionBy(\n",
        "            \"DeviceVendor\",\n",
        "            \"SourceIP\",\n",
        "            \"DestinationIP\",\n",
        "            \"DestinationPort\",\n",
        "        )\n",
        "        .orderBy(F.col(\"SourceIP\").cast(\"long\"))\n",
        "        .rowsBetween(-2, 0)\n",
        "    )\n",
        "\n",
        "csl_beacon_df = (\n",
        "        csl_beacon_ranked\n",
        "        .join(csl_dstip, [\"DestinationIP\"])\n",
        "        .withColumn(\"Timedeltalist\", F.collect_list(\"Timedeltacount\").over(w1))\n",
        "        .withColumn(\n",
        "            \"Timedeltalistcount\",\n",
        "            F.expr(\"AGGREGATE(Timedeltalist, DOUBLE(0), (acc, x) -> acc + x)\").cast(\n",
        "                \"long\"\n",
        "            ),\n",
        "        )\n",
        "        .withColumn(\n",
        "            \"Totaleventcount\",\n",
        "            F.sum(\"Timedeltalistcount\").over(\n",
        "                Window.partitionBy(\"SourceIP\", \"DestinationIP\", \"DestinationPort\")\n",
        "            ),\n",
        "        )\n",
        "        .withColumn(\n",
        "            \"Percentbeacon\",\n",
        "            (\n",
        "                F.col(\"Timedeltalistcount\")\n",
        "                / F.sum(\"Timedeltalistcount\").over(\n",
        "                    Window.partitionBy(\n",
        "                        \"DeviceVendor\",\n",
        "                        \"SourceIP\",\n",
        "                        \"DestinationIP\",\n",
        "                        \"DestinationPort\",\n",
        "                    )\n",
        "                )\n",
        "                * 100.0\n",
        "            ),\n",
        "        )\n",
        "        .filter(F.col(\"Percentbeacon\") > percent_beacon_threshold)\n",
        "        .filter(F.col(\"Totaleventcount\") > daily_event_count_threshold)\n",
        "        .orderBy(F.col(\"Percentbeacon\"), ascending=False)\n",
        "    )\n",
        "\n",
        "\n",
        "csl_beacon_df.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "source": [
        "## Export results from ADLS\n",
        "In this step, we will save  the results from previous step as single json file in ADLS. This file can be exported from ADLS to be used with native python session outside spark pool for more data analysis, visualization etc."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "outputs": [],
      "source": [
        "%%synapse\n",
        "dir_name = \"<dir-name>\"  # specify desired directory name\n",
        "new_path = adls_path + dir_name\n",
        "csl_beacon_pd = csl_beacon_df.coalesce(1).write.format(\"json\").save(new_path)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "source": [
        "## Stop Spark Session"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "outputs": [],
      "source": [
        "%synapse stop"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "source": [
        "## Export results from ADLS to local filesystem\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "gather": {
          "logged": 1632987153746
        },
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "outputs": [],
      "source": [
        "def initialize_storage_account(storage_account_name, storage_account_key):\n",
        "\n",
        "    try:\n",
        "        global service_client\n",
        "\n",
        "        service_client = DataLakeServiceClient(\n",
        "            account_url=\"{}://{}.dfs.core.windows.net\".format(\n",
        "                \"https\", storage_account_name\n",
        "            ),\n",
        "            credential=storage_account_key,\n",
        "        )\n",
        "\n",
        "    except Exception as e:\n",
        "        print(e)\n",
        "\n",
        "\n",
        "def list_directory_contents(container_name, input_path, file_type):\n",
        "    try:\n",
        "        file_system_client = service_client.get_file_system_client(\n",
        "            file_system=container_name\n",
        "        )\n",
        "\n",
        "        paths = file_system_client.get_paths(path=input_path)\n",
        "\n",
        "        pathlist = []\n",
        "\n",
        "        for path in paths:\n",
        "            pathlist.append(path.name) if path.name.endswith(file_type) else pathlist\n",
        "\n",
        "        return pathlist\n",
        "\n",
        "    except Exception as e:\n",
        "        print(e)\n",
        "\n",
        "\n",
        "def download_file_from_directory(container_name, input_path, input_file):\n",
        "    try:\n",
        "        file_system_client = service_client.get_file_system_client(\n",
        "            file_system=container_name\n",
        "        )\n",
        "\n",
        "        directory_client = file_system_client.get_directory_client(input_path)\n",
        "\n",
        "        local_file = open(\"output.json\", \"wb\")\n",
        "\n",
        "        file_client = directory_client.get_file_client(input_file)\n",
        "\n",
        "        download = file_client.download_file()\n",
        "\n",
        "        downloaded_bytes = download.readall()\n",
        "\n",
        "        local_file.write(downloaded_bytes)\n",
        "\n",
        "        local_file.close()\n",
        "\n",
        "    except Exception as e:\n",
        "        print(e)\n",
        "\n",
        "\n",
        "def json_normalize(input_file, output_file):\n",
        "    nwbeaconList = []\n",
        "    with open(input_file) as f:\n",
        "        for jsonObj in f:\n",
        "            nwbeaconDict = json.loads(jsonObj)\n",
        "            nwbeaconList.append(nwbeaconDict)\n",
        "\n",
        "    with open(output_file, \"w\") as write_file:\n",
        "        json.dump(nwbeaconList, write_file)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "source": [
        "### Download the files from ADLS\n",
        "\n",
        "In below sections, we will provide input details about ADLS account ad then use functions to connect , list contents and download results locally.\n",
        "\n",
        "If you need help in locating input details, follow below steps\n",
        "- Go to the https://web.azuresynapse.net and sign in to your workspace.\n",
        "- In Synapse Studio, click Data, select the Linked tab, and select the container under Azure Data Lake Storage Gen2.\n",
        "- Navigate to folder from the container, right click and select Properies.\n",
        "- Copy ABFSS path , extact the details and map to the input fields\n",
        "\n",
        "\n",
        "You can check [View account access keys](https://docs.microsoft.com/azure/storage/common/storage-account-keys-manage?tabs=azure-portal#view-account-access-keys) doc to find and retrieve your storage account keys for ADLS account.\n",
        "\n",
        "<p style=\"border: solid; padding: 5pt; color: black; background-color: #AA4000\">\n",
        "<b>Warning</b>: If you are storing secrets such as storage account keys in the notebook you should<br>\n",
        "probably opt to store either into msticpyconfig file on the compute instance or use<br.>\n",
        "Azure Key Vault to store the secrets.<br>\n",
        "Read more about using KeyVault\n",
        "<a href=https://msticpy.readthedocs.io/en/latest/getting_started/msticpyconfig.html#specifying-secrets-as-key-vault-secrets >in the MSTICPY docs</a>\n",
        "</p>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "gather": {
          "logged": 1633376635135
        },
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "outputs": [],
      "source": [
        "# Primary storage info\n",
        "account_name = \"<storage account name>\"  # fill in your primary account name\n",
        "container_name = \"<container name>\"  # fill in your container name\n",
        "subscription_id = \"<subscription id>\"\n",
        "resource_group = \"<resource-group>\" # fill in your resource gropup for ADLS account\n",
        "workspace_name = \"<Microsoft sentinel/Log Analytics workspace name>\" # fill in your la workspace name\n",
        "input_path = f\"WorkspaceResourceId=/subscriptions/{subscription_id}/resourcegroups/{resource_group}/providers/microsoft.operationalinsights/workspaces/\"\n",
        "\n",
        "adls_path = f\"abfss://{container_name}@{account_name}.dfs.core.windows.net/{input_path}/{workspace_name}\"\n",
        "\n",
        "dir_name = \"<dir-name>/\" #Replace the dirname previously specified to store results from spark\n",
        "account_key = \"<storage-account-key>\"  # Replace your storage account key"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "gather": {
          "logged": 1632987683046
        },
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "outputs": [],
      "source": [
        "new_path = input_path + dir_name\n",
        "\n",
        "initialize_storage_account(account_name, account_key)\n",
        "pathlist = list_directory_contents(container_name, new_path, \"json\")\n",
        "input_file = pathlist[0].split(\"/\")[-1]\n",
        "download_file_from_directory(container_name, new_path, input_file)\n",
        "\n",
        "json_normalize(\"output.json\", \"out_normalized.json\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "source": [
        "### Display results"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "gather": {
          "logged": 1632987687222
        },
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "outputs": [],
      "source": [
        "df = pd.read_json('out_normalized.json')\n",
        "\n",
        "df.head()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "source": [
        "## Enrich results \n",
        "In this section, we will enrich entities retrieved from network beaconing behavior such as IP information.\n",
        "Types of Enrichment which will beneficial in perfoming investigation will be IP Geolcation , Whois Registrar information and ThreatIntel lookups.\n",
        "\n",
        "For first time users, please refer `Getting Started Guide For Microsoft Sentinel ML Notebooks` and section [Create your configuration file](https://docs.microsoft.com/azure/sentinel/notebook-get-started#create-your-configuration-file) to create your `mstipyconfig`. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "source": [
        "### IP Geolocation Enrichment\n",
        "In this step, we will use msticpy geolocation capabilities using maxmind database. You will need maxmind API key to download the database.\n",
        "\n",
        "<div style=\"border: solid; padding: 5pt\"><b>Note:</b>\n",
        "You may see the GeoLite driver downloading its database the first time you run this.\n",
        "</div>\n",
        "<br>\n",
        "<details>\n",
        "    <summary>Learn more about MSTICPy GeoIP providers...</summary>\n",
        "    <p>\n",
        "    <a href=https://msticpy.readthedocs.io/en/latest/data_acquisition/GeoIPLookups.html >MSTICPy GeoIP Providers</a>\n",
        "    </p>\n",
        "</details>\n",
        "<br>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "gather": {
          "logged": 1632990019315
        },
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "outputs": [],
      "source": [
        "from msticpy.sectools.geoip import GeoLiteLookup\n",
        "\n",
        "iplocation = GeoLiteLookup()\n",
        "\n",
        "df = iplocation.df_lookup_ip(df, column=\"DestinationIP\")\n",
        "df.head()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "source": [
        "### Whois registration enrichment\n",
        "In this step, we can perform whois lokup on all public destination ips and populate additional information such as ASN. You can use this output to further filter known ASNs from the results."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "gather": {
          "logged": 1632989270369
        },
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "outputs": [],
      "source": [
        "num_ips = len(df[\"DestinationIP\"].unique())\n",
        "print(f\"Performing WhoIs lookups for {num_ips} IPs \", end=\"\")\n",
        "df[\"DestASN\"] = df.apply(lambda x: get_whois_info(x.DestinationIP, True), axis=1)\n",
        "df[\"DestASNFull\"] = df.apply(lambda x: x.DestASN[1], axis=1)\n",
        "df[\"DestASN\"] = df.apply(lambda x: x.DestASN[0], axis=1)\n",
        "\n",
        "#Display results\n",
        "df.head()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "source": [
        "### ThreatIntel Enrichment\n",
        "\n",
        "In this step, we can perform threatintel lookup using msticpy and open source TI providers such as IBM Xforce, VirusTotal, Greynoise etc. \n",
        "Below example shows performing lookup on single IP as well as bulk lookup on all ips using IBM Xforce TI Provider. \n",
        "<br>You will need to register with IBM Xforce and enter API keys into `mstipyconfig.yaml`\n",
        "\n",
        "<details>\n",
        "    <summary>Learn more...</summary>\n",
        "    <p>\n",
        "    </p>\n",
        "    <ul>\n",
        "        <li>More details are shown in the <i>A Tour of Cybersec notebook features</i> notebook</li>\n",
        "        <li><a href=https://msticpy.readthedocs.io/en/latest/data_acquisition/TIProviders.html >Threat Intel Lookups in MSTICPy</a></li>\n",
        "        <li> To learn more about adding TI sources, see the TI Provider setup in the <i>A Getting Started Guide For Microsoft Sentinel ML Notebooks</i> notebook\n",
        "    </ul>\n",
        "</details>\n",
        "<br>\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "gather": {
          "logged": 1632990440057
        },
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "outputs": [],
      "source": [
        "ti_lookup = TILookup()\n",
        "# Perform lookup on single IOC\n",
        "result = ti_lookup.lookup_ioc(observable=\"52.183.120.194\", providers=[\"XForce\"])\n",
        "ti_lookup.result_to_df(result)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "gather": {
          "logged": 1632991107856
        },
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "outputs": [],
      "source": [
        "# Flattening all the desnation IPs into comma separated list\n",
        "ip_list = df['DestinationIP'].astype(str).values.flatten().tolist()\n",
        "\n",
        "# Perform bulk lookup on all IPs with specified providers\n",
        "ti_resp = ti_lookup.lookup_iocs(data=ip_list, providers=[\"AzSTI\", \"XForce\"])\n",
        "select_ti = browse_results(ti_resp, severities=['high','warning'])\n",
        "select_ti"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "source": [
        "## Visualization\n",
        "MSTICpy also includes a feature to allow you to map locations, this can be particularily useful when looking at the distribution of remote network connections or other events. Below we plot the locations of destination IPs observed in our results."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "outputs": [],
      "source": [
        "from msticpy.nbtools import entityschema\n",
        "from msticpy.sectools.ip_utils import convert_to_ip_entities\n",
        "from msticpy.nbtools.foliummap import FoliumMap, get_map_center\n",
        "\n",
        "# Convert our IP addresses in string format into an ip address entity\n",
        "ip_entity = entityschema.IpAddress()\n",
        "ip_list = [convert_to_ip_entities(i)[0] for i in df['DestinationIP'].head(10)]\n",
        "\n",
        "# Get center location of all IP locaitons to center the map on\n",
        "location = get_map_center(ip_list)\n",
        "logon_map = FoliumMap(location=location, zoom_start=4)\n",
        "\n",
        "# Add location markers to our map and dsiplay it\n",
        "if len(ip_list) > 0:\n",
        "    logon_map.add_ip_cluster(ip_entities=ip_list)\n",
        "display(logon_map.folium_map)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "source": [
        "## Conclusion\n",
        "\n",
        "We originally started our hunting on very large datasets of firewall logs. Due to the sheer scale of data, we leveraged spark to load the data. \n",
        "<br>We then performed baselining on historical data and use it to further filter current day dataset. In the next step we performed various data transformation by using spark features such as paritioning, windowing, ranking datatset to find outbound network beaconing like behavior.\n",
        "<br> In order to analyze this data further, we enrich IP entities from result dataset with additional information such as Geolocation, whois registration and threat intel lookups. \n",
        "\n",
        "Analysts can perform further investigation on selected IP addresses from enrichment results by correlating various data sources available. \n",
        "You can then create incidents in Microsoft Sentinel and track investigation in it.\n"
      ]
    }
  ],
  "metadata": {
    "kernel_info": {
      "name": "python310-sdkv2"
    },
    "kernelspec": {
      "display_name": "Python 3.10 - SDK v2",
      "language": "python",
      "name": "python310-sdkv2"
    },
    "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.11.8"
    },
    "microsoft": {
      "host": {
        "AzureML": {
          "notebookHasBeenCompleted": true
        }
      }
    },
    "nteract": {
      "version": "nteract-front-end@1.0.0"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}