Data Collections

"""
This collection of functions can help you programatically manipulate Data
Collections in InsightCloudSec using our REST API.

These functions are intended to be copied and pasted into your code, or can be
imported and used directly in your tools.
"""

from collections import Mapping
import json
import os
import requests
from requests.api import request


# Assumes the presence of these environment variables:
# INSIGHTCLOUDSEC_API_USER_USERNAME -- user we want to authenticate as
# INSIGHTCLOUDSEC_API_USER_PASSWORD -- user's password
# INSIGHTCLOUDSEC_BASE_URL -- base url for the InsightCloudSec instance to run against
# INSIGHTCLOUDSEC_API_KEY -- API Key to interact with the InsightCloudSec REST API
USERNAME, PASSWORD, BASE_URL, API_KEY = (os.environ['INSIGHTCLOUDSEC_API_USER_USERNAME'],
                                os.environ['INSIGHTCLOUDSEC_API_USER_PASSWORD'],
                                os.environ['INSIGHTCLOUDSEC_BASE_URL'],
                                os.environ['INSIGHTCLOUDSEC_API_KEY'])
BASE_HEADERS = {'Content-Type': 'application/json;charset=UTF-8',
                'Accept': 'application/json',
                'Api-Key': API_KEY}


def create_data_collection(collection_name, collection_data=None):
    """
    Create a new data collection with name `collection_name`, optionally
    populated with the values in the dictionary `collection_data`.

    `collection_data` should be one of 2 things:
    - a dictionary mapping Data Collection values to descriptions. Using `None`
      as a description will result in the description being set to the empty
      string.
    - an iterable of strings, which will be inserted as values with no
      descriptions.

    For example, the following is a valid input:

    ```
    {
        'value one': 'description for value one',
        'value two': None,  # description will be set to the empty string
        'value three': 'description for value three'
    }
    ```
    Or you can simply pass a list like `['first value', 'second value']`, which
    is equivalent to passing `{'first value': None, 'second value': None}`
    """

    data = {'collection_name': collection_name,
            'collection_data': normalize_collection(collection_data)}

    return requests.post(
        url=requests.compat.urljoin(BASE_URL, '/v2/datacollections/'),
        headers=BASE_HEADERS,
        data=json.dumps(data)
    )

def update_data_collection(collection_id, collection_data):
    """
    Update the existing data collection with integer ID `collection_id` using
    the data in `collection_data`.

    `collection_data` should be a dictionary mapping values to descriptions.
    Any new key: description value pairs will be inserted into the data 
    collection; any whose key already exists in the Data Collection will be used 
    to update the existing description value.

    Descriptions may be `None`. A description equaling `None` in a new entry
    will result in an empty description. For existing entries, a description
    equaling `None` will result in no changes to an existing description;
    setting a description to the empty string must be done explicitly. For
    example:

    ```
    {
        # set the description for an existing value 'value one', or create a
        # new value with that description
        'value one': 'description for value one',
        # leave the description for an existing value 'value two' unchanged, or
        # create a value 'value two' with no description
        'value two': None,
        # empty the description for an existing value 'value three', or create
        # a new value with an empty description
        'value three': ''
    }
    ```
    Note that this operation does not remove any entries from the data
    collection.
    """
    url = requests.compat.urljoin(
        BASE_URL,
        requests.compat.urljoin('/v2/datacollections/', str(collection_id))
    )

    result = requests.post(
        url=url,
        headers=BASE_HEADERS,
        data=json.dumps({'collection_data': normalize_collection(collection_data)})
    )
    return result


def delete_data_collection_values(collection_id, values_to_delete):
    """
    Delete all entries with values in the iterable `values_to_delete`.

    Note that this is a 2-phase operation: this first checks that the values
    exist and gets their IDs within the collection, then sends the request to
    delete them. This means that calling this method concurrently with other
    data collection manipulation could have unexpected results.
    """
    url = requests.compat.urljoin(
        BASE_URL,
        requests.compat.urljoin('/v2/datacollections/', str(collection_id))
    )

    # phase 1: grab existing entries
    collection_result = requests.get(
        url=url, headers=BASE_HEADERS
    )
    existing_values_to_ids = {
        datum['value']: int(datum['id'])
        for datum in collection_result.json()['collection']['data']
    }

    # pre-deletion check: we should only try to delete entries that exist
    if not set(values_to_delete) < set(existing_values_to_ids):
        raise ValueError(
            'Some values to be deleted not in existing data '
            'collection: {}'.format(set(values_to_delete) - set(existing_values_to_ids))
        )
    # phase 2: delete specified entries
    return requests.delete(
        url=url,
        headers=BASE_HEADERS,
        data=json.dumps({
            'data_ids': [existing_values_to_ids[value] for value in values_to_delete]
        })
    )

def normalize_collection(collection_data):
    if isinstance(collection_data, Mapping):
        return collection_data
    return {datum: None for datum in collection_data}