import json import traceback from concurrent.futures import ThreadPoolExecutor, as_completed from dataclasses import dataclass, field from datetime import datetime, timedelta, UTC from enum import Enum from typing import Any import demistomock as demisto # noqa: F401 import urllib3 from CommonServerPython import * # noqa: F401 from CommonServerUserPython import * # noqa from ContentClientApiModule import * # Disable insecure warnings urllib3.disable_warnings() """ KOI Integration for fetching Alerts and Audit Logs from the KOI API. """ # region Constants and helpers # ================================= # Constants and helpers # ================================= INTEGRATION_NAME = "KOI" class ApiPaths: """Centralized KOI API endpoint paths. All paths are relative to the KOI base URL configured in integration parameters. Use the classmethods for parameterized routes (e.g., a specific policy or item) so URL construction lives in exactly one place. """ BASE = "/api/external/v2" ALERTS = f"{BASE}/alerts" AUDIT_LOGS = f"{BASE}/audit-logs" POLICIES = f"{BASE}/policies" ALLOWLIST = f"{BASE}/policies/allowlist" BLOCKLIST = f"{BASE}/policies/blocklist" INVENTORY = f"{BASE}/inventory" INVENTORY_SEARCH = f"{BASE}/inventory/search" @classmethod def policy(cls, policy_id: int) -> str: """Return the path for a specific policy by ID.""" return f"{cls.POLICIES}/{policy_id}" @classmethod def inventory_item(cls, item_id: str) -> str: """Return the path for a specific inventory item by ID.""" return f"{cls.INVENTORY}/{item_id}" @classmethod def inventory_item_endpoints(cls, item_id: str) -> str: """Return the path for the endpoints of a specific inventory item.""" return f"{cls.INVENTORY}/{item_id}/endpoints" class Config: """Global static configuration.""" VENDOR = "koi" PRODUCT = "koi" # Date format for API requests (ISO 8601) DATE_FORMAT = "%Y-%m-%dT%H:%M:%SZ" # Pagination DEFAULT_PAGE_SIZE = 50 MAX_PAGE_SIZE = 500 MAX_PAGES_PER_FETCH = 10 DEFAULT_PAGE = 1 DEFAULT_LIMIT = 50 MAX_LIMIT = 1000 # Fetch defaults DEFAULT_MAX_FETCH = 5000 # Default lookback time for first fetch or get-events command DEFAULT_FROM_TIME = "5 minutes ago" # API sort direction for chronological ordering SORT_DIRECTION = "asc" # Test module settings TEST_MODULE_LOOKBACK_MINUTES = 5 TEST_MODULE_MAX_EVENTS = 1 class LogType(Enum): """Enum to hold all configuration for different log types.""" ALERTS = ("alerts", "Alerts", ApiPaths.ALERTS) AUDIT = ("audit", "Audit", ApiPaths.AUDIT_LOGS) def __init__(self, type_string: str, title: str, api_endpoint: str): self.type_string = type_string self.title = title self.api_endpoint = api_endpoint # Valid audit log type filters VALID_AUDIT_TYPES = [ "approval_requests", "devices", "endpoints", "extensions", "firewall", "guardrails", "notifications", "policies", "remediation", "requests", "settings", "vetting", ] # Valid marketplace values for allowlist operations VALID_MARKETPLACES = [ "chocolatey", "chrome_web_store", "claude_desktop_extensions", "cursor", "docker", "edge_add_ons", "firefox_add_ons", "github_mcp_registry", "homebrew", "hugging_face", "jetbrains", "linux", "mac", "notepad++", "npm", "office_add_ins", "open_vsx_registry", "pypi", "visual_studio", "vscode", "windows", "windsurf", ] def get_formatted_utc_time(date_input: str | None) -> str: """Parse input and return the formatted UTC time string for KOI API. Args: date_input: Date string to parse (e.g., '3 days ago', '2024-01-01T00:00:00Z') Returns: Formatted UTC time string in ISO 8601 format. """ parsed_dt = parse_date_or_use_current(date_input) formatted_time = parsed_dt.strftime(Config.DATE_FORMAT) demisto.debug(f"[Date Helper] Input: '{date_input}' -> Output: '{formatted_time}' (UTC)") return formatted_time def parse_date_or_use_current(date_string: str | None) -> datetime: """Parse a date string or return current UTC datetime if parsing fails. Uses arg_to_datetime from CommonServerPython for consistent date parsing. Args: date_string: Date string to parse, or None to use current UTC time. Returns: Parsed datetime object in UTC. """ if not date_string: current_time = datetime.now(UTC) demisto.debug(f"[Date Helper] No input provided. Using current UTC: {current_time}") return current_time demisto.debug(f"[Date Helper] Attempting to parse date string: '{date_string}'") parsed_datetime = arg_to_datetime(arg=date_string, is_utc=True) if not parsed_datetime: demisto.debug(f"[Date Helper] Failed to parse '{date_string}'. Fallback to current UTC.") return datetime.now(UTC) demisto.debug(f"[Date Helper] Final parsed date: {parsed_datetime.isoformat()}") return parsed_datetime def get_log_types_from_titles(event_types_to_fetch: list[str]) -> list[LogType]: """Convert user-facing event type titles into LogType Enum members. Args: event_types_to_fetch: List of event type titles (e.g., ["Alerts", "Audit"]). Raises: DemistoException: If any of the provided event type titles are invalid. Returns: List of LogType Enum members. """ valid_titles = {lt.title for lt in LogType} invalid_types = [title for title in event_types_to_fetch if title not in valid_titles] if invalid_types: valid_options = ", ".join(sorted(valid_titles)) raise DemistoException( f"Invalid event type(s) provided: {invalid_types}. " f"Please select from the following list: {valid_options}" ) return [lt for lt in LogType if lt.title in event_types_to_fetch] def extract_time_from_event(event: dict, log_type: LogType) -> str | None: """Extract the time field value from an event based on log type. For alerts: finding_info.created_time (epoch ms) -> converted to ISO 8601. For audit logs: created_at (ISO 8601 string). Args: event: The event dictionary. log_type: The LogType Enum member. Returns: ISO 8601 formatted time string, or None if not found. """ if log_type == LogType.ALERTS: finding_info = event.get("finding_info", {}) created_time_ms = finding_info.get("created_time") if created_time_ms: try: dt = datetime.fromtimestamp(created_time_ms / 1000, tz=UTC) return dt.strftime(Config.DATE_FORMAT) except (ValueError, TypeError, OSError): demisto.debug(f"[Time Extract] Failed to parse alert created_time: {created_time_ms}") return None else: return event.get("created_at") return None def add_time_to_events(events: list[dict], log_type: LogType) -> None: """Add _time and source_log_type fields to events for XSIAM ingestion. Uses extract_time_from_event for consistent time extraction across all code paths. Args: events: List of event dictionaries to enrich. log_type: The LogType Enum member representing the source. """ for event in events: event_time = extract_time_from_event(event, log_type) if event_time: event["_time"] = event_time else: demisto.debug(f"[Event Time] WARNING: Event missing time field: {event.get('id', 'unknown')}") event["source_log_type"] = log_type.title def get_event_id(event: dict) -> str | None: """Extract the event ID from an event dictionary. Args: event: The event dictionary. Returns: The event ID string, or None if not found. """ for id_field in ("id", "alert_id", "log_id", "uuid"): event_id = event.get(id_field) if event_id: return str(event_id) return None def deduplicate_events(events: list[dict], last_fetched_ids: list[str]) -> list[dict]: """Remove already-processed events based on previously fetched IDs. Args: events: List of events to deduplicate. last_fetched_ids: List of event IDs from the previous run. Returns: List of new (non-duplicate) events. """ if not events: demisto.debug("[Dedup] No events to process") return events if not last_fetched_ids: demisto.debug("[Dedup] No deduplication needed (first run - no previous IDs)") return events demisto.debug(f"[Dedup] Checking {len(events)} events against {len(last_fetched_ids)} previously fetched IDs") fetched_ids_set = set(last_fetched_ids) new_events = [event for event in events if get_event_id(event) not in fetched_ids_set] skipped_count = len(events) - len(new_events) if skipped_count > 0: demisto.debug(f"[Dedup] Skipped {skipped_count} duplicates. {len(new_events)} new events remain.") else: demisto.debug("[Dedup] No duplicates found.") return new_events def parse_list_items_from_entry_id(entry_id: str) -> list[dict[str, Any]]: """Read and parse a JSON file from a War Room entry ID containing list items. The JSON file must contain a list of item objects, each with at least 'item_id' and 'marketplace'. Args: entry_id: The War Room entry ID of the uploaded JSON file. Returns: List of item dictionaries parsed from the JSON file. Raises: DemistoException: If the file cannot be read, parsed, or has invalid structure. """ try: filepath_result = demisto.getFilePath(entry_id) except Exception as e: raise DemistoException(f"Could not find file for entry ID '{entry_id}': {e}") if not filepath_result or not (file_path := filepath_result.get("path")): raise DemistoException(f"Entry ID '{entry_id}' is not a valid file entry.") demisto.debug(f"[File Parse] Reading items from file: {file_path}") try: with open(file_path, encoding="utf-8") as f: data = json.load(f) except json.JSONDecodeError as e: raise DemistoException(f"Failed to parse JSON file from entry ID '{entry_id}': {e}") except OSError as e: raise DemistoException(f"Failed to read file from entry ID '{entry_id}': {e}") if not isinstance(data, list): raise DemistoException( f"Invalid JSON structure in entry ID '{entry_id}': expected a list of items, got {type(data).__name__}." ) for i, item in enumerate(data): if not isinstance(item, dict): raise DemistoException(f"Invalid item at index {i}: expected a dictionary, got {type(item).__name__}.") if "item_id" not in item or "marketplace" not in item: raise DemistoException(f"Invalid item at index {i}: each item must contain 'item_id' and 'marketplace'.") if item["marketplace"] not in VALID_MARKETPLACES: raise DemistoException( f"Invalid marketplace '{item['marketplace']}' at index {i}. Valid values: {VALID_MARKETPLACES}" ) demisto.debug(f"[File Parse] Parsed {len(data)} items from entry ID '{entry_id}'") return data def resolve_items_from_args(args: dict[str, Any]) -> list[dict[str, Any]]: """Resolve list items from command arguments. Supports two input modes: - Bulk from file: 'items_list_raw_json_entry_id' with a War Room entry ID. - Single item: 'item_id' and 'marketplace' (with optional 'created_by' and 'notes'). File entry ID takes priority when both modes are provided. Args: args: Command arguments dictionary. Returns: List of item dictionaries. Raises: DemistoException: If neither mode provides valid input, or marketplace is invalid. """ entry_id: str | None = args.get("items_list_raw_json_entry_id") item_id: str | None = args.get("item_id") marketplace: str | None = args.get("marketplace") if entry_id: return parse_list_items_from_entry_id(entry_id) if item_id and marketplace: if marketplace not in VALID_MARKETPLACES: raise DemistoException(f"Invalid marketplace '{marketplace}'. Valid values: {VALID_MARKETPLACES}") item: dict[str, Any] = { "item_id": item_id, "marketplace": marketplace, } created_by: str | None = args.get("created_by") notes: str | None = args.get("notes") if created_by: item["created_by"] = created_by if notes: item["notes"] = notes return [item] raise DemistoException( "Either 'item_id' and 'marketplace' must be provided, or 'items_list_raw_json_entry_id' must be provided." ) def parse_filter_from_args(args: dict[str, Any]) -> dict[str, Any]: """Resolve a filter object from command arguments. Supports two input modes: - Inline JSON: 'filter_json' with a JSON string. - File upload: 'filter_raw_json_entry_id' with a War Room entry ID of a JSON file. File entry ID takes priority when both are provided. Args: args: Command arguments dictionary. Returns: Parsed filter dictionary. Raises: DemistoException: If no filter is provided, the JSON cannot be parsed, or the file cannot be read. """ entry_id: str | None = args.get("filter_raw_json_entry_id") filter_json: str | None = args.get("filter_json") if entry_id: try: filepath_result = demisto.getFilePath(entry_id) except Exception as e: raise DemistoException(f"Could not find file for entry ID '{entry_id}': {e}") if not filepath_result or "path" not in filepath_result: raise DemistoException(f"Entry ID '{entry_id}' is not a valid file entry.") file_path = filepath_result["path"] demisto.debug(f"[Filter Parse] Reading filter from file: {file_path}") try: with open(file_path, encoding="utf-8") as f: data = json.load(f) except json.JSONDecodeError as e: raise DemistoException(f"Failed to parse JSON filter file from entry ID '{entry_id}': {e}") except OSError as e: raise DemistoException(f"Failed to read filter file from entry ID '{entry_id}': {e}") if not isinstance(data, dict): raise DemistoException( f"Invalid filter JSON structure in entry ID '{entry_id}': " f"expected a dictionary, got {type(data).__name__}." ) demisto.debug(f"[Filter Parse] Parsed filter from file: {data}") return data if filter_json: try: data = json.loads(filter_json) except json.JSONDecodeError as e: raise DemistoException(f"Failed to parse filter_json: {e}") if not isinstance(data, dict): raise DemistoException(f"Invalid filter_json structure: expected a dictionary, got {type(data).__name__}.") demisto.debug(f"[Filter Parse] Parsed inline filter: {data}") return data raise DemistoException("Either 'filter_json' or 'filter_raw_json_entry_id' must be provided.") def parse_integration_params(params: dict[str, Any]) -> dict[str, Any]: """Parse and validate integration configuration parameters. Extracts connection settings from the raw demisto.params() dictionary and validates audit type filters if provided. Args: params: Raw parameters from demisto.params(). Returns: Validated configuration dictionary with keys: base_url, api_key, verify, proxy. Raises: DemistoException: If audit type filter contains invalid values. """ base_url = params.get("url", "https://api.prod.koi.security/").rstrip("/") api_key = params.get("api_key", {}) if isinstance(api_key, dict): api_key = api_key.get("password", "") verify_certificate = not argToBoolean(params.get("insecure", False)) proxy = argToBoolean(params.get("proxy", False)) # Validate audit types filter if provided audit_types_filter = argToList(params.get("audit_types_filter")) if audit_types_filter: invalid = [t for t in audit_types_filter if t not in VALID_AUDIT_TYPES] if invalid: raise DemistoException(f"Invalid audit log type(s): {invalid}. Valid types: {VALID_AUDIT_TYPES}") demisto.debug(f"[Config] URL: {base_url}") return { "base_url": base_url, "api_key": api_key, "verify": verify_certificate, "proxy": proxy, } # endregion # region Client # ================================= # Client # ================================= class Client(ContentClient): """KOI API client. Extends ContentClient with KOI-specific functionality including Bearer token authentication and API methods for alerts and audit logs. """ def __init__( self, base_url: str, api_key: str, verify: bool, proxy: bool, ): """Initialize the KOI client. Args: base_url: KOI API server URL. api_key: KOI API key for Bearer token authentication. verify: Whether to verify SSL certificates. proxy: Whether to use proxy settings. """ auth_handler = BearerTokenAuthHandler(token=api_key) retry_policy = RetryPolicy( # type: ignore[call-arg] max_attempts=4, retryable_status_codes=(429, 500, 502, 503, 504), ) super().__init__( base_url=base_url, verify=verify, proxy=proxy, auth_handler=auth_handler, client_name="KOI", timeout=60, retry_policy=retry_policy, ) def get_events_page( self, log_type: LogType, created_at_gte: str | None = None, created_at_lte: str | None = None, page: int = 1, page_size: int = Config.DEFAULT_PAGE_SIZE, audit_types: list[str] | None = None, ) -> list[dict]: """Fetch a single page of events from the KOI API. This is the single unified method used by all commands (test-module, fetch-events, get-events) to retrieve events from the API. Args: log_type: The LogType to fetch (ALERTS or AUDIT). created_at_gte: Filter events created at or after this datetime (ISO 8601). created_at_lte: Filter events created at or before this datetime (ISO 8601). page: Page number (1-based). page_size: Number of results per page (max 500). audit_types: Optional list of audit log types to filter by (only for AUDIT). Returns: List of event dictionaries from the API response. """ params: dict[str, Any] = { "page": page, "page_size": min(page_size, Config.MAX_PAGE_SIZE), "sort_direction": Config.SORT_DIRECTION, } if created_at_gte: params["created_at_gte"] = created_at_gte if created_at_lte: params["created_at_lte"] = created_at_lte if log_type == LogType.AUDIT and audit_types: params["types"] = ",".join(audit_types) demisto.debug(f"[API Fetch] {log_type.type_string} | Page: {page} | Params: {params}") response = self._http_request( method="GET", url_suffix=log_type.api_endpoint, params=params, ) events = response.get("alerts") or response.get("data") or response.get("items") or response.get("results") or [] demisto.debug(f"[API Fetch] {log_type.type_string} | Page {page}: {len(events)} events returned") return events def get_policies( self, page: int, page_size: int, ) -> dict[str, Any]: """Fetch a single page of policies from the Koi API. Args: page: Page number for pagination (1-based). page_size: Number of results per page (max 500). Returns: The full API response dictionary containing 'policies' list and 'total_count'. """ params: dict[str, Any] = { "page": page, "page_size": page_size, } demisto.debug(f"[API] Fetching policies | Params: {params}") response = self._http_request( method="GET", url_suffix=ApiPaths.POLICIES, params=params, ) demisto.debug("[API] Policies response received") return response def update_policy_status(self, policy_id: int, enabled: bool) -> dict[str, Any]: """Update the enabled/disabled status of a policy. Args: policy_id: The ID of the policy to update. enabled: Whether to enable (True) or disable (False) the policy. Returns: The full updated policy object from the API. """ url_suffix = ApiPaths.policy(policy_id) body: dict[str, Any] = {"enabled": enabled} demisto.debug(f"[API] Updating policy {policy_id} status to enabled={enabled}") response = self._http_request( method="PUT", url_suffix=url_suffix, json_data=body, ) demisto.debug(f"[API] Policy {policy_id} status updated successfully") return response def get_allowlist(self) -> dict[str, Any]: """Fetch all items in the allowlist from the Koi API. Returns: The full API response dictionary containing 'items' list. """ demisto.debug("[API] Fetching allowlist") response = self._http_request( method="GET", url_suffix=ApiPaths.ALLOWLIST, ) items = response.get("items", []) demisto.debug(f"[API] Allowlist response received: {len(items)} items") return response def get_blocklist(self) -> dict[str, Any]: """Fetch all items in the blocklist from the Koi API. Returns: The full API response dictionary containing 'items' list. """ demisto.debug("[API] Fetching blocklist") response = self._http_request( method="GET", url_suffix=ApiPaths.BLOCKLIST, ) items = response.get("items", []) demisto.debug(f"[API] Blocklist response received: {len(items)} items") return response def remove_allowlist_items( self, items: list[dict[str, Any]], ) -> None: """Remove one or more items from the global allowlist. Args: items: List of item dictionaries, each containing at least 'item_id' and 'marketplace'. """ body: dict[str, Any] = {"items": items} demisto.debug(f"[API] Removing {len(items)} allowlist item(s): {items}") self._http_request( method="DELETE", url_suffix=ApiPaths.ALLOWLIST, json_data=body, resp_type="response", ok_codes=(204,), ) demisto.debug(f"[API] Successfully removed {len(items)} allowlist item(s)") def add_allowlist_items( self, items: list[dict[str, Any]], ) -> None: """Add one or more items to the global allowlist. Args: items: List of item dictionaries, each containing at least 'item_id' and 'marketplace'. """ body: dict[str, Any] = {"items": items} demisto.debug(f"[API] Adding {len(items)} allowlist item(s): {items}") self._http_request( method="POST", url_suffix=ApiPaths.ALLOWLIST, json_data=body, resp_type="response", ok_codes=(204,), ) demisto.debug(f"[API] Successfully added {len(items)} allowlist item(s)") def remove_blocklist_items( self, items: list[dict[str, Any]], ) -> None: """Remove one or more items from the global blocklist. Args: items: List of item dictionaries, each containing at least 'item_id' and 'marketplace'. """ body: dict[str, Any] = {"items": items} demisto.debug(f"[API] Removing {len(items)} blocklist item(s): {items}") self._http_request( method="DELETE", url_suffix=ApiPaths.BLOCKLIST, json_data=body, resp_type="response", ok_codes=(204,), ) demisto.debug(f"[API] Successfully removed {len(items)} blocklist item(s)") def add_blocklist_items( self, items: list[dict[str, Any]], ) -> None: """Add one or more items to the global blocklist. Args: items: List of item dictionaries, each containing at least 'item_id' and 'marketplace'. """ body: dict[str, Any] = {"items": items} demisto.debug(f"[API] Adding {len(items)} blocklist item(s): {items}") self._http_request( method="POST", url_suffix=ApiPaths.BLOCKLIST, json_data=body, resp_type="response", ok_codes=(204,), ) demisto.debug(f"[API] Successfully added {len(items)} blocklist item(s)") def get_inventory( self, page: int, page_size: int, brew_category_koi: str | None = None, browser_category_koi: str | None = None, chocolatey_category_koi: str | None = None, device_id: str | None = None, finding_id: str | None = None, first_seen: str | None = None, ide_category_koi: str | None = None, installation_method: str | None = None, item_display_name: str | None = None, item_id: str | None = None, marketplace: str | None = None, platform: str | None = None, publisher_name: str | None = None, risk_level: str | None = None, software_category_koi: str | None = None, sort_by: str | None = None, sort_direction: str | None = None, view: str | None = None, ) -> dict[str, Any]: """Fetch a single page of inventory items from the Koi API. Args: page: Page number for pagination (1-based). page_size: Number of results per page (max 500). brew_category_koi: Filter by Homebrew package category (Koi classification). browser_category_koi: Filter by browser extension category (Koi classification). chocolatey_category_koi: Filter by Chocolatey package category (Koi classification). device_id: Filter devices by device id. finding_id: Filter devices by finding id. first_seen: Filter by first seen date (ISO 8601 format). ide_category_koi: Filter by IDE extension category (Koi classification). installation_method: Filter by installation method. item_display_name: Filter by item display name (case-insensitive partial match). item_id: Filter by item ID. marketplace: Filter by marketplace. platform: Filter by platform. publisher_name: Filter by publisher name (case-insensitive partial match). risk_level: Filter by risk level. software_category_koi: Filter by software category (Koi classification). sort_by: Column to sort by. sort_direction: Sort direction (asc or desc). view: Filter by predefined view (marketplace group). Returns: The full API response dictionary containing 'items' list and 'total_count'. """ params: dict[str, Any] = assign_params( page=page, page_size=page_size, brew_category_koi=brew_category_koi, browser_category_koi=browser_category_koi, chocolatey_category_koi=chocolatey_category_koi, device_id=device_id, finding_id=finding_id, first_seen=first_seen, ide_category_koi=ide_category_koi, installation_method=installation_method, item_display_name=item_display_name, item_id=item_id, marketplace=marketplace, platform=platform, publisher_name=publisher_name, risk_level=risk_level, software_category_koi=software_category_koi, sort_by=sort_by, sort_direction=sort_direction, view=view, ) demisto.debug(f"[API] Fetching inventory | Params: {params}") response = self._http_request( method="GET", url_suffix=ApiPaths.INVENTORY, params=params, ) demisto.debug("[API] Inventory response received") return response def get_inventory_item( self, item_id: str, marketplace: str, version: str, ) -> dict[str, Any]: """Fetch details for a specific inventory item from the Koi API. Args: item_id: Unique identifier for the item. marketplace: The marketplace where the item is hosted. version: The specific version of the item to retrieve. Returns: The full API response dictionary with item details. """ params: dict[str, Any] = { "marketplace": marketplace, "version": version, } url_suffix = ApiPaths.inventory_item(item_id) demisto.debug(f"[API] Fetching inventory item {item_id} | Params: {params}") response = self._http_request( method="GET", url_suffix=url_suffix, params=params, ) demisto.debug(f"[API] Inventory item {item_id} response received") return response def get_inventory_item_endpoints( self, item_id: str, marketplace: str, version: str, page: int, page_size: int, ) -> dict[str, Any]: """Fetch endpoints that have a specific inventory item installed. Args: item_id: Unique identifier for the item. marketplace: The marketplace where the item is hosted. version: The specific version of the item. page: Page number for pagination (1-based). page_size: Number of results per page (max 500). Returns: The full API response dictionary containing 'endpoints' list and 'total_count'. """ params: dict[str, Any] = { "marketplace": marketplace, "version": version, "page": page, "page_size": page_size, } url_suffix = ApiPaths.inventory_item_endpoints(item_id) demisto.debug(f"[API] Fetching endpoints for item {item_id} | Params: {params}") response = self._http_request( method="GET", url_suffix=url_suffix, params=params, ) demisto.debug(f"[API] Endpoints for item {item_id} response received") return response def search_inventory( self, page: int, page_size: int, filter_obj: dict[str, Any], sort_by: str | None = None, sort_direction: str | None = None, ) -> dict[str, Any]: """Search inventory items using advanced filters via POST. Args: page: Page number for pagination (1-based). page_size: Number of results per page (max 500). filter_obj: Filter object using query builder syntax. sort_by: Column to sort by. sort_direction: Sort direction (asc or desc). Returns: The full API response dictionary containing 'items' list and 'total_count'. """ body: dict[str, Any] = { "page": page, "page_size": page_size, "filter": filter_obj, } if sort_by: body["sort_by"] = sort_by if sort_direction: body["sort_direction"] = sort_direction demisto.debug(f"[API] Searching inventory | Body: {body}") response = self._http_request( method="POST", url_suffix=ApiPaths.INVENTORY_SEARCH, json_data=body, ) demisto.debug("[API] Inventory search response received") return response def send_events(self, events: list[dict]) -> None: """Send events to XSIAM using the ContentClient context. Wraps send_events_to_xsiam to keep event sending encapsulated within the client class for consistent logging and diagnostics. Args: events: List of event dicts to send. """ demisto.debug(f"[API] Sending {len(events)} events to XSIAM") send_events_to_xsiam(events=events, vendor=Config.VENDOR, product=Config.PRODUCT) demisto.debug(f"[API] Successfully sent {len(events)} events to XSIAM") # endregion # region Command implementations # ================================= # Command implementations # ================================= def test_module(client: Client) -> str: """Test API connectivity by fetching a small number of events. Args: client: The KOI client. Returns: 'ok' if test passed, otherwise raises an exception. """ demisto.debug("[Test Module] Starting...") try: utc_now = datetime.now(UTC) test_time = (utc_now - timedelta(minutes=Config.TEST_MODULE_LOOKBACK_MINUTES)).strftime(Config.DATE_FORMAT) demisto.debug(f"[Test Module] Fetching alerts from: {test_time}") fetch_events_with_pagination( client, log_type=LogType.ALERTS, created_after=test_time, max_events=Config.TEST_MODULE_MAX_EVENTS, ) demisto.debug("[Test Module] Success") return "ok" except Exception as error: error_msg = str(error) demisto.debug(f"[Test Module] Failed: {error_msg}") if "401" in error_msg or "403" in error_msg: return "Authorization Error: Verify your API Key." raise def fetch_events_with_pagination( client: Client, log_type: LogType, created_after: str, created_before: str | None = None, max_events: int = Config.DEFAULT_MAX_FETCH, audit_types: list[str] | None = None, ) -> list[dict]: """Fetch events with pagination support. This is the single unified pagination function used by all commands (test-module, fetch-events, get-events). Args: client: The KOI client. log_type: The LogType to fetch. created_after: Start time (ISO 8601). created_before: End time (ISO 8601) or None. max_events: Maximum number of events to fetch. audit_types: Optional list of audit log types to filter by. Returns: List of event dictionaries. """ events: list[dict] = [] page = 1 page_size = min(Config.MAX_PAGE_SIZE, max_events) demisto.debug( f"[Pagination Loop] Start | Type: {log_type.type_string} | Goal: {max_events} | " f"Time: {created_after} -> {created_before or 'Now'}" ) while len(events) < max_events: page_events = client.get_events_page( log_type=log_type, created_at_gte=created_after, created_at_lte=created_before, page=page, page_size=page_size, audit_types=audit_types if log_type == LogType.AUDIT else None, ) if not page_events: demisto.debug(f"[Pagination Loop] Page {page}: Empty. Stopping.") break events.extend(page_events) demisto.debug(f"[Pagination Loop] Page {page}: +{len(page_events)} events. Total: {len(events)}") if len(page_events) < page_size: demisto.debug("[Pagination Loop] Last page (partial). Stopping.") break page += 1 if page > Config.MAX_PAGES_PER_FETCH: demisto.debug(f"[Pagination Loop] Max page limit reached ({Config.MAX_PAGES_PER_FETCH}). Stopping.") break if len(events) >= max_events: demisto.debug(f"[Pagination Loop] Threshold reached ({len(events)} >= {max_events}). Stopping.") break # Slice to limit if len(events) > max_events: demisto.debug(f"[Pagination Result] Slicing {len(events)} events to limit {max_events}") events = events[:max_events] demisto.debug(f"[Pagination Result] Returning {len(events)} {log_type.type_string} events") return events def get_events_command(client: Client, args: dict, params: dict) -> CommandResults | str: """Manual command to get events for debugging/development. Args: client: The KOI client. args: Command arguments. params: Integration parameters. Returns: CommandResults or string message. """ demisto.debug("[Command] koi-get-events triggered") limit = int(args.get("limit", "50")) start_time_input = args.get("start_time", Config.DEFAULT_FROM_TIME) end_time_input = args.get("end_time") should_push_events = resolve_should_push_events(args) event_type_arg = argToList(args.get("event_type")) event_types_to_fetch = argToList(params.get("event_types_to_fetch", ["Alerts", "Audit"])) log_types = get_log_types_from_titles(event_type_arg if event_type_arg else event_types_to_fetch) created_after = get_formatted_utc_time(start_time_input) created_before = get_formatted_utc_time(end_time_input) if end_time_input else None audit_types_filter = argToList(params.get("audit_types_filter")) or None demisto.debug(f"[Command Params] From: {created_after}, To: {created_before}, Limit: {limit}, Push: {should_push_events}") all_events: list[dict] = [] for log_type in log_types: events = fetch_events_with_pagination( client, log_type=log_type, created_after=created_after, created_before=created_before, max_events=limit, audit_types=audit_types_filter if log_type == LogType.AUDIT else None, ) add_time_to_events(events, log_type) all_events.extend(events) demisto.debug(f"[Command Result] Total events retrieved: {len(all_events)}") if should_push_events and all_events: client.send_events(all_events) return f"Successfully retrieved and pushed {len(all_events)} events to XSIAM" readable_output = tableToMarkdown(f"{INTEGRATION_NAME} Events", all_events, removeNull=True) return CommandResults( readable_output=readable_output, outputs_prefix="KOI.Event", outputs_key_field="id", outputs=all_events, ) @dataclass class FetchResult: """Result of fetching events for a single log type.""" log_type: LogType new_events: list[dict] = field(default_factory=list) last_run_updates: dict[str, str | list[str]] = field(default_factory=dict) error: str | None = None def _fetch_single_log_type( client: Client, log_type: LogType, last_run: dict[str, str | list[str]], max_events: int, audit_types: list[str] | None, ) -> FetchResult: """Fetch and process events for a single log type. This function is executed in a separate thread by fetch_events_command via ThreadPoolExecutor, enabling parallel fetching of multiple log types. Each thread receives an immutable copy of last_run to avoid shared mutable state. The function handles its own errors — if an API call fails, the error is captured in FetchResult.error and the thread returns gracefully without affecting other threads. Thread safety: - Receives a dict copy of last_run (no shared mutable state). - Returns a FetchResult with last_run_updates (merged by the main thread after completion). - Uses demisto.debug() for logging (thread-safe in XSOAR runtime). Args: client: The KOI client (thread-safe — ContentClient uses httpx which is thread-safe). log_type: The LogType to fetch (ALERTS or AUDIT). last_run: Immutable copy of the current last_run state dict. max_events: Maximum events to fetch per type. audit_types: Optional audit type filter (only applied for AUDIT log type). Returns: FetchResult containing new_events, last_run_updates, and any error message. """ result = FetchResult(log_type=log_type) try: last_fetch_key = f"last_fetch_{log_type.type_string}" previous_ids_key = f"previous_ids_{log_type.type_string}" raw_timestamp = last_run.get(last_fetch_key) last_fetch_timestamp: str | None = raw_timestamp if isinstance(raw_timestamp, str) else None raw_ids = last_run.get(previous_ids_key) last_fetched_ids: list[str] = raw_ids if isinstance(raw_ids, list) else [] if last_fetch_timestamp: time_input = last_fetch_timestamp demisto.debug( f"[Fetch] {log_type.type_string}: Continuing from {time_input}. " f"Prev ID count: {len(last_fetched_ids)}" ) else: time_input = Config.DEFAULT_FROM_TIME demisto.debug(f"[Fetch] {log_type.type_string}: First run - starting from default time") created_after = get_formatted_utc_time(time_input) # Fetch events using the unified pagination function events = fetch_events_with_pagination( client, log_type=log_type, created_after=created_after, max_events=max_events, audit_types=audit_types if log_type == LogType.AUDIT else None, ) if not events: demisto.debug(f"[Fetch] {log_type.type_string}: No events found.") return result # Pre-compute time values to avoid redundant extract_time_from_event calls. # Events are already sorted chronologically by the API (sort_direction=asc). event_times: list[str] = [extract_time_from_event(event, log_type) or "" for event in events] # Deduplicate new_events = deduplicate_events(events, last_fetched_ids) if new_events: add_time_to_events(new_events, log_type) result.new_events = new_events demisto.debug(f"[Fetch] {log_type.type_string}: {len(new_events)} new events after dedup") else: demisto.debug(f"[Fetch] {log_type.type_string}: All events were duplicates.") # Update Last Run - always update based on ALL fetched events (not just new_events) new_last_run_time = event_times[-1] if event_times else None if new_last_run_time: # Collect IDs for the new high-water mark timestamp using pre-computed times ids_at_last_timestamp: list[str] = [ event_id for event, event_time in zip(events, event_times) if event_time == new_last_run_time and (event_id := get_event_id(event)) ] # If the HWM timestamp hasn't changed, merge with previous IDs to prevent duplicates if new_last_run_time == last_fetch_timestamp: ids_at_last_timestamp = list(set(last_fetched_ids) | set(ids_at_last_timestamp)) result.last_run_updates[last_fetch_key] = new_last_run_time result.last_run_updates[previous_ids_key] = ids_at_last_timestamp demisto.debug(f"[Fetch] {log_type.type_string}: State updated. New HWM: {new_last_run_time}") else: demisto.debug(f"[Fetch] {log_type.type_string}: Warning: Last event missing time. State not updated.") except Exception as e: result.error = str(e) demisto.debug(f"[Fetch] {log_type.type_string}: Error fetching events: {e!s}.") return result def fetch_events_command(client: Client) -> None: """Scheduled command to fetch events using parallel threads. Uses ThreadPoolExecutor to fetch all configured log types (Alerts, Audit) simultaneously. This ensures that if one type takes a long time or fails, the other type still completes within the XSOAR execution timeout. Architecture: 1. Single getLastRun() read at the start. 2. Each log type is fetched in a separate thread via _fetch_single_log_type(). Each thread receives an immutable copy of last_run (no shared mutable state). 3. After all threads complete, results are merged sequentially: - New events from successful types are collected. - last_run updates from successful types are applied. - Failed types are skipped (their previous state is preserved). 4. All events are sent to XSIAM in a single batch. 5. Single setLastRun() write at the end. Race condition prevention: - One getLastRun() call, one setLastRun() call. - Threads don't share mutable state — each gets a dict copy. - Merge happens after all threads complete (no concurrent writes). Args: client: The KOI client. """ params = demisto.params() max_events_to_fetch = int(params.get("max_fetch", Config.DEFAULT_MAX_FETCH)) event_types_to_fetch = argToList(params.get("event_types_to_fetch", ["Alerts", "Audit"])) log_types = get_log_types_from_titles(event_types_to_fetch) audit_types_filter = argToList(params.get("audit_types_filter")) or None # Single read of last_run state — no race condition last_run = demisto.getLastRun() demisto.debug(f"[Fetch] Starting with last_run: {last_run}") # Guard against an empty log_types selection — ThreadPoolExecutor(max_workers=0) raises ValueError. if not log_types: demisto.debug("[Fetch] No event types selected. Nothing to fetch. Preserving last_run as-is.") demisto.setLastRun(last_run) return # Fetch all log types in parallel so one slow type doesn't block the other results: list[FetchResult] = [] with ThreadPoolExecutor(max_workers=len(log_types)) as executor: futures = { executor.submit( _fetch_single_log_type, client=client, log_type=log_type, last_run=dict(last_run), max_events=max_events_to_fetch, audit_types=audit_types_filter, ): log_type for log_type in log_types } for future in as_completed(futures): log_type = futures[future] try: result = future.result() results.append(result) except Exception as e: demisto.debug(f"[Fetch] {log_type.type_string}: Thread failed: {e!s}") # Merge results — collect all new events and last_run updates all_new_events: list[dict] = [] updated_last_run: dict[str, str | list[str]] = dict(last_run) for result in results: if result.error: demisto.debug(f"[Fetch] {result.log_type.type_string}: Skipped due to error: {result.error}") continue all_new_events.extend(result.new_events) updated_last_run.update(result.last_run_updates) # Send all successfully fetched events to XSIAM if all_new_events: client.send_events(all_new_events) # Single write of last_run state — preserves progress from successful types demisto.setLastRun(updated_last_run) demisto.debug(f"[Fetch] Last run updated: {updated_last_run}") def koi_policy_list_command(client: Client, args: dict[str, Any]) -> CommandResults: """List policies with pagination support. Supports two modes: - Single page: provide 'page' and/or 'page_size' to fetch a specific page. - Auto-paginate: provide 'limit' to automatically paginate and collect up to 'limit' policies. If 'page' is provided, single-page mode is used (limit is ignored). If only 'limit' is provided, auto-pagination mode is used. Args: client: The KOI client. args: Command arguments (page, page_size, limit). Returns: CommandResults with the policy list. """ demisto.debug("[Command] koi-policy-list triggered") page_arg = arg_to_number(args.get("page")) page_size = arg_to_number(args.get("page_size")) or Config.DEFAULT_PAGE_SIZE limit_arg = arg_to_number(args.get("limit")) if page_size > Config.MAX_PAGE_SIZE: raise DemistoException(f"page_size ({page_size}) exceeds the maximum allowed value of {Config.MAX_PAGE_SIZE}.") if limit_arg and limit_arg > Config.MAX_LIMIT: raise DemistoException(f"limit ({limit_arg}) exceeds the maximum allowed value of {Config.MAX_LIMIT}.") if page_arg: # Single-page mode: fetch the requested page demisto.debug(f"[Command] Single-page mode: page={page_arg}, page_size={page_size}") response = client.get_policies(page=page_arg, page_size=page_size) policies = response.get("policies", []) total_count = response.get("total_count") demisto.debug(f"[Command Result] Retrieved {len(policies)} policies (total_count={total_count})") else: # Auto-paginate mode: fetch pages until limit is reached limit = limit_arg or Config.DEFAULT_LIMIT demisto.debug(f"[Command] Auto-paginate mode: limit={limit}") policies = _fetch_policies_with_pagination(client, limit=limit) readable_output = tableToMarkdown( f"{INTEGRATION_NAME} Policies", policies, headers=["id", "name", "description", "action", "enabled", "group_ids", "creator_fullname", "created_at", "updated_at"], headerTransform=string_to_table_header, ) return CommandResults( readable_output=readable_output, outputs_prefix="Koi.Policy", outputs_key_field="id", outputs=policies, ) def _fetch_policies_with_pagination( client: Client, limit: int, page_size: int = Config.MAX_PAGE_SIZE, ) -> list[dict]: """Auto-paginate through policies until limit is reached. Args: client: The Koi client. limit: Maximum total number of policies to collect. page_size: Number of results per API page. Returns: List of policy dictionaries. """ policies: list[dict] = [] page = Config.DEFAULT_PAGE while len(policies) < limit: response = client.get_policies(page=page, page_size=page_size) page_policies = response.get("policies", []) if not page_policies: demisto.debug(f"[Pagination] Page {page}: Empty. Stopping.") break policies.extend(page_policies) demisto.debug(f"[Pagination] Page {page}: +{len(page_policies)} policies. Total: {len(policies)}") if len(page_policies) < page_size: demisto.debug("[Pagination] Last page (partial). Stopping.") break page += 1 # Trim to limit if len(policies) > limit: demisto.debug(f"[Pagination] Trimming {len(policies)} policies to limit {limit}") policies = policies[:limit] demisto.debug(f"[Pagination] Returning {len(policies)} policies") return policies def koi_allowlist_get_command(client: Client, args: dict[str, Any]) -> CommandResults: """Retrieve all items in the allowlist. Args: client: The KOI client. args: Command arguments (unused, no inputs for this command). Returns: CommandResults with the allowlist items. """ demisto.debug("[Command] koi-allowlist-get triggered") response = client.get_allowlist() items = response.get("items", []) demisto.debug(f"[Command Result] Retrieved {len(items)} allowlist items") readable_output = tableToMarkdown( f"{INTEGRATION_NAME} Allowlist", items, headers=[ "item_id", "item_name", "item_display_name", "marketplace", "publisher_name", "package_name", "notes", "created_by", "created_at", ], headerTransform=string_to_table_header, ) return CommandResults( readable_output=readable_output, outputs_prefix="Koi.Allowlist", outputs_key_field="item_id", outputs=items, ) def koi_allowlist_items_remove_command(client: Client, args: dict[str, Any]) -> CommandResults: """Remove one or more items from the global allowlist. Supports two input modes: - Single item: provide 'item_id' and 'marketplace' (with optional 'created_by' and 'notes'). - Bulk from file: provide 'items_list_raw_json_entry_id' with a War Room entry ID of a JSON file containing a list of item objects. Args: client: The KOI client. args: Command arguments. Returns: CommandResults with a success message. """ demisto.debug("[Command] koi-allowlist-items-remove triggered") items = resolve_items_from_args(args) client.remove_allowlist_items(items) item_count = len(items) demisto.debug(f"[Command Result] {item_count} allowlist item(s) removed successfully") if item_count == 1: readable = f"Allowlist item '{items[0]['item_id']}' (marketplace: {items[0]['marketplace']}) was removed successfully." else: readable = f"{item_count} allowlist items were removed successfully." return CommandResults(readable_output=readable) def koi_allowlist_items_add_command(client: Client, args: dict[str, Any]) -> CommandResults: """Add one or more items to the global allowlist. Supports two input modes: - Single item: provide 'item_id' and 'marketplace' (with optional 'created_by' and 'notes'). - Bulk from file: provide 'items_list_raw_json_entry_id' with a War Room entry ID of a JSON file containing a list of item objects. Args: client: The KOI client. args: Command arguments. Returns: CommandResults with a success message. """ demisto.debug("[Command] koi-allowlist-items-add triggered") items = resolve_items_from_args(args) client.add_allowlist_items(items) item_count = len(items) demisto.debug(f"[Command Result] {item_count} allowlist item(s) added successfully") if item_count == 1: readable = f"Allowlist item '{items[0]['item_id']}' (marketplace: {items[0]['marketplace']}) was added successfully." else: readable = f"{item_count} allowlist items were added successfully." return CommandResults(readable_output=readable) def koi_blocklist_get_command(client: Client, args: dict[str, Any]) -> CommandResults: """Retrieve all items in the blocklist. Args: client: The KOI client. args: Command arguments (unused, no inputs for this command). Returns: CommandResults with the blocklist items. """ demisto.debug("[Command] koi-blocklist-get triggered") response = client.get_blocklist() items = response.get("items", []) demisto.debug(f"[Command Result] Retrieved {len(items)} blocklist items") readable_output = tableToMarkdown( f"{INTEGRATION_NAME} Blocklist", items, headers=[ "item_id", "item_name", "item_display_name", "marketplace", "publisher_name", "package_name", "notes", "created_by", "created_at", ], headerTransform=string_to_table_header, ) return CommandResults( readable_output=readable_output, outputs_prefix="Koi.Blocklist", outputs_key_field="item_id", outputs=items, ) def koi_blocklist_items_remove_command(client: Client, args: dict[str, Any]) -> CommandResults: """Remove one or more items from the global blocklist. Supports two input modes: - Single item: provide 'item_id' and 'marketplace' (with optional 'created_by' and 'notes'). - Bulk from file: provide 'items_list_raw_json_entry_id' with a War Room entry ID of a JSON file containing a list of item objects. Args: client: The KOI client. args: Command arguments. Returns: CommandResults with a success message. """ demisto.debug("[Command] koi-blocklist-items-remove triggered") items = resolve_items_from_args(args) client.remove_blocklist_items(items) item_count = len(items) demisto.debug(f"[Command Result] {item_count} blocklist item(s) removed successfully") if item_count == 1: readable = f"Blocklist item '{items[0]['item_id']}' (marketplace: {items[0]['marketplace']}) was removed successfully." else: readable = f"{item_count} blocklist items were removed successfully." return CommandResults(readable_output=readable) def koi_blocklist_items_add_command(client: Client, args: dict[str, Any]) -> CommandResults: """Add one or more items to the global blocklist. Supports two input modes: - Single item: provide 'item_id' and 'marketplace' (with optional 'created_by' and 'notes'). - Bulk from file: provide 'items_list_raw_json_entry_id' with a War Room entry ID of a JSON file containing a list of item objects. Args: client: The KOI client. args: Command arguments. Returns: CommandResults with a success message. """ demisto.debug("[Command] koi-blocklist-items-add triggered") items = resolve_items_from_args(args) client.add_blocklist_items(items) item_count = len(items) demisto.debug(f"[Command Result] {item_count} blocklist item(s) added successfully") if item_count == 1: readable = f"Blocklist item '{items[0]['item_id']}' (marketplace: {items[0]['marketplace']}) was added successfully." else: readable = f"{item_count} blocklist items were added successfully." return CommandResults(readable_output=readable) def koi_policy_status_update_command(client: Client, args: dict[str, Any]) -> CommandResults: """Enable or disable a policy by ID. Args: client: The KOI client. args: Command arguments (policy_id, enabled). Returns: CommandResults with the updated policy. """ demisto.debug("[Command] koi-policy-status-update triggered") policy_id = arg_to_number(args.get("policy_id")) if policy_id is None: raise DemistoException("policy_id is required and must be a valid integer.") enabled = argToBoolean(args.get("enabled")) response = client.update_policy_status(policy_id=policy_id, enabled=enabled) status_text = "enabled" if enabled else "disabled" demisto.debug(f"[Command Result] Policy {policy_id} {status_text} successfully") readable_output = tableToMarkdown( f"{INTEGRATION_NAME} Policy Updated", response, headers=[ "id", "name", "description", "action", "enabled", "group_ids", "creator_fullname", "created_at", "updated_at", ], headerTransform=string_to_table_header, ) return CommandResults( readable_output=readable_output, outputs_prefix="Koi.Policy", outputs_key_field="id", outputs=response, ) def koi_inventory_list_command(client: Client, args: dict[str, Any]) -> CommandResults: """List inventory items with pagination and filtering support. Supports two modes: - Single page: provide 'page' and/or 'page_size' to fetch a specific page. - Auto-paginate: provide 'limit' to automatically paginate and collect up to 'limit' items. If 'page' is provided, single-page mode is used (limit is ignored). If only 'limit' is provided, auto-pagination mode is used. Args: client: The KOI client. args: Command arguments including pagination and filter parameters. Returns: CommandResults with the inventory item list. """ demisto.debug("[Command] koi-inventory-list triggered") page_arg = arg_to_number(args.get("page")) page_size = arg_to_number(args.get("page_size")) or Config.DEFAULT_PAGE_SIZE limit_arg = arg_to_number(args.get("limit")) if page_size > Config.MAX_PAGE_SIZE: raise DemistoException(f"page_size ({page_size}) exceeds the maximum allowed value of {Config.MAX_PAGE_SIZE}.") if limit_arg and limit_arg > Config.MAX_LIMIT: raise DemistoException(f"limit ({limit_arg}) exceeds the maximum allowed value of {Config.MAX_LIMIT}.") # Extract filter arguments filter_kwargs: dict[str, Any] = assign_params( brew_category_koi=args.get("brew_category_koi"), browser_category_koi=args.get("browser_category_koi"), chocolatey_category_koi=args.get("chocolatey_category_koi"), device_id=args.get("device_id"), finding_id=args.get("finding_id"), first_seen=args.get("first_seen"), ide_category_koi=args.get("ide_category_koi"), installation_method=args.get("installation_method"), item_display_name=args.get("item_display_name"), item_id=args.get("item_id"), marketplace=args.get("marketplace"), platform=args.get("platform"), publisher_name=args.get("publisher_name"), risk_level=args.get("risk_level"), software_category_koi=args.get("software_category_koi"), sort_by=args.get("sort_by"), sort_direction=args.get("sort_direction"), view=args.get("view"), ) if page_arg: # Single-page mode: fetch the requested page demisto.debug(f"[Command] Single-page mode: page={page_arg}, page_size={page_size}") response = client.get_inventory(page=page_arg, page_size=page_size, **filter_kwargs) items = response.get("items", []) total_count = response.get("total_count") demisto.debug(f"[Command Result] Retrieved {len(items)} inventory items (total_count={total_count})") else: # Auto-paginate mode: fetch pages until limit is reached limit = limit_arg or Config.DEFAULT_LIMIT demisto.debug(f"[Command] Auto-paginate mode: limit={limit}") items = _fetch_inventory_with_pagination(client, limit=limit, filter_kwargs=filter_kwargs) readable_output = tableToMarkdown( f"{INTEGRATION_NAME} Inventory", items, headers=[ "item_id", "item_display_name", "marketplace", "platforms", "publisher_name", "risk", "risk_level", "version", "status", "endpoint_count", "installs_count", "installation_method", "is_first_party", "is_signed", "first_seen", "last_seen", "last_used", "released_at", "short_description", "categories", "findings", "brew_category_koi", "browser_category_koi", "chocolatey_category_koi", "ide_category_koi", "software_category_koi", "governed_details", ], headerTransform=string_to_table_header, ) return CommandResults( readable_output=readable_output, outputs_prefix="Koi.Inventory", outputs_key_field="item_id", outputs=items, ) def _fetch_inventory_with_pagination( client: Client, limit: int, filter_kwargs: dict[str, Any], page_size: int = Config.MAX_PAGE_SIZE, ) -> list[dict]: """Auto-paginate through inventory items until limit is reached. Args: client: The Koi client. limit: Maximum total number of items to collect. filter_kwargs: Filter parameters to pass to the API. page_size: Number of results per API page. Returns: List of inventory item dictionaries. """ items: list[dict] = [] page = Config.DEFAULT_PAGE while len(items) < limit: response = client.get_inventory(page=page, page_size=page_size, **filter_kwargs) page_items = response.get("items", []) if not page_items: demisto.debug(f"[Pagination] Page {page}: Empty. Stopping.") break items.extend(page_items) demisto.debug(f"[Pagination] Page {page}: +{len(page_items)} items. Total: {len(items)}") if len(page_items) < page_size: demisto.debug("[Pagination] Last page (partial). Stopping.") break page += 1 # Trim to limit if len(items) > limit: demisto.debug(f"[Pagination] Trimming {len(items)} items to limit {limit}") items = items[:limit] demisto.debug(f"[Pagination] Returning {len(items)} inventory items") return items def koi_inventory_item_get_command(client: Client, args: dict[str, Any]) -> CommandResults: """Retrieve comprehensive details for a specific inventory item. Args: client: The KOI client. args: Command arguments (item_id, marketplace, version). Returns: CommandResults with the inventory item details. """ demisto.debug("[Command] koi-inventory-item-get triggered") item_id: str = args["item_id"] marketplace: str = args["marketplace"] version: str = args["version"] response = client.get_inventory_item( item_id=item_id, marketplace=marketplace, version=version, ) demisto.debug(f"[Command Result] Retrieved inventory item {item_id}") readable_output = tableToMarkdown( f"{INTEGRATION_NAME} Inventory Item", response, headers=[ "item_id", "item_display_name", "marketplace", "platforms", "publisher_name", "risk", "risk_level", "version", "status", "endpoint_count", "installs_count", "installation_method", "is_first_party", "is_signed", "first_seen", "last_seen", "last_used", "released_at", "short_description", "categories", "findings", "brew_category_koi", "browser_category_koi", "chocolatey_category_koi", "ide_category_koi", "software_category_koi", "governed_details", ], headerTransform=string_to_table_header, ) return CommandResults( readable_output=readable_output, outputs_prefix="Koi.Inventory", outputs_key_field="item_id", outputs=response, ) def koi_inventory_search_command(client: Client, args: dict[str, Any]) -> CommandResults: """Search inventory items using advanced filters. Supports two modes: - Single page: provide 'page' and/or 'page_size' to fetch a specific page. - Auto-paginate: provide 'limit' to automatically paginate and collect up to 'limit' items. If 'page' is provided, single-page mode is used (limit is ignored). If only 'limit' is provided, auto-pagination mode is used. Args: client: The KOI client. args: Command arguments including filter, pagination, and sorting parameters. Returns: CommandResults with the search results. """ demisto.debug("[Command] koi-inventory-search triggered") page_arg = arg_to_number(args.get("page")) page_size = arg_to_number(args.get("page_size")) or Config.DEFAULT_PAGE_SIZE limit_arg = arg_to_number(args.get("limit")) if page_size > Config.MAX_PAGE_SIZE: raise DemistoException(f"page_size ({page_size}) exceeds the maximum allowed value of {Config.MAX_PAGE_SIZE}.") if limit_arg and limit_arg > Config.MAX_LIMIT: raise DemistoException(f"limit ({limit_arg}) exceeds the maximum allowed value of {Config.MAX_LIMIT}.") filter_obj: dict[str, Any] = parse_filter_from_args(args) sort_by: str | None = args.get("sort_by") sort_direction: str | None = args.get("sort_direction") if page_arg: # Single-page mode demisto.debug(f"[Command] Single-page mode: page={page_arg}, page_size={page_size}") response = client.search_inventory( page=page_arg, page_size=page_size, filter_obj=filter_obj, sort_by=sort_by, sort_direction=sort_direction, ) items = response.get("items", []) total_count = response.get("total_count") demisto.debug(f"[Command Result] Retrieved {len(items)} items (total_count={total_count})") else: # Auto-paginate mode limit = limit_arg or Config.DEFAULT_LIMIT demisto.debug(f"[Command] Auto-paginate mode: limit={limit}") items = _search_inventory_with_pagination( client, limit=limit, filter_obj=filter_obj, sort_by=sort_by, sort_direction=sort_direction, ) readable_output = tableToMarkdown( f"{INTEGRATION_NAME} Inventory Search", items, headers=[ "item_id", "item_display_name", "marketplace", "platforms", "publisher_name", "risk", "risk_level", "version", "status", "endpoint_count", "installs_count", "installation_method", "is_first_party", "is_signed", "first_seen", "last_seen", "last_used", "released_at", "short_description", "categories", "findings", "brew_category_koi", "browser_category_koi", "chocolatey_category_koi", "ide_category_koi", "software_category_koi", "governed_details", ], headerTransform=string_to_table_header, ) return CommandResults( readable_output=readable_output, outputs_prefix="Koi.Inventory", outputs_key_field="item_id", outputs=items, ) def _search_inventory_with_pagination( client: Client, limit: int, filter_obj: dict[str, Any], sort_by: str | None = None, sort_direction: str | None = None, page_size: int = Config.MAX_PAGE_SIZE, ) -> list[dict]: """Auto-paginate through inventory search results until limit is reached. Args: client: The Koi client. limit: Maximum total number of items to collect. filter_obj: Filter object for the search. sort_by: Column to sort by. sort_direction: Sort direction. page_size: Number of results per API page. Returns: List of inventory item dictionaries. """ items: list[dict] = [] page = Config.DEFAULT_PAGE while len(items) < limit: response = client.search_inventory( page=page, page_size=page_size, filter_obj=filter_obj, sort_by=sort_by, sort_direction=sort_direction, ) page_items = response.get("items", []) if not page_items: demisto.debug(f"[Pagination] Page {page}: Empty. Stopping.") break items.extend(page_items) demisto.debug(f"[Pagination] Page {page}: +{len(page_items)} items. Total: {len(items)}") if len(page_items) < page_size: demisto.debug("[Pagination] Last page (partial). Stopping.") break page += 1 # Trim to limit if len(items) > limit: demisto.debug(f"[Pagination] Trimming {len(items)} items to limit {limit}") items = items[:limit] demisto.debug(f"[Pagination] Returning {len(items)} inventory search results") return items def koi_inventory_item_endpoints_list_command(client: Client, args: dict[str, Any]) -> CommandResults: """List endpoints that have a specific inventory item installed. Supports two modes: - Single page: provide 'page' and/or 'page_size' to fetch a specific page. - Auto-paginate: provide 'limit' to automatically paginate and collect up to 'limit' endpoints. If 'page' is provided, single-page mode is used (limit is ignored). If only 'limit' is provided, auto-pagination mode is used. Args: client: The KOI client. args: Command arguments (item_id, marketplace, version, page, page_size, limit). Returns: CommandResults with the endpoint list. """ demisto.debug("[Command] koi-inventory-item-endpoints-list triggered") item_id: str = args["item_id"] marketplace: str = args["marketplace"] version: str = args["version"] page_arg = arg_to_number(args.get("page")) page_size = arg_to_number(args.get("page_size")) or Config.DEFAULT_PAGE_SIZE limit_arg = arg_to_number(args.get("limit")) if page_size > Config.MAX_PAGE_SIZE: raise DemistoException(f"page_size ({page_size}) exceeds the maximum allowed value of {Config.MAX_PAGE_SIZE}.") if limit_arg and limit_arg > Config.MAX_LIMIT: raise DemistoException(f"limit ({limit_arg}) exceeds the maximum allowed value of {Config.MAX_LIMIT}.") if page_arg: # Single-page mode demisto.debug(f"[Command] Single-page mode: page={page_arg}, page_size={page_size}") response = client.get_inventory_item_endpoints( item_id=item_id, marketplace=marketplace, version=version, page=page_arg, page_size=page_size, ) endpoints = response.get("endpoints", []) total_count = response.get("total_count") demisto.debug(f"[Command Result] Retrieved {len(endpoints)} endpoints (total_count={total_count})") else: # Auto-paginate mode limit = limit_arg or Config.DEFAULT_LIMIT demisto.debug(f"[Command] Auto-paginate mode: limit={limit}") endpoints = _fetch_item_endpoints_with_pagination( client, item_id=item_id, marketplace=marketplace, version=version, limit=limit, ) readable_output = tableToMarkdown( f"{INTEGRATION_NAME} Inventory Item Endpoints", endpoints, headers=[ "id", "hostname", "os", "platform", "serial", "last_logged_on_user", "activation_status", "path", "first_seen", "last_seen", ], headerTransform=string_to_table_header, ) return CommandResults( readable_output=readable_output, outputs_prefix="Koi.Inventory.Endpoint", outputs_key_field="id", outputs=endpoints, ) def _fetch_item_endpoints_with_pagination( client: Client, item_id: str, marketplace: str, version: str, limit: int, page_size: int = Config.MAX_PAGE_SIZE, ) -> list[dict]: """Auto-paginate through item endpoints until limit is reached. Args: client: The Koi client. item_id: Unique identifier for the item. marketplace: The marketplace where the item is hosted. version: The specific version of the item. limit: Maximum total number of endpoints to collect. page_size: Number of results per API page. Returns: List of endpoint dictionaries. """ endpoints: list[dict] = [] page = Config.DEFAULT_PAGE while len(endpoints) < limit: response = client.get_inventory_item_endpoints( item_id=item_id, marketplace=marketplace, version=version, page=page, page_size=page_size, ) page_endpoints = response.get("endpoints", []) if not page_endpoints: demisto.debug(f"[Pagination] Page {page}: Empty. Stopping.") break endpoints.extend(page_endpoints) demisto.debug(f"[Pagination] Page {page}: +{len(page_endpoints)} endpoints. Total: {len(endpoints)}") if len(page_endpoints) < page_size: demisto.debug("[Pagination] Last page (partial). Stopping.") break page += 1 # Trim to limit if len(endpoints) > limit: demisto.debug(f"[Pagination] Trimming {len(endpoints)} endpoints to limit {limit}") endpoints = endpoints[:limit] demisto.debug(f"[Pagination] Returning {len(endpoints)} endpoints") return endpoints # endregion # region Main router # ================================= # Main router # ================================= COMMAND_MAP: dict[str, Any] = { "test-module": test_module, "koi-get-events": get_events_command, "fetch-events": fetch_events_command, "koi-policy-list": koi_policy_list_command, "koi-allowlist-get": koi_allowlist_get_command, "koi-allowlist-items-remove": koi_allowlist_items_remove_command, "koi-allowlist-items-add": koi_allowlist_items_add_command, "koi-blocklist-get": koi_blocklist_get_command, "koi-blocklist-items-remove": koi_blocklist_items_remove_command, "koi-blocklist-items-add": koi_blocklist_items_add_command, "koi-policy-status-update": koi_policy_status_update_command, "koi-inventory-list": koi_inventory_list_command, "koi-inventory-item-get": koi_inventory_item_get_command, "koi-inventory-search": koi_inventory_search_command, "koi-inventory-item-endpoints-list": koi_inventory_item_endpoints_list_command, } def main() -> None: """Main entry point for KOI integration.""" demisto.debug(f"{INTEGRATION_NAME} integration started") command = demisto.command() try: if command not in COMMAND_MAP: raise DemistoException(f"Command '{command}' is not implemented") params = demisto.params() args = demisto.args() config = parse_integration_params(params) client = Client( base_url=config["base_url"], api_key=config["api_key"], verify=config["verify"], proxy=config["proxy"], ) command_func = COMMAND_MAP[command] if command == "test-module": result = command_func(client) return_results(result) elif command == "fetch-events": command_func(client) elif command == "koi-get-events": result = command_func(client, args, params) return_results(result) else: result = command_func(client, args) return_results(result) except Exception as error: error_msg = f"Failed to execute {command}. Error: {error!s}" demisto.error(f"{error_msg}\n{traceback.format_exc()}") return_error(error_msg) demisto.debug(f"{INTEGRATION_NAME} integration finished") if __name__ in ("__main__", "__builtin__", "builtins"): main()