Deprecation of pick in place of Prompt Toolkit

falcon_toolkit/common/auth_backends/public_mssp.py

@@ -5,17 +5,16 @@
 """
 import os
 
-from typing import Dict, List, Optional
+from typing import Dict, Optional
 
 import keyring
-import pick
 
 from caracara import Client
 
 from falcon_toolkit.common.auth import AuthBackend
 from falcon_toolkit.common.auth_backends.utils import advanced_options_wizard
 from falcon_toolkit.common.constants import KEYRING_SERVICE_NAME
-from falcon_toolkit.common.utils import fancy_input
+from falcon_toolkit.common.utils import fancy_input, choose_cid
 
 
 class PublicCloudFlightControlParentCIDBackend(AuthBackend):
@@ -118,16 +117,10 @@ def authenticate(self) -> Client:
             )[chosen_cid_str]
         else:
             child_cids_data = parent_client.flight_control.get_child_cid_data(cids=child_cids)
-
-            options: List[pick.Option] = []
-            for child_cid_str, child_cid_data in child_cids_data.items():
-                child_cid_name = child_cid_data['name']
-                option_text = f"{child_cid_str}: {child_cid_name}"
-                option = pick.Option(label=option_text, value=child_cid_str)
-                options.append(option)
-
-            chosen_option, _ = pick.pick(options, "Please select a CID to connect to")
-            chosen_cid_str = chosen_option.value
+            chosen_cid_str = choose_cid(
+                cids=child_cids_data,
+                prompt_text="MSSP Child CID Search"
+            )
             chosen_cid = child_cids_data[chosen_cid_str]
 
         chosen_cid_name = chosen_cid['name']

falcon_toolkit/common/auth_backends/utils.py

@@ -5,16 +5,17 @@
 designed to avoid excessive code-reuse between similar implementations of auth backend.
 This file provides:
 - A list of all public CrowdStrike clouds
-- A cloud selection function to allow a user to choose a cloud via pick
+- A cloud selection function to allow a user to choose a cloud via Prompt Toolkit
 - Advanced options configuration for overriding cloud, TLS validation, etc.
 """
 from typing import (
     Dict,
     List,
     NamedTuple,
+    Tuple,
 )
 
-import pick
+from caracara.common.csdialog import csradiolist_dialog
 
 from falcon_toolkit.common.utils import fancy_input
 
@@ -29,13 +30,17 @@
 
 
 def cloud_choice() -> str:
-    """Configure a selection of clouds and allow the user to choose one via pick."""
-    cloud_choices: List[pick.Option] = []
+    """Configure a selection of clouds and allow the user to choose one via Prompt Toolkit."""
+    cloud_choices: List[Tuple] = []
     for cloud_id, cloud_description in CLOUDS.items():
-        cloud_choices.append(pick.Option(cloud_description, cloud_id))
+        cloud_choices.append((cloud_id, cloud_description))
 
-    chosen_option, _ = pick.pick(cloud_choices, title="Please choose a Falcon cloud")
-    chosen_falcon_cloud: str = chosen_option.value
+    chosen_falcon_cloud: str = csradiolist_dialog(
+        title="Falcon Cloud Selection",
+        text="Please choose a Falcon cloud",
+        cancel_text=None,
+        values=cloud_choices,
+    ).run()
 
     return chosen_falcon_cloud
 
@@ -56,12 +61,16 @@ def advanced_options_wizard() -> AdvancedOptionsType:
 
     cloud_name = cloud_choice()
 
-    tls_verify_options: List[pick.Option] = [
-        pick.Option("Verify SSL/TLS certificates (recommended!)", value=True),
-        pick.Option("Do not verify SSL/TLS certificates (not recommended)", False),
+    tls_verify_options = [
+        (True, "Verify SSL/TLS certificates (recommended!)"),
+        (False, "Do not verify SSL/TLS certificates (not recommended)"),
     ]
-    chosen_ssl_verify, _ = pick.pick(tls_verify_options, title="Verify SSL/TLS certificates?")
-    ssl_verify: bool = chosen_ssl_verify.value
+    ssl_verify: bool = csradiolist_dialog(
+        title="Connection Security",
+        text="Enable SSL/TLS certificate verification?",
+        cancel_text=None,
+        values=tls_verify_options,
+    ).run()
 
     proxy_dict = None
     proxy_url_input = fancy_input("HTTPS proxy URL (leave blank if not needed): ", loop=False)

falcon_toolkit/common/utils.py

@@ -5,19 +5,24 @@
 """
 import os
 
+from typing import Dict, Iterable
+
 from colorama import (
     Fore,
     Style,
 )
+from prompt_toolkit import prompt
+from prompt_toolkit.completion import CompleteEvent, Completer, Completion
+from prompt_toolkit.document import Document
 
 from falcon_toolkit.common.constants import LOG_SUB_DIR
 
 
-def fancy_input(prompt: str, loop: bool = True):
+def fancy_input(prompt_text: str, loop: bool = True):
     """Request user input (with colour). Optionally loop until the input is not blank."""
     inputted = False
     colour_prompt = Style.BRIGHT + Fore.BLUE + \
-        prompt + Fore.RESET + Style.RESET_ALL
+        prompt_text + Fore.RESET + Style.RESET_ALL
 
     while not inputted:
         data = input(colour_prompt)
@@ -27,11 +32,11 @@ def fancy_input(prompt: str, loop: bool = True):
     return data
 
 
-def fancy_input_int(prompt: str) -> int:
+def fancy_input_int(prompt_text: str) -> int:
     """Request an integer from the user (with colour), and loop until the input is valid."""
     valid_input = False
     while not valid_input:
-        typed_input = fancy_input(prompt, loop=True)
+        typed_input = fancy_input(prompt_text, loop=True)
         if typed_input.isdigit():
             valid_input = True
 
@@ -64,3 +69,42 @@ def filename_safe_string(unsafe_string: str) -> str:
     clean_string = safe_string.replace(' ', '_')
 
     return clean_string
+
+
+class CIDCompleter(Completer):
+    """Prompt Toolkit Completer that provides a searchable list of CIDs."""
+    def __init__(self, data_dict: Dict[str, Dict]):
+        self.data_dict = data_dict
+
+    def get_completions(
+        self,
+        document: Document,
+        complete_event: CompleteEvent,
+    ) -> Iterable[Completion]:
+        for cid, cid_data in self.data_dict.items():
+            cid_name = cid_data["name"]
+            cloud_name = cid_data.get("cloud_name")
+            if cloud_name:
+                display_meta = f"{cid_name} [{cloud_name}]"
+            else:
+                display_meta = cid_name
+
+            word_lower = document.current_line.lower()
+            if word_lower in cid or word_lower in display_meta.lower():
+                yield Completion(
+                    cid,
+                    start_position=-len(document.current_line),
+                    display=cid,
+                    display_meta=display_meta,
+                )
+
+
+def choose_cid(cids: Dict[str, Dict], prompt_text="CID Search") -> str:
+    """Choose a CID from a dictionary of CIDs via Prompt Toolkit and return the CID string."""
+    cid_completer = CIDCompleter(data_dict=cids)
+    chosen_cid = None
+    while chosen_cid not in cids:
+        chosen_cid = prompt(f"{prompt_text} >", completer=cid_completer)
+
+    print(chosen_cid)
+    return chosen_cid

falcon_toolkit/falcon.py

@@ -31,8 +31,8 @@
 import sys
 
 import click
-import pick
 
+from caracara.common.csdialog import csradiolist_dialog
 from caracara_filters.dialects import DIALECTS
 from colorama import (
     deinit as colorama_deinit,
@@ -136,47 +136,45 @@ def cli(
     ):
         # The user has used Falcon Toolkit before, and uses the default directory, so we
         # offer to move the configuration folder for them.
-        choice, _ = pick.pick(
-            options=[
-                pick.Option(
-                    f"Please move my current folder contents to {config_path}",
-                    "MOVE_FOLDER",
-                ),
-                pick.Option(
-                    (
-                        f"Leave my old folder ({OLD_DEFAULT_CONFIG_DIR}) alone "
-                        f"and create a new one at {config_path}"
-                    ),
-                    "LEAVE_ALONE",
-                ),
-                pick.Option(
-                    (
-                        f"Leave my old folder ({OLD_DEFAULT_CONFIG_DIR}) alone, "
-                        f"create a new one at {config_path}, and copy my configuration "
-                        "file there."
-                    ),
-                    "COPY_CONFIG_ONLY",
-                ),
-                pick.Option(
-                    "Exit Falcon Toolkit and do nothing",
-                    "ABORT",
-                ),
-            ],
-            title=(
-                "As of Falcon Toolkit 3.3.0, the configuration directory has moved to a "
-                "platform-specific data configuration directory."
+        option_pairs = [
+            (
+                "MOVE_FOLDER",
+                f"Please move my current folder contents to {config_path}",
+            ),
+            (
+                "LEAVE_ALONE",
+                f"Leave my old folder ({OLD_DEFAULT_CONFIG_DIR}) alone "
+                f"and create a new one at {config_path}",
+            ),
+            (
+                "COPY_CONFIG_ONLY",
+                f"create a new one at {config_path}, and copy my configuration file there",
             )
-        )
+        ]
+        choice = csradiolist_dialog(
+            title="Falcon Toolkit Configuration Directory",
+            text=(
+                "As of Falcon Toolkit 3.3.0, the configuration directory has moved to a "
+                "platform-specific data configuration directory. Please choose how you "
+                "would like the Toolkit to proceed, or press Abort to exit the program "
+                "without making any changes."
+            ),
+            cancel_text="Cancel",
+            values=option_pairs,
+        ).run()
+        if choice is None:
+            click.echo(click.style("Exiting the Toolkit without making changes.", bold=True))
+            sys.exit(1)
 
-        if choice.value == "MOVE_FOLDER":
+        if choice == "MOVE_FOLDER":
             click.echo(f"Moving {OLD_DEFAULT_CONFIG_DIR} to {config_path}")
             os.rename(OLD_DEFAULT_CONFIG_DIR, config_path)
-        elif choice.value == "LEAVE_ALONE":
+        elif choice == "LEAVE_ALONE":
             click.echo(
                 f"Creating a new, empty data directory at {config_path} "
                 "and leaving the original folder alone"
             )
-        elif choice.value == "COPY_CONFIG_ONLY":
+        elif choice == "COPY_CONFIG_ONLY":
             click.echo(
                 f"Creating a new, empty data directory at {config_path} "
                 "and copying the current configuration there"

falcon_toolkit/policies/cli.py

@@ -4,13 +4,14 @@
 of the logic itself is contained in other files, including policies.py
 """
 import os
+import sys
 
 from typing import List
 
 import click
-import pick
 
 from caracara import Client
+from caracara.common.csdialog import csradiolist_dialog
 from caracara.common.policy_wrapper import Policy
 
 from click_option_group import (
@@ -93,16 +94,23 @@ def policies_export(ctx: click.Context):
     policies_api: PoliciesApiModule = ctx.obj['policies_api']
     policies_type: str = ctx.obj['policies_type']
     click.echo("Loading policies...")
-    policies = policies_api.describe_policies()
+    policies: List[Policy] = policies_api.describe_policies()
 
-    options: List[pick.Option] = []
+    options = []
     for policy in policies:
         option_text = f"{policy.name} [{policy.platform_name}]"
-        option = pick.Option(label=option_text, value=policy)
-        options.append(option)
+        options.append((policy, option_text))
+
+    chosen_policy = csradiolist_dialog(
+        title="Policy Selection",
+        text="Please choose a policy to export",
+        values=options,
+    ).run()
+
+    if chosen_policy is None:
+        click.echo("No option chosen; aborting.")
+        sys.exit(1)
 
-    chosen_option, _ = pick.pick(options, "Please choose a policy to export")
-    chosen_policy: Policy = chosen_option.value
     default_filename = f"{chosen_policy.name}.json"
     reasonable_filename = False
     while not reasonable_filename: