from dataclasses import dataclass from datetime import datetime, timedelta, UTC import time from typing import Any, Literal, TypedDict import demistomock as demisto import requests import urllib3 from CommonServerPython import * from CommonServerUserPython import * from pydantic import BaseModel, Field from pydantic_core import ValidationError urllib3.disable_warnings() XSOAR_DATE_FORMAT = "%Y-%m-%dT%H:%M:%SZ" NZ_DATE_FORMAT = "%Y-%m-%dT%H:%M:%S" MAX_PAGE_SIZE = 100 DEFAULT_MAX_FETCH = 200 DEFAULT_FIRST_FETCH = "7 days" AUTH_ENDPOINT = "v1/auth" GRAPHQL_ENDPOINT = "v1/graphql" HELLO_WORLD_QUERY = "query HelloWorld { hello }" WEAKNESSES_PAGE_QUERY = """ query WeaknessesPage($page_input: PageInput!) { weaknesses_page(page_input: $page_input) { weaknesses { uuid created_at vuln_id vuln_name vuln_short_name vuln_category vuln_cisa_kev vuln_known_ransomware_campaign_use ip has_proof score severity affected_asset_uuid affected_asset_display_name attack_paths_count op_id } page_info { page_num page_size } } } """ class Credentials(BaseModel): password: str class Params(BaseModel): url: str credentials: Credentials insecure: bool = False proxy: bool = False max_fetch: int = DEFAULT_MAX_FETCH first_fetch: str = DEFAULT_FIRST_FETCH class Token(BaseModel): token: str class Incident(TypedDict): name: str occurred: str rawJSON: str dbotMirrorId: str class Weakness(BaseModel): uuid: str created_at: str vuln_id: str vuln_name: str | None vuln_short_name: str | None vuln_category: str | None vuln_cisa_kev: bool | None vuln_known_ransomware_campaign_use: bool | None ip: str | None has_proof: bool | None score: float | None severity: str | None affected_asset_uuid: str | None affected_asset_display_name: str | None attack_paths_count: int op_id: str | None def to_incident(self) -> Incident: """Convert a Weakness to an XSOAR Incident.""" asset_name = self.affected_asset_display_name or "Unknown Asset" vuln_name = self.vuln_short_name or self.vuln_name or self.vuln_id data = self.model_dump() data["incident_type"] = "NodeZero Weakness" return Incident( name=f"{vuln_name} on {asset_name}", occurred=self.created_at if self.created_at.endswith("Z") else f"{self.created_at}Z", dbotMirrorId=self.uuid, rawJSON=json.dumps(data), ) class PageInfo(BaseModel): page_num: int page_size: int class WeaknessesPage(BaseModel): weaknesses: list[Weakness] page_info: PageInfo @classmethod def from_graphql_response(cls, response: dict) -> "WeaknessesPage": data = response.get("data", {}).get("weaknesses_page", {}) return cls( weaknesses=[Weakness.model_validate(w) for w in data.get("weaknesses", [])], page_info=PageInfo.model_validate(data.get("page_info", {})), ) def _default_since_date() -> str: return (datetime.now(UTC) - timedelta(days=7)).strftime(NZ_DATE_FORMAT) class GetWeaknessesArgs(BaseModel): """Parsed arguments for the nodezero-get-weaknesses command.""" since_date: str = Field(default_factory=_default_since_date) limit: int = Field(default=50, le=1000) class LastRun(BaseModel): """State persisted between fetch-incidents runs.""" last_fetch_date: str | None = None # ISO 8601 format last_ids: list[str] = [] # UUIDs of incidents at last_fetch_date timestamp @dataclass(frozen=True) class UnauthenticatedError(Exception): message: str = "Authentication failed: invalid or expired API key" @dataclass(frozen=True) class TokenExpiredError(Exception): pass @dataclass(frozen=True) class ResponseValidationError(Exception): validation_error: ValidationError @dataclass(frozen=True) class ParametersValidationError(Exception): validation_error: ValidationError class Client(BaseClient): def __init__( self, base_url: str, api_key: str, verify: bool, proxy: bool, ): super().__init__(base_url=base_url, verify=verify, proxy=proxy) self._base_url: str = base_url self._verify: bool = verify self._api_key: str = api_key self._jwt: str | None = None self._expiry: int = 0 self.load_integration_context() def load_integration_context(self) -> None: context = get_integration_context() if context: self._jwt = context["jwt"] self._expiry = context["expiry"] def is_authenticated(self) -> bool: return self._jwt is not None and self._expiry > int(time.time()) def authenticate(self) -> None: url = f"{self._base_url}/{AUTH_ENDPOINT}" json = {"key": self._api_key} now = int(time.time()) response = self._session.post(url, json=json) self._raise_for_status(response) content = Token.model_validate_json(response.content) self._jwt = content.token self._expiry = now + 3600 - 60 # 1 hr minus 60s buffer to prevent race conditions _ = set_integration_context({"jwt": self._jwt, "expiry": self._expiry}) def _graphql_request(self, query: str, variables: dict[str, Any]) -> requests.Response: url = f"{self._base_url}/{GRAPHQL_ENDPOINT}" headers = {"Authorization": f"Bearer {self._jwt}"} json = {"query": query, "variables": variables} resp = self._session.post(url, json=json, headers=headers) self._raise_for_status(resp) return resp def _raise_for_status(self, response: requests.Response) -> None: if response.status_code == 401: raise UnauthenticatedError response.raise_for_status() def hello_world(self) -> None: """Send hello_world request to check availability. If no error is raised, the call was successful """ _ = self._graphql_request(HELLO_WORLD_QUERY, {}) def query_weaknesses_page( self, page_num: int, page_size: int, since_date: str, ) -> WeaknessesPage: """Fetch a page of HIGH/CRITICAL OPEN weaknesses since the given date. Args: page_num: 1-indexed page number page_size: Number of results per page since_date: Fetch weaknesses first seen on or after this date (YYYY-MM-DD) Returns: WeaknessesPage with weaknesses list and page_info """ variables = { "page_input": { "page_num": page_num, "page_size": page_size, "filter_by_inputs": [ {"field_name": "severity", "values": ["HIGH", "CRITICAL"]}, {"field_name": "created_at", "greater_than_or_equal": since_date}, ], "order_by": "created_at", "sort_order": "ASC", } } resp = self._graphql_request(WEAKNESSES_PAGE_QUERY, variables) try: return WeaknessesPage.from_graphql_response(resp.json()) except ValidationError as e: raise ResponseValidationError(e) def authenticate(client: Client): if not client.is_authenticated(): client.authenticate() def fetch_all_weaknesses_pages( client: Client, since_date: str, max_fetch: int = DEFAULT_MAX_FETCH, ) -> list[Weakness]: """Fetch all HIGH/CRITICAL OPEN weaknesses since the given date, paginating as needed. Args: client: NodeZero API client since_date: Fetch weaknesses first seen on or after this date (YYYY-MM-DD) max_fetch: Maximum number of weaknesses to fetch Returns: List of WeaknessesPages, up to max_fetch items """ all_weaknesses: list[Weakness] = [] page_num = 1 page_size = min(MAX_PAGE_SIZE, max_fetch) while len(all_weaknesses) < max_fetch: page = client.query_weaknesses_page( page_num=page_num, page_size=page_size, since_date=since_date, ) all_weaknesses.extend(page.weaknesses) if len(page.weaknesses) < page_size: break page_num += 1 return all_weaknesses[:max_fetch] """ COMMAND FUNCTIONS """ def get_weaknesses_command( client: Client, args: GetWeaknessesArgs, ) -> CommandResults: """Return a table of HIGH/CRITICAL weaknesses from NodeZero. Args: client: NodeZero API client. args: Parsed command arguments. Returns: CommandResults with human-readable table and context data. """ authenticate(client) weaknesses = fetch_all_weaknesses_pages( client=client, since_date=args.since_date, max_fetch=args.limit, ) outputs = [w.model_dump() for w in weaknesses] headers = [ "uuid", "severity", "vuln_id", "vuln_short_name", "ip", "affected_asset_display_name", "score", "created_at", ] return CommandResults( readable_output=tableToMarkdown( "NodeZero Weaknesses", outputs, headers=headers, removeNull=True, ), outputs_prefix="NodeZero.Weakness", outputs_key_field="uuid", outputs=outputs, ) def _test_module(client: Client) -> Literal["ok"]: """Tests API connectivity and authentication. Returning 'ok' indicates that the integration works like it is supposed to. Connection to the service is successful. Raises: exceptions if something goes wrong. Args: client: client to use Returns: 'ok' if test passed, anything else will fail the test. """ authenticate(client) client.hello_world() return "ok" def dedup_by_ids(weaknesses: list[Weakness], ids_to_skip: list[str]) -> list[Weakness]: """Filter out weaknesses that were already fetched in a previous run. Args: weaknesses: List of weaknesses to filter ids_to_skip: UUIDs of weaknesses to exclude Returns: List of weaknesses not in ids_to_skip """ if not ids_to_skip: return weaknesses skip_set = set(ids_to_skip) return [w for w in weaknesses if w.uuid not in skip_set] def fetch_incidents( client: Client, max_fetch: int = DEFAULT_MAX_FETCH, first_fetch_time: datetime | None = None, ) -> tuple[list[Incident], LastRun]: """Fetch HIGH/CRITICAL OPEN weaknesses as incidents. On first run, fetches weaknesses from `first_fetch_time`. On subsequent runs, fetches weaknesses since the last fetch date. Args: client: NodeZero API client max_fetch: Maximum number of incidents to fetch first_fetch_time: Datetime to fetch from on first run Returns: Tuple of (incidents list, new last_run state) """ authenticate(client) last_run = LastRun.model_validate(demisto.getLastRun()) if last_run.last_fetch_date: since_date = last_run.last_fetch_date elif first_fetch_time: since_date = first_fetch_time.strftime(NZ_DATE_FORMAT) else: since_date = (datetime.now(UTC) - timedelta(days=7)).strftime(NZ_DATE_FORMAT) demisto.debug(f"Fetching weaknesses since {since_date}") weaknesses = fetch_all_weaknesses_pages( client=client, since_date=since_date, max_fetch=max_fetch, ) pre_dedup_count = len(weaknesses) weaknesses = dedup_by_ids(weaknesses, last_run.last_ids) demisto.debug(f"After deduplication: {len(weaknesses)} weaknesses (filtered {pre_dedup_count - len(weaknesses)})") if not weaknesses: if pre_dedup_count > 0 and last_run.last_fetch_date: # All results were duplicates — advance timestamp by 1s to prevent infinite loop advanced = datetime.strptime(last_run.last_fetch_date, NZ_DATE_FORMAT) + timedelta(seconds=1) next_run = LastRun(last_fetch_date=advanced.strftime(NZ_DATE_FORMAT), last_ids=[]) demisto.debug(f"All results deduped; advancing fetch date to {next_run.last_fetch_date}") return [], next_run return [], last_run incidents = [w.to_incident() for w in weaknesses] latest_timestamp = max(w.created_at for w in weaknesses) last_ids = [w.uuid for w in weaknesses if w.created_at == latest_timestamp] next_run = LastRun(last_fetch_date=latest_timestamp, last_ids=last_ids) demisto.debug(f"Fetched {len(incidents)} incidents, next fetch from {latest_timestamp}") return incidents, next_run def params() -> Params: try: return Params.model_validate(demisto.params()) except ValidationError as e: raise ParametersValidationError(e) def main(): _args = demisto.args() command = demisto.command() try: params = Params.model_validate(demisto.params()) demisto.debug(f"Command being called is {command}") client = Client( base_url=params.url, api_key=params.credentials.password, verify=not params.insecure, proxy=params.proxy, ) if command == "test-module": return_results(_test_module(client)) elif command == "nodezero-get-weaknesses": weakness_args = GetWeaknessesArgs.model_validate(_args) return_results(get_weaknesses_command(client, weakness_args)) elif command == "fetch-incidents": first_fetch_time = arg_to_datetime( arg=params.first_fetch, arg_name="First fetch time", required=True, ) incidents, next_run = fetch_incidents( client, max_fetch=params.max_fetch, first_fetch_time=first_fetch_time, ) demisto.setLastRun(next_run.model_dump()) demisto.incidents(incidents) # type: ignore[arg-type] else: raise NotImplementedError(f"Command {command} is not implemented") except Exception as e: return_error(f"Failed to execute {command} command.\nError:\n{str(e)}") if __name__ in ("__main__", "__builtin__", "builtins"): # pragma: no cover main()