DHS Feed v2

import demistomock as demisto  # noqa: F401
from CommonServerPython import *  # noqa: F401
from pytz import utc
from TAXII2ApiModule import *  # noqa: E402
from taxii2client.common import _ensure_datetime_to_string

""" CONSTANTS """

COMPLEX_OBSERVATION_MODE_SKIP = "Skip indicators with more than a single observation"
MAX_FETCH_INTERVAL = "48 hours"
DEFAULT_FETCH_INTERVAL = "24 hours"
DEFAULT_LIMIT_PER_REQUEST = 1000  # DHS default limit

""" HELPER FUNCTIONS """


def get_datetime(given_interval: Union[str, datetime]) -> datetime:
    """
    Receives an interval and returns the corresponding datetime.
    """
    if isinstance(given_interval, datetime):
        return given_interval
    date = dateparser.parse(given_interval, date_formats=[TAXII_TIME_FORMAT])
    if not date:
        raise DemistoException("Given time interval is not in a valid format.")
    return date.replace(tzinfo=utc)  # type: ignore[union-attr]


def get_limited_interval(
    given_interval: Union[str, datetime], fetch_interval: Optional[Union[str, datetime]] = MAX_FETCH_INTERVAL
) -> datetime:
    """
    Returns the closer time between the two time intervals given.
    """
    given_interval: datetime = get_datetime(given_interval)
    fetch_interval: datetime = get_datetime(fetch_interval or MAX_FETCH_INTERVAL)
    return max(given_interval, fetch_interval)  # later time is bigger


def fetch_one_collection(client: Taxii2FeedClient, limit: int, initial_interval: datetime, last_run_ctx: Optional[dict] = None):
    demisto.debug(f"Fetching collection {client.collection_to_fetch.id=}")
    last_fetch_time = last_run_ctx.get(client.collection_to_fetch.id) if last_run_ctx else None
    # initial_interval gets here limited so no need to check limitation with default value
    added_after: datetime = get_limited_interval(initial_interval, last_fetch_time)

    indicators = client.build_iterator(limit, added_after=added_after)

    if last_run_ctx is not None:  # in case we got {}, we want to set it because we are in fetch incident run
        last_run_ctx[client.collection_to_fetch.id] = _ensure_datetime_to_string(
            client.last_fetched_indicator__modified if client.last_fetched_indicator__modified else added_after
        )

    return indicators, last_run_ctx


def fetch_all_collections(client: Taxii2FeedClient, limit: int, initial_interval: datetime, last_run_ctx: Optional[dict] = None):
    indicators: list = []
    for collection in client.collections:  # type: ignore[attr-defined]
        client.collection_to_fetch = collection
        fetched_iocs, last_run_ctx = fetch_one_collection(client, limit, initial_interval, last_run_ctx)
        indicators.extend(fetched_iocs)

        if limit >= 0:
            limit -= len(fetched_iocs)
            if limit <= 0:
                break

    return indicators, last_run_ctx


""" COMMAND FUNCTIONS """


def command_test_module(client: Taxii2FeedClient, initial_interval: str):
    if get_datetime(MAX_FETCH_INTERVAL) > get_datetime(initial_interval):
        return 'Due to DHS API limitations, "First Fetch Time" is limited to 48 hours.'

    try:
        client.initialise()
        if client.collections:
            get_indicators_command(client, {"limit": "1", "added_after": get_limited_interval("6 hours", initial_interval)})
        else:
            return "Could not connect to server"
    except requests.exceptions.ConnectTimeout:
        return (
            "Connection Timeout Error - potential reasons might be that the 'Discovery Service URL' parameter"
            " is incorrect or that the server is not accessible from your host."
        )
    except requests.exceptions.SSLError:
        return (
            "SSL Certificate Verification Failed - try selecting 'Trust any certificate' checkbox in"
            " the instance configuration. If this doesn't work, verify that your certificate and key are valid  and matching."
        )
    except requests.exceptions.ProxyError:
        return (
            "Proxy Error - if the 'Use system proxy' checkbox in the integration configuration is"
            " selected, try clearing the checkbox."
        )
    except requests.exceptions.ConnectionError:
        return (
            "Verify that the server URL parameter is correct and that you have access to the server from your host."
            " Run the test again."
        )
    except requests.HTTPError:
        return (
            "HTTP error - check your certificate and key, and that you are trying to reach a valid URL and API root."
            " Wait and run the test again."
        )

    return "ok"


def fetch_indicators_command(
    client: Taxii2FeedClient, limit: int, last_run_ctx: dict, initial_interval: str = DEFAULT_FETCH_INTERVAL
) -> tuple[list, dict]:
    """
    Fetch indicators from TAXII 2 server
    :param client: Taxii2FeedClient
    :param limit: upper limit of indicators to fetch
    :param last_run_ctx: last run dict with {collection_id: last_run_time string}
    :param initial_interval: initial interval in human readable format
    :return: indicators in cortex TIM format, updated last_run_ctx
    """
    initial_interval: datetime = get_limited_interval(get_datetime(initial_interval or DEFAULT_FETCH_INTERVAL))

    if client.collection_to_fetch:
        indicators, last_run_ctx = fetch_one_collection(client, limit, initial_interval, last_run_ctx)  # type: ignore[arg-type]
    else:
        indicators, last_run_ctx = fetch_all_collections(client, limit, initial_interval, last_run_ctx)  # type: ignore[arg-type]

    return indicators, last_run_ctx


def get_indicators_command(
    client: Taxii2FeedClient, args: Dict[str, Any]
) -> Union[CommandResults, Dict[str, List[Optional[str]]]]:
    """
    Fetch indicators from TAXII 2 server
    :param client: Taxii2FeedClient
    :param args: Dict that holds
        raw: When set to 'true' will return only rawJSON
        limit: upper limit of indicators to fetch
        added_after: added after time string in parse_date_range format
    :return: indicators in cortex TIM format
    """
    limit = arg_to_number(args.get("limit")) or 10
    raw = argToBoolean(args.get("raw", "false"))
    max_fetch_datetime = get_datetime(MAX_FETCH_INTERVAL)
    added_after: datetime = get_datetime(args.get("added_after", DEFAULT_FETCH_INTERVAL))
    if max_fetch_datetime > added_after:
        raise DemistoException('Due to DHS API limitations, "added_after" is limited to 48 hours.')

    if client.collection_to_fetch:
        indicators = client.build_iterator(limit, added_after=added_after)
    else:
        indicators, _ = fetch_all_collections(client, limit, added_after)  # type: ignore[arg-type]

    if raw:
        return {"indicators": [x.get("rawJSON") for x in indicators]}

    return CommandResults(
        readable_output=f"Found {len(indicators)} results added after {_ensure_datetime_to_string(added_after)} UTC:\n"
        + tableToMarkdown(name="DHS Indicators", t=indicators, headers=["value", "type"], removeNull=True),
        outputs_prefix="DHS.Indicators",
        outputs_key_field="value",
        outputs=indicators,
        raw_response=indicators,
    )


def get_collections_command(client: Taxii2FeedClient) -> CommandResults:
    """
    Get the available collections in the DHS server
    """
    collections = []
    for collection in client.collections:  # type: ignore[attr-defined]
        collections.append({"Name": collection.title, "ID": collection.id})
    return CommandResults(
        readable_output=tableToMarkdown("DHS Server Collections", t=collections, headers=["Name", "ID"]),
        outputs_prefix="DHS.Collections",
        outputs_key_field="ID",
        outputs=collections,
    )


""" MAIN FUNCTION """


def main():  # pragma: no cover
    params = demisto.params()
    url = params.get("url", "https://ais2.cisa.dhs.gov/taxii2/")
    key = params.get("key", {}).get("password")
    certificate = params.get("certificate")
    verify_certificate = not params.get("insecure", False)
    proxies = handle_proxy()

    collection_to_fetch = params.get("collection_to_fetch")
    skip_complex_mode = params.get("observation_operator_mode") == COMPLEX_OBSERVATION_MODE_SKIP
    feed_tags = argToList(params.get("feedTags"))
    tlp_color = params.get("tlp_color", "")
    objects_to_fetch = params.get("objects_to_fetch", [])

    initial_interval = params.get("initial_interval", DEFAULT_FETCH_INTERVAL)
    limit = arg_to_number(params.get("limit", -1)) or -1
    limit_per_request = arg_to_number(params.get("limit_per_request", DEFAULT_LIMIT_PER_REQUEST)) or DEFAULT_LIMIT_PER_REQUEST
    default_api_root = params.get("default_api_root", "public")

    command = demisto.command()
    demisto.info(f"Command being called is {command}")

    try:
        client = Taxii2FeedClient(
            url=url,
            collection_to_fetch=collection_to_fetch,
            proxies=proxies,
            verify=verify_certificate,
            objects_to_fetch=objects_to_fetch,
            skip_complex_mode=skip_complex_mode,
            tags=feed_tags,
            limit_per_request=limit_per_request,
            tlp_color=tlp_color,
            certificate=certificate,
            key=key,
            default_api_root=default_api_root,
        )

        start_time = time.time()
        if command == "test-module":
            return_results(command_test_module(client, initial_interval))

        elif command == "fetch-indicators":
            client.initialise()
            last_run_indicators = demisto.getLastRun()
            indicators, last_run_indicators = fetch_indicators_command(client, limit, last_run_indicators, initial_interval)
            for iter_ in batch(indicators, batch_size=2000):
                demisto.createIndicators(iter_)

            demisto.setLastRun(last_run_indicators)

        elif command == "dhs-get-indicators":
            client.initialise()
            return_results(get_indicators_command(client, demisto.args()))

        elif command == "dhs-get-collections":
            client.initialise()
            return_results(get_collections_command(client))

        else:
            raise NotImplementedError(f"{command} command is not implemented.")

        demisto.debug(f"Running {command} took {round(time.time() - start_time)}sec")

    except Exception as error:
        error_msg = str(error)
        if isinstance(error, requests.exceptions.SSLError):
            error_msg = (
                "Encountered an HTTPS certificate error. This error can be ignored by enabling "
                '"Trust any certificate (not secure)" in the instance configuration.'
            )
        elif isinstance(error, requests.HTTPError):
            error_msg = (
                "Encountered an HTTP error. Please check your certificate and key, and that you are trying to reach a "
                'valid URL and API root. If this occurs when the test works, increase the "Max STIX Objects Per Poll" '
                'in the instance configuration, reduce the "Max Indicators Per Fetch" in the instance configuration or '
                'reduce the "limit" in the command argument.'
            )
        return_error(error_msg, error)


""" ENTRY POINT """

if __name__ in ("__main__", "builtins"):
    main()