|
| 1 | +from os import environ |
| 2 | +from pathlib import Path |
| 3 | +from typing import List |
| 4 | + |
| 5 | +import click |
| 6 | +import requests |
| 7 | +from boltons.iterutils import chunked |
| 8 | +from rich.panel import Panel |
| 9 | +from rich.progress import track |
| 10 | + |
| 11 | +from _incydr_cli import console |
| 12 | +from _incydr_cli import logging_options |
| 13 | +from _incydr_cli import render |
| 14 | +from _incydr_cli.cmds.models import AgentCSV, AgentJSON |
| 15 | +from _incydr_cli.cmds.options.output_options import columns_option |
| 16 | +from _incydr_cli.cmds.options.output_options import input_format_option |
| 17 | +from _incydr_cli.cmds.options.output_options import single_format_option |
| 18 | +from _incydr_cli.cmds.options.output_options import SingleFormat |
| 19 | +from _incydr_cli.cmds.options.output_options import table_format_option |
| 20 | +from _incydr_cli.cmds.options.output_options import TableFormat |
| 21 | +from _incydr_cli.core import IncydrCommand |
| 22 | +from _incydr_cli.core import IncydrGroup |
| 23 | +from _incydr_sdk.core.client import Client |
| 24 | +from _incydr_sdk.agents.models import Agent |
| 25 | +from _incydr_sdk.utils import model_as_card |
| 26 | + |
| 27 | + |
| 28 | +@click.group(cls=IncydrGroup) |
| 29 | +@logging_options |
| 30 | +def agents(): |
| 31 | + """View and manage Incydr agents.""" |
| 32 | + |
| 33 | + |
| 34 | +@agents.command("list", cls=IncydrCommand) |
| 35 | +@click.option( |
| 36 | + "--active/--inactive", |
| 37 | + default=None, |
| 38 | + help="Filter by active or inactive agents. Defaults to returning both when when neither option is passed.", |
| 39 | +) |
| 40 | +@table_format_option |
| 41 | +@columns_option |
| 42 | +@logging_options |
| 43 | +def list_( |
| 44 | + active: bool = None, |
| 45 | + format_: TableFormat = None, |
| 46 | + columns: str = None, |
| 47 | +): |
| 48 | + """ |
| 49 | + List agents. |
| 50 | + """ |
| 51 | + client = Client() |
| 52 | + agents = client.agents.v1.iter_all(active=active) |
| 53 | + |
| 54 | + if format_ == TableFormat.table: |
| 55 | + render.table(Agent, agents, columns=columns, flat=False) |
| 56 | + elif format_ == TableFormat.csv: |
| 57 | + render.csv(Agent, agents, columns=columns, flat=True) |
| 58 | + elif format_ == TableFormat.json_pretty: |
| 59 | + for agent in agents: |
| 60 | + console.print_json(agent.json()) |
| 61 | + else: |
| 62 | + for agent in agents: |
| 63 | + click.echo(agent.json()) |
| 64 | + |
| 65 | + |
| 66 | +@agents.command(cls=IncydrCommand) |
| 67 | +@click.argument("agent_id") |
| 68 | +@single_format_option |
| 69 | +@logging_options |
| 70 | +def show( |
| 71 | + agent_id: str, |
| 72 | + format_: SingleFormat, |
| 73 | +): |
| 74 | + """ |
| 75 | + Show details for a single agent. |
| 76 | + """ |
| 77 | + client = Client() |
| 78 | + agent = client.agents.v1.get_agent(agent_id) |
| 79 | + |
| 80 | + if format_ == SingleFormat.rich and client.settings.use_rich: |
| 81 | + console.print(Panel.fit(model_as_card(agent), title=f"Agent {agent.agent_id}")) |
| 82 | + elif format_ == SingleFormat.json_pretty: |
| 83 | + console.print_json(agent.json()) |
| 84 | + else: |
| 85 | + click.echo(agent.json()) |
| 86 | + |
| 87 | + |
| 88 | +@agents.command(cls=IncydrCommand) |
| 89 | +@click.argument("file", type=click.File()) |
| 90 | +@input_format_option |
| 91 | +@logging_options |
| 92 | +def bulk_activate(file: Path, format_: str): |
| 93 | + """ |
| 94 | + Activate a group of agents from a file (CSV or JSON-LINES formatted). |
| 95 | +
|
| 96 | + \b |
| 97 | + Use `-` as filename to read from stdin. |
| 98 | +
|
| 99 | + Input files require a header (for CSV input) or JSON key for each object (for JSON-LINES input) to identify |
| 100 | + which agent ID to activate. |
| 101 | +
|
| 102 | + Header and JSON key values that are accepted are: agent_id, agentId, or guid |
| 103 | + """ |
| 104 | + chunk_size = ( |
| 105 | + environ.get("incydr_batch_size") or environ.get("INCYDR_BATCH_SIZE") or 50 |
| 106 | + ) |
| 107 | + try: |
| 108 | + chunk_size = int(chunk_size) |
| 109 | + except ValueError: |
| 110 | + console.print( |
| 111 | + f"INCYDR_BATCH_SIZE environment variable must be an integer, found: '{chunk_size}'" |
| 112 | + ) |
| 113 | + return |
| 114 | + |
| 115 | + # parse CSV or JSON input |
| 116 | + if format_ == "csv": |
| 117 | + models = AgentCSV.parse_csv(file) |
| 118 | + else: |
| 119 | + models = AgentJSON.parse_json_lines(file) |
| 120 | + try: |
| 121 | + agent_ids = [agent.agent_id for agent in models] |
| 122 | + except ValueError as err: |
| 123 | + console.print(err) |
| 124 | + return |
| 125 | + |
| 126 | + # validate we got at least one agent_id |
| 127 | + num_agents = len(agent_ids) |
| 128 | + if num_agents < 1: |
| 129 | + console.print(f"[red]No agent IDs found in {format_} input.") |
| 130 | + return |
| 131 | + |
| 132 | + client = Client() |
| 133 | + batches = chunked(agent_ids, size=chunk_size) |
| 134 | + for batch in track(batches, description="Activating agents...", console=console): |
| 135 | + process_batch(client, batch, activate=True) |
| 136 | + |
| 137 | + |
| 138 | +@agents.command(cls=IncydrCommand) |
| 139 | +@click.argument("file", type=click.File()) |
| 140 | +@input_format_option |
| 141 | +@logging_options |
| 142 | +def bulk_deactivate(file: Path, format_: str): |
| 143 | + """ |
| 144 | + Deactivate a group of agents from a file (CSV or JSON-LINES formatted). |
| 145 | +
|
| 146 | + \b |
| 147 | + Use `-` as filename to read from stdin. |
| 148 | +
|
| 149 | + Input files require a header (for CSV input) or JSON key for each object (for JSON-LINES input) to identify |
| 150 | + which agent ID to deactivate. |
| 151 | +
|
| 152 | + Header and JSON key values that are accepted are: agent_id, agentId, or guid |
| 153 | + """ |
| 154 | + chunk_size = ( |
| 155 | + environ.get("incydr_batch_size") or environ.get("INCYDR_BATCH_SIZE") or 50 |
| 156 | + ) |
| 157 | + try: |
| 158 | + chunk_size = int(chunk_size) |
| 159 | + except ValueError: |
| 160 | + console.print( |
| 161 | + f"INCYDR_BATCH_SIZE environment variable must be an integer, found: '{chunk_size}'" |
| 162 | + ) |
| 163 | + return |
| 164 | + |
| 165 | + # parse CSV or JSON input |
| 166 | + if format_ == "csv": |
| 167 | + models = AgentCSV.parse_csv(file) |
| 168 | + else: |
| 169 | + models = AgentJSON.parse_json_lines(file) |
| 170 | + try: |
| 171 | + agent_ids = [agent.agent_id for agent in models] |
| 172 | + except ValueError as err: |
| 173 | + console.print(err) |
| 174 | + return |
| 175 | + |
| 176 | + # validate we got at least one agent_id |
| 177 | + num_agents = len(agent_ids) |
| 178 | + if num_agents < 1: |
| 179 | + console.print(f"[red]No agent IDs found in {format_} input.") |
| 180 | + return |
| 181 | + |
| 182 | + client = Client() |
| 183 | + batches = chunked(agent_ids, size=chunk_size) |
| 184 | + for batch in track(batches, description="Deactivating agents...", console=console): |
| 185 | + process_batch(client, batch, activate=False) |
| 186 | + |
| 187 | + |
| 188 | +def process_batch(client: Client, batch: List[str], activate: bool): |
| 189 | + action = "activation" if activate else "deactivation" |
| 190 | + api_call = client.agents.v1.activate if activate else client.agents.v1.deactivate |
| 191 | + process_individually = False |
| 192 | + try: |
| 193 | + api_call(batch) |
| 194 | + except requests.HTTPError as err: |
| 195 | + if err.response.status_code == 404: |
| 196 | + invalid_agent_ids = err.response.json().get("agentsNotFound") |
| 197 | + if invalid_agent_ids is None: |
| 198 | + console.print( |
| 199 | + f"[red]Unknown 404 error processing batch of {len(batch)} agent {action}s." |
| 200 | + ) |
| 201 | + process_individually = True |
| 202 | + else: |
| 203 | + console.print( |
| 204 | + f"[red]404 Error processing batch of {len(batch)} agent {action}s, agent_ids not found:[/red] {invalid_agent_ids}" |
| 205 | + ) |
| 206 | + batch = list(set(batch) - set(invalid_agent_ids)) |
| 207 | + if len(batch) > 0: |
| 208 | + console.print(f"Removing invalid agent_ids and retrying...") |
| 209 | + try: |
| 210 | + api_call(batch) |
| 211 | + except requests.HTTPError as err: |
| 212 | + console.print(f"[red]Error retrying batch. {err.response.text}") |
| 213 | + process_individually = True |
| 214 | + else: |
| 215 | + console.print( |
| 216 | + f"[red]Unknown error processing batch of {len(batch)} agent {action}s." |
| 217 | + ) |
| 218 | + process_individually = True |
| 219 | + if process_individually and len(batch) > 1: |
| 220 | + console.print(f"Trying agent {action} for this batch individually.") |
| 221 | + for agent_id in batch: |
| 222 | + try: |
| 223 | + api_call(agent_id) |
| 224 | + except requests.HTTPError as err: |
| 225 | + msg = f"Failed to process {action} for {agent_id}: {err.response.text}" |
| 226 | + client.settings.logger.error(msg) |
| 227 | + console.print(msg) |
0 commit comments