Adding GitHub Audit Log Extractor (#16)

* Adding Audit Log Extractor

* updating events -> actions

* formatting changes

* formatting changes

* updating tests

* adding cov to makefile and splitting url to 2 lines

* adding support for lookback_period and changing timestamp key

* updating code to resolve tests

* updating formatting

* format

* more formatting

* updating pipeline

* making changes based on comments

* keeping everything consistent

---------

Co-authored-by: rreddy15 <[email protected]>

Author: rreddy2

Makefile

@@ -28,3 +28,6 @@ lint: fmt
 test:
 	poetry run pytest
 
+.PHONY: coverage
+coverage:
+	poetry run coverage html

nodestream_github/__init__.py

@@ -1,3 +1,4 @@
+from .audit import GithubAuditLogExtractor
 from .interpretations import (
     RepositoryRelationshipInterpretation,
     UserRelationshipInterpretation,
@@ -9,6 +10,7 @@
 from .users import GithubUserExtractor
 
 __all__ = (
+    "GithubAuditLogExtractor",
     "GithubOrganizationsExtractor",
     "GithubPlugin",
     "GithubReposExtractor",

nodestream_github/audit.py

@@ -0,0 +1,48 @@
+"""
+Nodestream Extractor that extracts audit logs from the GitHub REST API.
+
+Developed using Enterprise Server 3.12
+https://docs.github.com/en/[email protected]/rest?apiVersion=2022-11-28
+"""
+
+from collections.abc import AsyncGenerator
+from typing import Any
+
+from nodestream.pipeline import Extractor
+
+from .client import GithubRestApiClient
+from .logging import get_plugin_logger
+from .types import GithubAuditLog
+
+logger = get_plugin_logger(__name__)
+
+
+class GithubAuditLogExtractor(Extractor):
+    """
+    Extracts audit logs from the GitHub REST API.
+    You can pass the enterprise_name, actions and lookback_period to the extractor
+    along with the regular GitHub parameters.
+
+    lookback_period can contain keys for days, months, and/or years as ints
+    actions can be found in the GitHub documentation
+    https://docs.github.com/en/[email protected]/admin/monitoring-activity-in-your-enterprise/reviewing-audit-logs-for-your-enterprise/searching-the-audit-log-for-your-enterprise#search-based-on-the-action-performed
+    """
+
+    def __init__(
+        self,
+        enterprise_name: str,
+        actions: list[str] | None = None,
+        lookback_period: dict[str, int] | None = None,
+        **github_client_kwargs: dict[str, Any] | None,
+    ):
+        self.enterprise_name = enterprise_name
+        self.client = GithubRestApiClient(**github_client_kwargs)
+        self.lookback_period = lookback_period
+        self.actions = actions
+
+    async def extract_records(self) -> AsyncGenerator[GithubAuditLog]:
+        async for audit in self.client.fetch_enterprise_audit_log(
+            self.enterprise_name, self.actions, self.lookback_period
+        ):
+            audit["timestamp"] = audit.pop("@timestamp")
+            yield audit

nodestream_github/client/githubclient.py

@@ -6,10 +6,12 @@
 import json
 import logging
 from collections.abc import AsyncGenerator
+from datetime import UTC, datetime
 from enum import Enum
 from typing import Any
 
 import httpx
+from dateutil.relativedelta import relativedelta
 from limits import RateLimitItem, RateLimitItemPerMinute
 from limits.aio.storage import MemoryStorage
 from limits.aio.strategies import MovingWindowRateLimiter, RateLimiter
@@ -328,6 +330,34 @@ async def fetch_all_organizations(self) -> AsyncGenerator[types.GithubOrg]:
         except httpx.HTTPError as e:
             _fetch_problem("all organizations", e)
 
+    async def fetch_enterprise_audit_log(
+        self, enterprise_name: str, actions: list[str], lookback_period: dict[str, int]
+    ) -> AsyncGenerator[types.GithubAuditLog]:
+        """Fetches enterprise-wide audit log data
+
+        https://docs.github.com/en/enterprise-cloud@latest/rest/enterprise-admin/audit-log?apiVersion=2022-11-28#get-the-audit-log-for-an-enterprise
+        """
+        try:
+            # adding action-based filtering
+            actions_phrase = " ".join(f"action:{action}" for action in actions)
+            # adding lookback_period based filtering
+            date_filter = (
+                f" created:>={(datetime.now(tz=UTC) - relativedelta(**lookback_period))
+                .strftime('%Y-%m-%d')}"
+                if lookback_period
+                else ""
+            )
+            search_phrase = f"{actions_phrase}{date_filter}"
+
+            params = {"phrase": search_phrase} if search_phrase else {}
+
+            async for audit in self._get_paginated(
+                f"enterprises/{enterprise_name}/audit-log", params=params
+            ):
+                yield audit
+        except httpx.HTTPError as e:
+            _fetch_problem("audit log", e)
+
     async def fetch_full_org(self, org_login: str) -> types.GithubOrg | None:
         """Fetches the complete org record.
 

nodestream_github/github_audit.yaml

@@ -0,0 +1,26 @@
+- implementation: nodestream_github:GithubAuditLogExtractor
+  arguments:
+    github_hostname: !config 'github_hostname'
+    auth_token: !config 'auth_token'
+    user_agent: !config 'user_agent'
+    enterprise_name: 'test-enterprise'
+    actions:
+      - protected_branch.create
+      - repo.download_zip
+    lookback_period:
+      days: 1
+
+- implementation: nodestream.interpreting:Interpreter
+  arguments:
+    interpretations:
+    - type: source_node
+      node_type: BranchProtectionPolicyChange
+      key:
+        timestamp: !jmespath 'timestamp'
+        actor: !jmespath 'actor'
+        action: !jmespath 'action'
+    - type: properties
+      properties:
+        org: !jmespath 'org'
+        repo: !jmespath 'repo'
+        created_at: !jmespath 'created_at'

nodestream_github/types/__init__.py

@@ -1,4 +1,5 @@
 from .github import (
+    GithubAuditLog,
     GithubOrg,
     GithubOrgSummary,
     GithubRepo,
@@ -23,6 +24,7 @@
     "GithubRepo",
     "GithubTeam",
     "GithubTeamSummary",
+    "GithubAuditLog",
     "GithubUser",
     "HeaderTypes",
     "JSONType",

nodestream_github/types/github.py

@@ -11,6 +11,7 @@
 Webhook: TypeAlias = JSONType
 GithubTeam: TypeAlias = JSONType
 GithubTeamSummary: TypeAlias = JSONType
+GithubAuditLog: TypeAlias = JSONType
 
 LanguageRecord: TypeAlias = JSONType
 OrgRecord: TypeAlias = JSONType

tests/data/audit.py

@@ -0,0 +1,62 @@
+GITHUB_AUDIT = [
+    {
+        "@timestamp": 1606929874512,
+        "action": "team.add_member",
+        "actor": "octocat",
+        "created_at": 1606929874512,
+        "_document_id": "xJJFlFOhQ6b-5vaAFy9Rjw",
+        "org": "octo-corp",
+        "team": "octo-corp/example-team",
+        "user": "monalisa",
+    },
+    {
+        "@timestamp": 1606507117008,
+        "action": "org.create",
+        "actor": "octocat",
+        "created_at": 1606507117008,
+        "_document_id": "Vqvg6kZ4MYqwWRKFDzlMoQ",
+        "org": "octocat-test-org",
+    },
+    {
+        "@timestamp": 1605719148837,
+        "action": "repo.destroy",
+        "actor": "monalisa",
+        "created_at": 1605719148837,
+        "_document_id": "LwW2vpJZCDS-WUmo9Z-ifw",
+        "org": "mona-org",
+        "repo": "mona-org/mona-test-repo",
+        "visibility": "private",
+    },
+]
+
+
+GITHUB_EXPECTED_OUTPUT = [
+    {
+        "timestamp": 1606929874512,
+        "action": "team.add_member",
+        "actor": "octocat",
+        "created_at": 1606929874512,
+        "_document_id": "xJJFlFOhQ6b-5vaAFy9Rjw",
+        "org": "octo-corp",
+        "team": "octo-corp/example-team",
+        "user": "monalisa",
+    },
+    {
+        "timestamp": 1606507117008,
+        "action": "org.create",
+        "actor": "octocat",
+        "created_at": 1606507117008,
+        "_document_id": "Vqvg6kZ4MYqwWRKFDzlMoQ",
+        "org": "octocat-test-org",
+    },
+    {
+        "timestamp": 1605719148837,
+        "action": "repo.destroy",
+        "actor": "monalisa",
+        "created_at": 1605719148837,
+        "_document_id": "LwW2vpJZCDS-WUmo9Z-ifw",
+        "org": "mona-org",
+        "repo": "mona-org/mona-test-repo",
+        "visibility": "private",
+    },
+]

tests/mocks/githubrest.py

@@ -189,3 +189,10 @@ def get_repos_for_user(
             url=f"{self.base_url}/users/{user_login}/repos?per_page=100&{type_param}",
             **kwargs,
         )
+
+    def get_enterprise_audit_logs(self, **kwargs: dict[str, Any]):
+        url = (
+            f"{self.base_url}/enterprises/test-enterprise"
+            f"/audit-log?per_page=100&phrase=action:protected_branch.create"
+        )
+        self.add_response(url=url, **kwargs)

tests/test_audit.py

@@ -0,0 +1,32 @@
+import pytest
+
+from nodestream_github import GithubAuditLogExtractor
+from tests.data.audit import GITHUB_AUDIT, GITHUB_EXPECTED_OUTPUT
+from tests.mocks.githubrest import (
+    DEFAULT_HOSTNAME,
+    DEFAULT_PER_PAGE,
+    GithubHttpxMock,
+)
+
+
+@pytest.fixture
+def audit_extractor() -> GithubAuditLogExtractor:
+    return GithubAuditLogExtractor(
+        auth_token="test-token",
+        github_hostname=DEFAULT_HOSTNAME,
+        user_agent="test-agent",
+        max_retries=0,
+        per_page=DEFAULT_PER_PAGE,
+        enterprise_name="test-enterprise",
+        actions=["protected_branch.create"],
+    )
+
+
+@pytest.mark.asyncio
+async def test_get_audit(
+    audit_extractor: GithubAuditLogExtractor, gh_rest_mock: GithubHttpxMock
+):
+    gh_rest_mock.get_enterprise_audit_logs(status_code=200, json=GITHUB_AUDIT)
+
+    all_records = [record async for record in audit_extractor.extract_records()]
+    assert all_records == GITHUB_EXPECTED_OUTPUT