diff --git a/feedly/api_client/enterprise/indicators_of_compromise.py b/feedly/api_client/enterprise/indicators_of_compromise.py index 3526844..8d2484d 100644 --- a/feedly/api_client/enterprise/indicators_of_compromise.py +++ b/feedly/api_client/enterprise/indicators_of_compromise.py @@ -4,9 +4,9 @@ from datetime import datetime from enum import Enum from itertools import chain +from requests import Response from typing import ClassVar, Dict, Generic, Iterable, List, Optional, TypeVar from urllib.parse import parse_qs -from requests import Response from feedly.api_client.data import Streamable from feedly.api_client.session import FeedlySession @@ -24,38 +24,61 @@ class IoCDownloaderABC(ABC, Generic[T]): RELATIVE_URL = "/v3/enterprise/ioc" FORMAT: ClassVar[str] - def __init__(self, session: FeedlySession, newer_than: Optional[datetime], stream_id: str): + def __init__( + self, + session: FeedlySession, + newer_than: Optional[datetime], + older_than: Optional[datetime], + stream_id: str, + max_batches: Optional[int] = None, + ): """ Use this class to export the contextualized IoCs from a stream. Enterprise/personals feeds/boards are supported (see dedicated methods below). The IoCs will be returned along with their context and relationships in a dictionary representing a valid STIX v2.1 Bundle object. https://docs.oasis-open.org/cti/stix/v2.1/os/stix-v2.1-os.html#_gms872kuzdmg Use the newer_than parameter to filter articles that are newer than your last call. + Use the older_than parameter to filter articles that are older than your last call. + Use the max_batches parameter to limit the number of batches/pages to retrieve. :param session: The authenticated session to use to make the api calls :param newer_than: Only articles newer than this parameter will be used. If None only one call will be make, and the continuation will be ignored + :param older_than: Only articles older than this parameter will be used. If None only one call will be make, + and the continuation will be ignored + :param max_batches: Maximum number of batches to retrieve. If None, will continue until no more data is available """ self.newer_than = newer_than + self.older_than = older_than self.session = session self.stream_id = stream_id + self.max_batches = max_batches def download_all(self) -> List[T]: return self._merge(self.stream_bundles()) def stream_bundles(self) -> Iterable[T]: continuation = None + batch_count = 0 while True: - resp = self.session.make_api_request( - f"{self.RELATIVE_URL}", - params={ - "newerThan": int(self.newer_than.timestamp()) if self.newer_than else None, - "continuation": continuation, - "streamId": self.stream_id, - "format": self.FORMAT, - }, - ) + params = { + "continuation": continuation, + "streamId": self.stream_id, + "format": self.FORMAT, + } + if self.newer_than: + params["newerThan"] = int(self.newer_than.timestamp()) + if self.older_than: + params["olderThan"] = int(self.older_than.timestamp()) + + resp = self.session.make_api_request(f"{self.RELATIVE_URL}", params=params) yield self._parse_response(resp) + batch_count += 1 + + # Check if we've reached max_batches limit + if self.max_batches and batch_count >= self.max_batches: + return + if not self.newer_than or "link" not in resp.headers: return next_url = resp.headers["link"][1:].split(">")[0] @@ -70,21 +93,35 @@ def _merge(self, resp_jsons: Iterable[T]) -> T: class IoCDownloaderBuilder: - def __init__(self, session: FeedlySession, format: IoCFormat, newer_than: Optional[datetime] = None): + def __init__( + self, + session: FeedlySession, + format: IoCFormat, + newer_than: Optional[datetime] = None, + older_than: Optional[datetime] = None, + max_batches: Optional[int] = None, + ): """ Use this class to export the contextualized IoCs from a stream. Enterprise/personals feeds/boards are supported (see dedicated methods below). The IoCs will be returned along with their context and relationships in a dictionary representing a valid STIX v2.1 Bundle object. https://docs.oasis-open.org/cti/stix/v2.1/os/stix-v2.1-os.html#_gms872kuzdmg Use the newer_than parameter to filter articles that are newer than your last call. + Use the older_than parameter to filter articles that are older than your last call. + Use the max_batches parameter to limit the number of batches/pages to retrieve. :param session: The authenticated session to use to make the api calls :param newer_than: Only articles newer than this parameter will be used. If None only one call will be make, and the continuation will be ignored + :param older_than: Only articles older than this parameter will be used. If None only one call will be make, + and the continuation will be ignored + :param max_batches: Maximum number of batches to retrieve. If None, will continue until no more data is available """ self.session = session self.format = format self.newer_than = newer_than + self.older_than = older_than + self.max_batches = max_batches self.session.api_host = "https://cloud.feedly.com" self.user = self.session.user @@ -113,7 +150,13 @@ def from_stream_id(self, stream_id: str) -> IoCDownloaderABC: IoCFormat.STIX: StixIoCDownloader, IoCFormat.CSV: CsvIoCDownloader, } - return format2class[self.format](session=self.session, newer_than=self.newer_than, stream_id=stream_id) + return format2class[self.format]( + session=self.session, + newer_than=self.newer_than, + older_than=self.older_than, + stream_id=stream_id, + max_batches=self.max_batches, + ) class StixIoCDownloader(IoCDownloaderABC[Dict]): diff --git a/setup.py b/setup.py index 1fa7f3f..bc25cfd 100644 --- a/setup.py +++ b/setup.py @@ -18,7 +18,7 @@ EMAIL = "ml@feedly.com" AUTHOR = "Feedly" REQUIRES_PYTHON = ">=3.6.0" -VERSION = "0.26" +VERSION = "0.27" # What packages are required for this module to be executed? with open("requirements.txt") as f: