{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SEC405: Scalable, Automated Anomaly Detection with Amazon GuardDuty and SageMaker\n",
    "\n",
    "## Using IP Insights to score security findings\n",
    "-------\n",
    "[Return to the workshop repository](https://github.com/aws-samples/aws-security-workshops/edit/master/detection-ml-wksp/)\n",
    "\n",
    "Amazon SageMaker IP Insights is an unsupervised anomaly detection algorithm for susicipous IP addresses that uses statistical modeling and neural networks to capture associations between online resources (such as account IDs or hostnames) and IPv4 addresses. Under the hood, it learns vector representations for online resources and IP addresses. This essentially means that if the vector representing an IP address and an online resource are close together, then it is likey for that IP address to access that online resource, even if it has never accessed it before.\n",
    "\n",
    "In this notebook, we use the Amazon SageMaker IP Insights algorithm to train a model using the `<principal ID, IP address`> tuples we generated from the CloudTrail log data, and then use the model to perform inference on the same type of tuples generated from GuardDuty findings to determine how unusual it is to see a particular IP address for a given principal involved with a finding.\n",
    "\n",
    "After running this notebook, you should be able to:\n",
    "\n",
    "- obtain, transform, and store data for use in Amazon SageMaker,\n",
    "- create an AWS SageMaker training job to produce an IP Insights model,\n",
    "- use the model to perform inference with an Amazon SageMaker endpoint.\n",
    "\n",
    "If you would like to know more, please check out the [SageMaker IP Inisghts Documentation](https://docs.aws.amazon.com/sagemaker/latest/dg/ip-insights.html).\n",
    "\n",
    "## Setup\n",
    "------\n",
    "*This notebook was created and tested on a ml.m4.xlarge notebook instance. We recommend using the same, but other instance types should still work.*\n",
    "\n",
    "The following is a cell that contains Python code and can be executed by clicking the button above labelled \"Run\". When a cell is running, you will see a star in the parentheses to the left (e.g., `In [*]`), and when it has completed you will see a number in the parentheses. Each click of \"Run\" will execute the next cell in the notebook.\n",
    "\n",
    "Go ahead and click **Run** now. You should see the text in the `print` statement get printed just beneath the cell.\n",
    "\n",
    "All of these cells share the same interpreter, so if a cell imports modules, like this one does, those modules will be available to every subsequent cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "import botocore\n",
    "import os\n",
    "import sagemaker\n",
    "\n",
    "print(\"Welcome to IP Insights!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Configure Amazon S3 Bucket\n",
    "\n",
    "Before going further, we to specify the S3 bucket that SageMaker will use for input and output data for the model, which will be the bucket where our training and inference tuples from CloudTrail logs and GuardDuty findings, respectively, are located. Edit the following cell to specify the name of the bucket and then run it; you do not need to change the prefix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Specify the full name of your \"sec405-tuplesbucket\" here\n",
    "bucket = 'sec405-tuplesbucket-########'\n",
    "prefix = ''"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, run the next cell to complete the setup."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "execution_role = sagemaker.get_execution_role()\n",
    "\n",
    "# Check if the bucket exists\n",
    "try:\n",
    "    boto3.Session().client('s3').head_bucket(Bucket=bucket)\n",
    "except botocore.exceptions.ParamValidationError as e:\n",
    "    print('Hey! You either forgot to specify your S3 bucket'\n",
    "          ' or you gave your bucket an invalid name!')\n",
    "except botocore.exceptions.ClientError as e:\n",
    "    if e.response['Error']['Code'] == '403':\n",
    "        print(\"Hey! You don't have permission to access the bucket, {}.\".format(bucket))\n",
    "    elif e.response['Error']['Code'] == '404':\n",
    "        print(\"Hey! Your bucket, {}, doesn't exist!\".format(bucket))\n",
    "    else:\n",
    "        raise\n",
    "else:\n",
    "    print('Training input/output will be stored in: s3://{}/{}'.format(bucket, prefix))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training\n",
    "\n",
    "Execute the two cells below to start training. Training should take several minutes to complete. You can look at various training metrics in the log as the model trains. These logs are also available in CloudWatch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.amazon.amazon_estimator import get_image_uri\n",
    "\n",
    "image = get_image_uri(boto3.Session().region_name, 'ipinsights')\n",
    "\n",
    "\n",
    "# Configure SageMaker IP Insights input channels\n",
    "train_key = os.path.join(prefix, 'train', 'cloudtrail_tuples.csv')\n",
    "s3_train_data = 's3://{}/{}'.format(bucket, train_key)\n",
    "\n",
    "input_data = {\n",
    "    'train': sagemaker.session.s3_input(s3_train_data, distribution='FullyReplicated', content_type='text/csv')\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# Set up the estimator with training job configuration\n",
    "ip_insights = sagemaker.estimator.Estimator(\n",
    "    image, \n",
    "    execution_role, \n",
    "    train_instance_count=1, \n",
    "    train_instance_type='ml.m4.xlarge',\n",
    "    output_path='s3://{}/{}/output'.format(bucket, prefix),\n",
    "    sagemaker_session=sagemaker.Session())\n",
    "\n",
    "# Configure algorithm-specific hyperparameters\n",
    "ip_insights.set_hyperparameters(\n",
    "    num_entity_vectors='20000',\n",
    "    random_negative_sampling_rate='5',\n",
    "    vector_dim='128', \n",
    "    mini_batch_size='1000',\n",
    "    epochs='5',\n",
    "    learning_rate='0.01',\n",
    ")\n",
    "\n",
    "# Start the training job (should take 3-4 minutes to complete)  \n",
    "ip_insights.fit(input_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Training job name: {}'.format(ip_insights.latest_training_job.job_name))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Deployment\n",
    "Execute the cell below to deploy the trained model on an endpoint for inference. It should take 5-7 minutes to spin up the instance and deploy the model (the horizontal dashed line represents progress, and it will print an exclamation point \\[!\\] when it is complete)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "predictor = ip_insights.deploy(\n",
    "    initial_instance_count=1,\n",
    "    instance_type='ml.m4.xlarge'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Endpoint name: {}'.format(predictor.endpoint))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Inference\n",
    "We can pass data in a variety of formats to our inference endpoint. In this example, we will pass CSV-formmated data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.predictor import csv_serializer, json_deserializer\n",
    "\n",
    "predictor.content_type = 'text/csv'\n",
    "predictor.serializer = csv_serializer\n",
    "predictor.accept = 'application/json'\n",
    "predictor.deserializer = json_deserializer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When queried by a principal and an IPAddress, the model returns a score (called 'dot_product') which indicates how expected that event is. In other words, the higher the dot_product, the more normal the event is. First let's run inference on the training (normal) data for sanity check."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "# Run inference on training (normal) data for sanity check\n",
    "s3_infer_data = 's3://{}/{}'.format(bucket, train_key)\n",
    "inference_data = pd.read_csv(s3_infer_data, header=None)\n",
    "inference_data.head()\n",
    "predictor.predict(inference_data.values)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's run inference on the GuardDuty findings. Notice that the scores are much lower than the normal scores."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run inference on GuardDuty findings\n",
    "infer_key = os.path.join(prefix, 'infer', 'guardduty_tuples.csv')\n",
    "s3_infer_data = 's3://{}/{}'.format(bucket, infer_key)\n",
    "inference_data = pd.read_csv(s3_infer_data, header=None)\n",
    "inference_data.head()\n",
    "predictor.predict(inference_data.values)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Clean-up\n",
    "\n",
    "To clean up resources created during the workshop, please see the [Cleaning up](https://github.com/aws-samples/aws-security-workshops/blob/cff322dab7cc0b9d71c4f1575c7016389b9dbe64/detection-ml-wksp/README.md) section in the workshop README guide."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_mxnet_p36",
   "language": "python",
   "name": "conda_mxnet_p36"
  },
  "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}