Add multi tenancy support (#99)

Author: xvaier

packages/flare/bin/constants.py

@@ -11,17 +11,23 @@
 
 class PasswordKeys(Enum):
     API_KEY = "api_key"
-    TENANT_ID = "tenant_id"
+    TENANT_IDS = "tenant_ids"
     INGEST_FULL_EVENT_DATA = "ingest_full_event_data"
     SEVERITIES_FILTER = "severities_filter"
     SOURCE_TYPES_FILTER = "source_types_filter"
 
 
 class DataStoreKeys(Enum):
-    LAST_INGESTED_TENANT_ID = "last_ingested_tenant_id"
     START_DATE = "start_date"
     TIMESTAMP_LAST_FETCH = "timestamp_last_fetch"
 
+    SECTION_METADATA = "metadata"
+    SECTION_TENANT_DATA = "tenant_data"
+
     @staticmethod
     def get_next_token(tenant_id: int) -> str:
         return f"next_{tenant_id}"
+
+    @staticmethod
+    def get_earliest_ingested(tenant_id: int) -> str:
+        return f"timestamp_earliest_ingested_{tenant_id}"

packages/flare/bin/cron_job_ingest_events.py

@@ -43,7 +43,7 @@ def main(
         return
 
     api_key = get_api_key(storage_passwords=storage_passwords)
-    tenant_id = get_tenant_id(storage_passwords=storage_passwords)
+    tenant_ids = get_tenant_ids(storage_passwords=storage_passwords)
     ingest_full_event_data = get_ingest_full_event_data(
         storage_passwords=storage_passwords
     )
@@ -52,36 +52,71 @@ def main(
 
     data_store.set_last_fetch(datetime.now(timezone.utc))
 
-    # If the tenant has changed, update the start date so that future requests will be based off today
-    # If you switch tenants, this will avoid the old tenant from ingesting all the events before today and the day
-    # that tenant was switched in the first place.
-    if data_store.get_last_tenant_id() != tenant_id:
-        data_store.set_start_date((datetime.now(timezone.utc) - timedelta(days=30)))
+    total_events_fetched_count = 0
 
-    data_store.set_last_tenant_id(tenant_id)
+    for tenant_id in tenant_ids:
+        events_fetched_count = 0
 
-    events_fetched_count = 0
-    for event, next_token in fetch_feed(
-        logger=logger,
-        api_key=api_key,
-        tenant_id=tenant_id,
-        ingest_full_event_data=ingest_full_event_data,
-        severities=severities_filter,
-        source_types=source_types_filter,
-        flare_api_cls=flare_api_cls,
-        data_store=data_store,
-    ):
-        data_store.set_last_fetch(datetime.now(timezone.utc))
+        # The earliest ingested date serves as a low water mark to look
+        # for identifiers 30 days prior to the day a tenant was first configured.
+        start_date = data_store.get_earliest_ingested_by_tenant(tenant_id)
+        if not start_date:
+            start_date = datetime.now(timezone.utc) - timedelta(days=30)
+            data_store.set_earliest_ingested_by_tenant(tenant_id, start_date)
+
+        for event, next_token in fetch_feed(
+            logger=logger,
+            api_key=api_key,
+            tenant_id=tenant_id,
+            ingest_full_event_data=ingest_full_event_data,
+            severities=severities_filter,
+            source_types=source_types_filter,
+            flare_api_cls=flare_api_cls,
+            data_store=data_store,
+        ):
+            data_store.set_last_fetch(datetime.now(timezone.utc))
+
+            data_store.set_next_by_tenant(tenant_id, next_token)
+
+            event["tenant_id"] = tenant_id
 
-        data_store.set_next_by_tenant(tenant_id, next_token)
+            # stdout is picked up by splunk and this is how events
+            # are ingested after being retrieved from Flare.
+            print(json.dumps(event), flush=True)
 
-        # stdout is picked up by splunk and this is how events
-        # are ingested after being retrieved from Flare.
-        print(json.dumps(event), flush=True)
+            events_fetched_count += 1
+        logger.info(f"Fetched {events_fetched_count} events on tenant {tenant_id}")
+        total_events_fetched_count += events_fetched_count
 
-        events_fetched_count += 1
+    logger.info(f"Fetched {events_fetched_count} events across all tenants")
+
+
+def fetch_feed(
+    logger: Logger,
+    api_key: str,
+    tenant_id: int,
+    ingest_full_event_data: bool,
+    severities: list[str],
+    source_types: list[str],
+    flare_api_cls: type[FlareAPI],
+    data_store: ConfigDataStore,
+) -> Iterator[tuple[dict, str]]:
+    flare_api: FlareAPI = flare_api_cls(api_key=api_key, tenant_id=tenant_id)
 
-    logger.info(f"Fetched {events_fetched_count} events")
+    try:
+        next = data_store.get_next_by_tenant(tenant_id)
+        start_date = data_store.get_earliest_ingested_by_tenant(tenant_id)
+        logger.info(f"Fetching {tenant_id=}, {next=}, {start_date=}")
+        for event_next in flare_api.fetch_feed_events(
+            next=next,
+            start_date=start_date,
+            ingest_full_event_data=ingest_full_event_data,
+            severities=severities,
+            source_types=source_types,
+        ):
+            yield event_next
+    except Exception as e:
+        logger.error(f"Exception={e}")
 
 
 def get_storage_password_value(
@@ -103,18 +138,20 @@ def get_api_key(storage_passwords: client.StoragePasswords) -> str:
     return api_key
 
 
-def get_tenant_id(storage_passwords: client.StoragePasswords) -> int:
-    stored_tenant_id = get_storage_password_value(
-        storage_passwords=storage_passwords, password_key=PasswordKeys.TENANT_ID.value
+def get_tenant_ids(storage_passwords: client.StoragePasswords) -> list[int]:
+    stored_tenant_ids = get_storage_password_value(
+        storage_passwords=storage_passwords, password_key=PasswordKeys.TENANT_IDS.value
     )
     try:
-        tenant_id = int(stored_tenant_id) if stored_tenant_id is not None else None
+        tenant_ids: Optional[list[int]] = (
+            json.loads(stored_tenant_ids) if stored_tenant_ids else None
+        )
     except Exception:
         pass
 
-    if not tenant_id:
-        raise Exception("Tenant ID not found")
-    return tenant_id
+    if not tenant_ids:
+        raise Exception("Tenant IDs not found")
+    return tenant_ids
 
 
 def get_ingest_full_event_data(storage_passwords: client.StoragePasswords) -> bool:
@@ -151,34 +188,6 @@ def get_source_types_filter(storage_passwords: client.StoragePasswords) -> list[
     return []
 
 
-def fetch_feed(
-    logger: Logger,
-    api_key: str,
-    tenant_id: int,
-    ingest_full_event_data: bool,
-    severities: list[str],
-    source_types: list[str],
-    flare_api_cls: type[FlareAPI],
-    data_store: ConfigDataStore,
-) -> Iterator[tuple[dict, str]]:
-    try:
-        flare_api: FlareAPI = flare_api_cls(api_key=api_key, tenant_id=tenant_id)
-
-        next = data_store.get_next_by_tenant(tenant_id)
-        start_date = data_store.get_start_date()
-        logger.info(f"Fetching {tenant_id=}, {next=}, {start_date=}")
-        for event_next in flare_api.fetch_feed_events(
-            next=next,
-            start_date=start_date,
-            ingest_full_event_data=ingest_full_event_data,
-            severities=severities,
-            source_types=source_types,
-        ):
-            yield event_next
-    except Exception as e:
-        logger.error(f"Exception={e}")
-
-
 def get_splunk_service(logger: Logger, token: str) -> client.Service:
     try:
         splunk_service = client.connect(

packages/flare/bin/data_store.py

@@ -16,14 +16,14 @@
 
 class ConfigDataStore:
     def __init__(self) -> None:
-        config_store = configparser.ConfigParser()
+        config_store = configparser.RawConfigParser()
         config_store.read(config_path)
 
         # Add data sections
-        if "metadata" not in config_store.sections():
-            config_store.add_section("metadata")
-        if "next_tokens" not in config_store.sections():
-            config_store.add_section("next_tokens")
+        if DataStoreKeys.SECTION_METADATA.value not in config_store.sections():
+            config_store.add_section(DataStoreKeys.SECTION_METADATA.value)
+        if DataStoreKeys.SECTION_TENANT_DATA.value not in config_store.sections():
+            config_store.add_section(DataStoreKeys.SECTION_TENANT_DATA.value)
         self._store = config_store
 
     def _commit(self) -> None:
@@ -33,47 +33,12 @@ def _commit(self) -> None:
     def _sync(self) -> None:
         self._store.read(config_path)
 
-    def get_last_tenant_id(self) -> Optional[int]:
-        self._sync()
-        last_ingested_tenant_id = self._store.get(
-            "metadata", DataStoreKeys.LAST_INGESTED_TENANT_ID.value, fallback=None
-        )
-
-        try:
-            return int(last_ingested_tenant_id) if last_ingested_tenant_id else None
-        except Exception:
-            pass
-        return None
-
-    def set_last_tenant_id(self, tenant_id: int) -> None:
-        self._store.set(
-            "metadata", DataStoreKeys.LAST_INGESTED_TENANT_ID.value, str(tenant_id)
-        )
-        self._commit()
-
-    def get_start_date(self) -> Optional[datetime]:
-        self._sync()
-        start_date = self._store.get(
-            "metadata", DataStoreKeys.START_DATE.value, fallback=None
-        )
-
-        if start_date:
-            try:
-                return datetime.fromisoformat(start_date)
-            except Exception:
-                pass
-        return None
-
-    def set_start_date(self, start_date: datetime) -> None:
-        self._store.set(
-            "metadata", DataStoreKeys.START_DATE.value, start_date.isoformat()
-        )
-        self._commit()
-
     def get_last_fetch(self) -> Optional[datetime]:
         self._sync()
         last_fetched = self._store.get(
-            "metadata", DataStoreKeys.TIMESTAMP_LAST_FETCH.value, fallback=None
+            DataStoreKeys.SECTION_METADATA.value,
+            DataStoreKeys.TIMESTAMP_LAST_FETCH.value,
+            fallback=None,
         )
 
         if last_fetched:
@@ -85,7 +50,7 @@ def get_last_fetch(self) -> Optional[datetime]:
 
     def set_last_fetch(self, last_fetch: datetime) -> None:
         self._store.set(
-            "metadata",
+            DataStoreKeys.SECTION_METADATA.value,
             DataStoreKeys.TIMESTAMP_LAST_FETCH.value,
             last_fetch.isoformat(),
         )
@@ -94,7 +59,7 @@ def set_last_fetch(self, last_fetch: datetime) -> None:
     def get_next_by_tenant(self, tenant_id: int) -> Optional[str]:
         self._sync()
         return self._store.get(
-            "next_tokens",
+            DataStoreKeys.SECTION_TENANT_DATA.value,
             DataStoreKeys.get_next_token(tenant_id=tenant_id),
             fallback=None,
         )
@@ -104,5 +69,33 @@ def set_next_by_tenant(self, tenant_id: int, next: Optional[str]) -> None:
             return
 
         self._store.set(
-            "next_tokens", DataStoreKeys.get_next_token(tenant_id=tenant_id), next
+            DataStoreKeys.SECTION_TENANT_DATA.value,
+            DataStoreKeys.get_next_token(tenant_id=tenant_id),
+            next,
+        )
+        self._commit()
+
+    def get_earliest_ingested_by_tenant(self, tenant_id: int) -> Optional[datetime]:
+        self._sync()
+        earliest_ingested = self._store.get(
+            DataStoreKeys.SECTION_TENANT_DATA.value,
+            DataStoreKeys.get_earliest_ingested(tenant_id=tenant_id),
+            fallback=None,
         )
+
+        if earliest_ingested:
+            try:
+                return datetime.fromisoformat(earliest_ingested)
+            except Exception:
+                pass
+        return None
+
+    def set_earliest_ingested_by_tenant(
+        self, tenant_id: int, earliest_ingested: datetime
+    ) -> None:
+        self._store.set(
+            DataStoreKeys.SECTION_TENANT_DATA.value,
+            DataStoreKeys.get_earliest_ingested(tenant_id=tenant_id),
+            earliest_ingested.isoformat(),
+        )
+        self._commit()

packages/flare/bin/flare_external_requests.py

@@ -90,13 +90,11 @@ def handle_GET(self) -> None:
 
         data_store = ConfigDataStore()
         last_fetched_timestamp = data_store.get_last_fetch()
-        last_tenant_id = data_store.get_last_tenant_id()
 
         status_resp = {
             "last_fetched_at": last_fetched_timestamp.isoformat()
             if last_fetched_timestamp is not None
-            else None,
-            "last_tenant_id": last_tenant_id,
+            else None
         }
         logger.debug(f"FlareIngestionStatus: {status_resp}")
         self.response.setHeader("Content-Type", "application/json")

packages/flare/tests/bin/conftest.py

@@ -116,7 +116,7 @@ def mock_env(mock_config_file: Path) -> Generator[None, None, None]:
 @pytest.fixture
 def data_store(mock_env: None) -> Generator[ConfigDataStore, None, None]:
     # Creates an instance of ConfigDataStore with mocked dependencies.
-    with mock.patch("configparser.ConfigParser.read") as mock_read:
+    with mock.patch("configparser.RawConfigParser.read") as mock_read:
         mock_read.return_value = None
         store = ConfigDataStore()
         store._commit = lambda: None

packages/flare/tests/bin/test_data_store.py

@@ -3,17 +3,6 @@
 from datetime import timezone
 
 
-def test_get_and_set_last_tenant_id(data_store: ConfigDataStore) -> None:
-    data_store.set_last_tenant_id(123)
-    assert data_store.get_last_tenant_id() == 123
-
-
-def test_get_and_set_start_date(data_store: ConfigDataStore) -> None:
-    date = datetime(2024, 3, 6, 12, 0, 0, tzinfo=timezone.utc)
-    data_store.set_start_date(date)
-    assert data_store.get_start_date() == date
-
-
 def test_get_and_set_last_fetch(data_store: ConfigDataStore) -> None:
     date = datetime(2024, 3, 6, 14, 0, 0, tzinfo=timezone.utc)
     data_store.set_last_fetch(date)
@@ -25,3 +14,10 @@ def test_get_and_set_next_by_tenant(data_store: ConfigDataStore) -> None:
     next_token = "next_token_value"
     data_store.set_next_by_tenant(tenant_id, next_token)
     assert data_store.get_next_by_tenant(tenant_id) == "next_token_value"
+
+
+def test_get_and_set_earliest_ingested_by_tenant(data_store: ConfigDataStore) -> None:
+    tenant_id = 789
+    date = datetime(2024, 3, 6, 14, 0, 0, tzinfo=timezone.utc)
+    data_store.set_earliest_ingested_by_tenant(tenant_id, date)
+    assert data_store.get_earliest_ingested_by_tenant(tenant_id) == date