Adds maintenance token retrieval (#156)

caracara/client.py

@@ -31,6 +31,7 @@
     PreventionPoliciesApiModule,
     ResponsePoliciesApiModule,
     RTRApiModule,
+    SensorUpdatePoliciesApiModule,
     UsersApiModule,
 )
 
@@ -158,6 +159,8 @@ def __init__(  # pylint: disable=R0913,R0914,R0915
         self.response_policies = ResponsePoliciesApiModule(self.api_authentication)
         self.logger.debug("Setting up the RTR module")
         self.rtr = RTRApiModule(self.api_authentication)
+        self.logger.debug("Setting up the Sensor Update Policies module")
+        self.sensor_update_policies = SensorUpdatePoliciesApiModule(self.api_authentication)
         self.logger.debug("Setting up the Users module")
         self.users = UsersApiModule(self.api_authentication)
 

caracara/modules/__init__.py

@@ -11,6 +11,7 @@
     'PreventionPoliciesApiModule',
     'ResponsePoliciesApiModule',
     'RTRApiModule',
+    'SensorUpdatePoliciesApiModule',
     'UsersApiModule',
 ]
 
@@ -20,4 +21,5 @@
 from caracara.modules.prevention_policies import PreventionPoliciesApiModule
 from caracara.modules.response_policies import ResponsePoliciesApiModule
 from caracara.modules.rtr import RTRApiModule
+from caracara.modules.sensor_update_policies import SensorUpdatePoliciesApiModule
 from caracara.modules.users import UsersApiModule

caracara/modules/sensor_update_policies/__init__.py

@@ -0,0 +1,7 @@
+"""Caracara Sensor Policies API Module."""
+__all__ = [
+    'SensorUpdatePoliciesApiModule',
+]
+
+from caracara.modules.sensor_update_policies.sensor_update_policies \
+    import SensorUpdatePoliciesApiModule

caracara/modules/sensor_update_policies/sensor_update_policies.py

@@ -0,0 +1,87 @@
+"""Sensor Update Policy APIs.
+
+This module allows for interaction with the Sensor Update policies, and facilitates
+the retrieval of maintenance tokens.
+"""
+import datetime
+
+from typing import Optional
+
+from falconpy import (
+    OAuth2,
+    SensorUpdatePolicies,
+)
+
+from caracara.common.exceptions import BaseCaracaraError
+from caracara.common.module import FalconApiModule
+
+
+class SensorUpdatePoliciesApiModule(FalconApiModule):
+    """This module facilitates interactions with the Sensor Update Policies API."""
+
+    name = "CrowdStrike Sensor Update Policies API Module"
+    help = "Interact with the Sensor Update Policies API, and get maintenance tokens"
+
+    def __init__(self, api_authentication: OAuth2):
+        """Construct an instance of the SensorUpdatePoliciesApiModule."""
+        super().__init__(api_authentication)
+
+        self.logger.debug("Configuring the FalconPy Sensor Update Policies module")
+        self.sensor_update_policies_api = SensorUpdatePolicies(auth_object=self.api_authentication)
+
+    def get_maintenance_token(
+        self,
+        device_id: str,
+        audit_message: Optional[str] = None,
+    ) -> str:
+        """Get the maintenance token for a device.
+
+        Arguments
+        ---------
+        device_id: str, required
+            A Falcon Device ID (AID) to get the maintenance token for.
+
+        audit_message: str or None
+            A text string containing an audit message for the retrieval of the maintenance token.
+            If left blank / None, Caracara will generate one for you.
+
+        Returns
+        -------
+        str: The maintenance token.
+        """
+        if audit_message is None:
+            timestamp = datetime.datetime.now(tz=datetime.timezone.utc,).strftime(
+                "%Y-%m-%d %H:%M:%S"
+            )
+
+            audit_message = f"Generated via Caracara at {timestamp} UTC"
+
+        response = self.sensor_update_policies_api.reveal_uninstall_token(
+            audit_message=audit_message,
+            device_id=device_id,
+        )
+
+        if response['status_code'] == 200:
+            body = response['body']
+            return body['resources'][0]['uninstall_token']
+
+        raise BaseCaracaraError(
+            "The API operation failed to generate a maintenance token."
+            "Check that you have the Sensor Update Policies - Write permission "
+            "enabled on your API token."
+        )
+
+    def get_bulk_maintenance_token(self, audit_message: Optional[str] = None) -> str:
+        """Get the bulk maintenance token for a Falcon tenant.
+
+        Arguments
+        ---------
+        audit_message: str or None
+            A text string containing an audit message for the retrieval of the maintenance token.
+            If left blank / None, Caracara will generate one for you.
+
+        Returns
+        -------
+        str: The bulk maintenance token.
+        """
+        return self.get_maintenance_token(device_id="MAINTENANCE", audit_message=audit_message)

examples/common/__init__.py

@@ -15,11 +15,13 @@
     NoSessionsConnected,
 )
 from examples.common.filter_loader import parse_filter_list
+from examples.common.prompts import choose_item
 from examples.common.timer import Timer
 from examples.common.utils import pretty_print
 
 __all__ = [
     "caracara_example",
+    "choose_item",
     "parse_filter_list",
     "pretty_print",
     "MissingArgument",

examples/common/prompts.py

@@ -0,0 +1,53 @@
+"""
+Caracara Example Collection.
+
+This file contains a pretty prompt_toolkit-based completer for devices.
+"""
+from typing import Dict, Iterable
+
+from prompt_toolkit import prompt
+from prompt_toolkit.completion import CompleteEvent, Completer, Completion
+from prompt_toolkit.document import Document
+
+
+class PrettyCompleter(Completer):
+    """Prompt Toolkit Completer that provides a searchable list of items."""
+
+    def __init__(self, data_dict: Dict[str, Dict]):
+        """Create a new completer based on a dictionary that maps items to meta strings.
+
+        Each item should be in the format:
+        {
+            "id1": "label",
+            "id2": "label2",
+            etc.
+        }
+        """
+        self.data_dict = data_dict
+
+    def get_completions(
+        self,
+        document: Document,
+        complete_event: CompleteEvent,
+    ) -> Iterable[Completion]:
+        """Yield items that match the entered search string."""
+        for item_id, item_label in self.data_dict.items():
+            word_lower = document.current_line.lower()
+            if word_lower in item_id.lower() or word_lower in item_label.lower():
+                yield Completion(
+                    item_id,
+                    start_position=-len(document.current_line),
+                    display=item_id,
+                    display_meta=item_label,
+                )
+
+
+def choose_item(items: Dict[str, Dict], prompt_text="Item Search") -> str:
+    """Visually choose an item ID from a dict of IDs mapped to strings."""
+    completer = PrettyCompleter(data_dict=items)
+    chosen_id = None
+    while chosen_id not in items:
+        chosen_id = prompt(f"{prompt_text} >", completer=completer)
+
+    print(chosen_id)
+    return chosen_id

examples/sensor_update_policies/get_maintenance_token.py

@@ -0,0 +1,36 @@
+#!/usr/bin/env python3
+"""
+Caracara Examples Collection.
+
+get_maintenance_token.py
+
+This example will allow you to search for a system in the CID and get its maintenance token.
+"""
+from typing import Dict
+
+from caracara import Client
+
+from examples.common import (
+    caracara_example,
+    choose_item,
+)
+
+
+@caracara_example
+def get_maintenance_token(**kwargs):
+    """Search for a system and get its maintenance token."""
+    client: Client = kwargs['client']
+
+    print("Getting all devices in the Falcon tenant")
+    devices: Dict[str, Dict] = client.hosts.describe_devices()
+
+    id_name_mapping = {
+        "MAINTENANCE": "Bulk Maintenance Token",
+    }
+    for device_id, device_data in devices.items():
+        id_name_mapping[device_id] = device_data['hostname']
+
+    chosen_id = choose_item(id_name_mapping, prompt_text="Search for a Device")
+
+    maintenance_token = client.sensor_update_policies.get_maintenance_token(chosen_id)
+    print(maintenance_token)