proxilion 0.0.1__py3-none-any.whl

proxilion/audit/exporters/gcp_storage.py

@@ -0,0 +1,570 @@
+"""
+Google Cloud Storage exporter for Proxilion audit logs.
+
+Supports exporting audit logs to:
+- Google Cloud Storage
+- BigQuery (for analytics)
+
+Uses google-cloud-storage if available, falls back to urllib with OAuth2.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import time
+import urllib.error
+import urllib.parse
+import urllib.request
+from datetime import datetime, timezone
+from typing import Any
+
+from proxilion.audit.events import AuditEventV2
+from proxilion.audit.exporters.cloud_base import (
+    BaseCloudExporter,
+    CloudExporterConfig,
+    ExportBatch,
+    ExportResult,
+)
+
+logger = logging.getLogger(__name__)
+
+# Check for google-cloud-storage availability
+try:
+    from google.cloud import storage as gcs
+    from google.oauth2 import service_account
+    HAS_GCS = True
+except ImportError:
+    HAS_GCS = False
+
+# Check for google-cloud-bigquery availability
+try:
+    from google.cloud import bigquery
+    HAS_BIGQUERY = True
+except ImportError:
+    HAS_BIGQUERY = False
+
+
+class GCSExporter(BaseCloudExporter):
+    """
+    Export audit logs to Google Cloud Storage.
+
+    Uses google-cloud-storage if installed, otherwise falls back to
+    urllib with Application Default Credentials or service account.
+
+    Example:
+        >>> config = CloudExporterConfig(
+        ...     provider="gcp",
+        ...     bucket_name="my-audit-logs",
+        ...     prefix="proxilion/prod/",
+        ... )
+        >>> exporter = GCSExporter(config)
+        >>> result = exporter.export(events)
+    """
+
+    def __init__(self, config: CloudExporterConfig) -> None:
+        """
+        Initialize the GCS exporter.
+
+        Args:
+            config: Exporter configuration.
+        """
+        super().__init__(config)
+        self._client = None
+        self._bucket = None
+        self._access_token: str | None = None
+        self._token_expiry: float = 0
+        self._initialize_client()
+
+    def _initialize_client(self) -> None:
+        """Initialize the GCS client."""
+        if HAS_GCS:
+            self._init_gcs_client()
+        else:
+            self._init_urllib_client()
+
+    def _init_gcs_client(self) -> None:
+        """Initialize google-cloud-storage client."""
+        if self.config.credentials_path:
+            credentials = service_account.Credentials.from_service_account_file(
+                self.config.credentials_path,
+                scopes=["https://www.googleapis.com/auth/devstorage.read_write"],
+            )
+            self._client = gcs.Client(credentials=credentials)
+        else:
+            # Use Application Default Credentials
+            self._client = gcs.Client()
+
+        self._bucket = self._client.bucket(self.config.bucket_name)
+
+    def _init_urllib_client(self) -> None:
+        """Initialize urllib-based client."""
+        # For urllib fallback, we need credentials
+        if self.config.credentials_path:
+            self._load_service_account()
+
+    def _load_service_account(self) -> None:
+        """Load service account credentials from file."""
+        try:
+            with open(self.config.credentials_path) as f:
+                self._service_account = json.load(f)
+        except Exception as e:
+            logger.warning(f"Failed to load service account: {e}")
+            self._service_account = None
+
+    def _get_access_token(self) -> str:
+        """
+        Get OAuth2 access token.
+
+        Returns:
+            Access token string.
+        """
+        # Check if we have a cached valid token
+        if self._access_token and time.time() < self._token_expiry:
+            return self._access_token
+
+        if not hasattr(self, "_service_account") or not self._service_account:
+            raise ValueError(
+                "GCP credentials not configured. Set credentials_path or use "
+                "Application Default Credentials with google-cloud-storage."
+            )
+
+        # Create JWT for service account
+        from base64 import urlsafe_b64encode
+
+        now = int(time.time())
+        expiry = now + 3600  # 1 hour
+
+        header = {"alg": "RS256", "typ": "JWT"}
+        payload = {
+            "iss": self._service_account["client_email"],
+            "scope": "https://www.googleapis.com/auth/devstorage.read_write",
+            "aud": "https://oauth2.googleapis.com/token",
+            "iat": now,
+            "exp": expiry,
+        }
+
+        # Encode header and payload
+        header_b64 = urlsafe_b64encode(json.dumps(header).encode()).rstrip(b"=")
+        payload_b64 = urlsafe_b64encode(json.dumps(payload).encode()).rstrip(b"=")
+        signing_input = header_b64 + b"." + payload_b64
+
+        # Sign with RSA-SHA256
+        try:
+            from cryptography.hazmat.primitives import hashes, serialization
+            from cryptography.hazmat.primitives.asymmetric import padding
+
+            private_key = serialization.load_pem_private_key(
+                self._service_account["private_key"].encode(),
+                password=None,
+            )
+            signature = private_key.sign(
+                signing_input,
+                padding.PKCS1v15(),
+                hashes.SHA256(),
+            )
+            signature_b64 = urlsafe_b64encode(signature).rstrip(b"=")
+        except ImportError as e:
+            raise ImportError(
+                "cryptography package required for service account auth. "
+                "Install with: pip install cryptography"
+            ) from e
+
+        jwt = signing_input + b"." + signature_b64
+
+        # Exchange JWT for access token
+        token_url = "https://oauth2.googleapis.com/token"
+        data = urllib.parse.urlencode({
+            "grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer",
+            "assertion": jwt.decode(),
+        }).encode()
+
+        request = urllib.request.Request(token_url, data=data)
+        request.add_header("Content-Type", "application/x-www-form-urlencoded")
+
+        with urllib.request.urlopen(request, timeout=30) as response:
+            token_data = json.loads(response.read())
+            self._access_token = token_data["access_token"]
+            self._token_expiry = time.time() + token_data.get("expires_in", 3600) - 60
+
+        return self._access_token
+
+    def export_batch(self, batch: ExportBatch) -> ExportResult:
+        """
+        Export a batch to GCS.
+
+        Args:
+            batch: The batch to export.
+
+        Returns:
+            ExportResult with success/failure information.
+        """
+        start_time = time.time()
+
+        try:
+            # Prepare data
+            data = batch.to_bytes(self.config.compression)
+            key = self.generate_key(batch.created_at, batch.batch_id)
+            checksum = self.compute_checksum(data)
+
+            # Upload with retry
+            self.with_retry(self._upload_object, key, data)
+
+            duration_ms = (time.time() - start_time) * 1000
+
+            logger.info(
+                f"Exported {batch.event_count} events to gs://{self.config.bucket_name}/{key}"
+            )
+
+            return ExportResult(
+                success=True,
+                events_exported=batch.event_count,
+                batch_id=batch.batch_id,
+                destination=f"gs://{self.config.bucket_name}/{key}",
+                duration_ms=duration_ms,
+                bytes_written=len(data),
+                checksum=checksum,
+            )
+
+        except Exception as e:
+            duration_ms = (time.time() - start_time) * 1000
+            logger.error(f"Failed to export batch {batch.batch_id}: {e}")
+
+            return ExportResult(
+                success=False,
+                events_exported=0,
+                batch_id=batch.batch_id,
+                error=str(e),
+                duration_ms=duration_ms,
+            )
+
+    def _upload_object(self, key: str, data: bytes) -> None:
+        """
+        Upload an object to GCS.
+
+        Args:
+            key: Object key.
+            data: Object data.
+        """
+        if HAS_GCS:
+            self._upload_gcs(key, data)
+        else:
+            self._upload_urllib(key, data)
+
+    def _upload_gcs(self, key: str, data: bytes) -> None:
+        """Upload using google-cloud-storage."""
+        blob = self._bucket.blob(key)
+        blob.upload_from_string(
+            data,
+            content_type=self.get_content_type(),
+        )
+
+    def _upload_urllib(self, key: str, data: bytes) -> None:
+        """Upload using urllib with OAuth2."""
+        access_token = self._get_access_token()
+
+        # Build upload URL
+        url = (
+            f"https://storage.googleapis.com/upload/storage/v1/b/"
+            f"{urllib.parse.quote(self.config.bucket_name)}/o"
+            f"?uploadType=media&name={urllib.parse.quote(key)}"
+        )
+
+        headers = {
+            "Authorization": f"Bearer {access_token}",
+            "Content-Type": self.get_content_type(),
+            "Content-Length": str(len(data)),
+        }
+
+        request = urllib.request.Request(url, data=data, headers=headers, method="POST")
+
+        try:
+            with urllib.request.urlopen(
+                request, timeout=self.config.read_timeout
+            ) as response:
+                if response.status not in (200, 201):
+                    raise ValueError(f"GCS upload failed with status {response.status}")
+        except urllib.error.HTTPError as e:
+            raise ValueError(f"GCS upload failed: {e.code} {e.reason}") from e
+
+    def health_check(self) -> bool:
+        """
+        Check if we can connect to GCS.
+
+        Returns:
+            True if healthy.
+        """
+        try:
+            if HAS_GCS:
+                self._bucket.reload()
+            else:
+                # Try to get bucket metadata
+                access_token = self._get_access_token()
+                url = (
+                    f"https://storage.googleapis.com/storage/v1/b/"
+                    f"{urllib.parse.quote(self.config.bucket_name)}"
+                )
+                headers = {"Authorization": f"Bearer {access_token}"}
+                request = urllib.request.Request(url, headers=headers)
+
+                with urllib.request.urlopen(request, timeout=10) as response:
+                    return response.status == 200
+
+            return True
+        except Exception as e:
+            logger.warning(f"GCS health check failed: {e}")
+            return False
+
+    def list_exports(
+        self,
+        start_date: datetime | None = None,
+        end_date: datetime | None = None,
+        max_results: int = 1000,
+    ) -> list[str]:
+        """
+        List exported files in the bucket.
+
+        Args:
+            start_date: Filter to exports after this date.
+            end_date: Filter to exports before this date.
+            max_results: Maximum number of results to return.
+
+        Returns:
+            List of object names.
+        """
+        if not HAS_GCS:
+            raise NotImplementedError("list_exports requires google-cloud-storage")
+
+        prefix = self.config.prefix
+        if start_date:
+            prefix += f"{start_date.year:04d}/"
+
+        blobs = self._bucket.list_blobs(prefix=prefix, max_results=max_results)
+
+        names = []
+        for blob in blobs:
+            # Filter by date if needed
+            if end_date:
+                parts = blob.name.split("/")
+                if len(parts) >= 4:
+                    try:
+                        year = int(parts[-4])
+                        month = int(parts[-3])
+                        day = int(parts[-2])
+                        blob_date = datetime(year, month, day, tzinfo=timezone.utc)
+                        if blob_date > end_date:
+                            continue
+                    except (ValueError, IndexError):
+                        pass
+
+            names.append(blob.name)
+
+            if len(names) >= max_results:
+                break
+
+        return names
+
+
+class BigQueryExporter(BaseCloudExporter):
+    """
+    Export audit logs directly to BigQuery.
+
+    Provides streaming insert for real-time analytics.
+    Requires google-cloud-bigquery package.
+
+    Example:
+        >>> config = CloudExporterConfig(
+        ...     provider="gcp",
+        ...     bucket_name="my-project.my_dataset.audit_logs",
+        ... )
+        >>> exporter = BigQueryExporter(config)
+        >>> result = exporter.export(events)
+    """
+
+    def __init__(
+        self,
+        config: CloudExporterConfig,
+        project_id: str | None = None,
+        dataset_id: str | None = None,
+        table_id: str | None = None,
+        create_table: bool = True,
+    ) -> None:
+        """
+        Initialize the BigQuery exporter.
+
+        Args:
+            config: Exporter configuration.
+            project_id: GCP project ID.
+            dataset_id: BigQuery dataset ID.
+            table_id: BigQuery table ID.
+            create_table: Create table if it doesn't exist.
+        """
+        super().__init__(config)
+
+        if not HAS_BIGQUERY:
+            raise ImportError(
+                "google-cloud-bigquery required for BigQueryExporter. "
+                "Install with: pip install google-cloud-bigquery"
+            )
+
+        # Parse table reference from bucket_name if in format project.dataset.table
+        if "." in config.bucket_name:
+            parts = config.bucket_name.split(".")
+            self.project_id = parts[0] if len(parts) > 0 else project_id
+            self.dataset_id = parts[1] if len(parts) > 1 else dataset_id
+            self.table_id = parts[2] if len(parts) > 2 else table_id
+        else:
+            self.project_id = project_id
+            self.dataset_id = dataset_id
+            self.table_id = table_id or config.bucket_name
+
+        self.create_table = create_table
+        self._client = None
+        self._table = None
+        self._initialize_client()
+
+    def _initialize_client(self) -> None:
+        """Initialize BigQuery client."""
+        if self.config.credentials_path:
+            from google.oauth2 import service_account
+            credentials = service_account.Credentials.from_service_account_file(
+                self.config.credentials_path,
+            )
+            self._client = bigquery.Client(
+                credentials=credentials,
+                project=self.project_id,
+            )
+        else:
+            self._client = bigquery.Client(project=self.project_id)
+
+        # Get or create table
+        table_ref = f"{self.project_id}.{self.dataset_id}.{self.table_id}"
+
+        try:
+            self._table = self._client.get_table(table_ref)
+        except Exception:
+            if self.create_table:
+                self._create_table(table_ref)
+            else:
+                raise
+
+    def _create_table(self, table_ref: str) -> None:
+        """Create the BigQuery table."""
+        schema = [
+            bigquery.SchemaField("event_id", "STRING", mode="REQUIRED"),
+            bigquery.SchemaField("timestamp", "TIMESTAMP", mode="REQUIRED"),
+            bigquery.SchemaField("sequence_number", "INTEGER"),
+            bigquery.SchemaField("event_type", "STRING"),
+            bigquery.SchemaField("user_id", "STRING"),
+            bigquery.SchemaField("user_roles", "STRING", mode="REPEATED"),
+            bigquery.SchemaField("session_id", "STRING"),
+            bigquery.SchemaField("agent_id", "STRING"),
+            bigquery.SchemaField("tool_name", "STRING"),
+            bigquery.SchemaField("tool_arguments", "JSON"),
+            bigquery.SchemaField("authorization_allowed", "BOOLEAN"),
+            bigquery.SchemaField("authorization_reason", "STRING"),
+            bigquery.SchemaField("policies_evaluated", "STRING", mode="REPEATED"),
+            bigquery.SchemaField("event_hash", "STRING"),
+            bigquery.SchemaField("previous_hash", "STRING"),
+        ]
+
+        table = bigquery.Table(table_ref, schema=schema)
+
+        # Enable partitioning by timestamp
+        table.time_partitioning = bigquery.TimePartitioning(
+            type_=bigquery.TimePartitioningType.DAY,
+            field="timestamp",
+        )
+
+        self._table = self._client.create_table(table)
+        logger.info(f"Created BigQuery table: {table_ref}")
+
+    def _event_to_row(self, event: AuditEventV2) -> dict[str, Any]:
+        """Convert an audit event to a BigQuery row."""
+        data = event.data
+        return {
+            "event_id": event.event_id,
+            "timestamp": event.timestamp.isoformat(),
+            "sequence_number": event.sequence_number,
+            "event_type": data.event_type.value,
+            "user_id": data.user_id,
+            "user_roles": data.user_roles or [],
+            "session_id": data.session_id,
+            "agent_id": data.agent_id,
+            "tool_name": data.tool_name,
+            "tool_arguments": json.dumps(data.tool_arguments) if data.tool_arguments else None,
+            "authorization_allowed": data.authorization_allowed,
+            "authorization_reason": data.authorization_reason,
+            "policies_evaluated": data.policies_evaluated or [],
+            "event_hash": event.event_hash,
+            "previous_hash": event.previous_hash,
+        }
+
+    def export_batch(self, batch: ExportBatch) -> ExportResult:
+        """
+        Export a batch to BigQuery.
+
+        Args:
+            batch: The batch to export.
+
+        Returns:
+            ExportResult with success/failure information.
+        """
+        start_time = time.time()
+
+        try:
+            # Convert events to rows
+            rows = [self._event_to_row(event) for event in batch.events]
+
+            # Insert rows
+            errors = self._client.insert_rows_json(self._table, rows)
+
+            duration_ms = (time.time() - start_time) * 1000
+
+            if errors:
+                error_msg = f"BigQuery insert errors: {errors}"
+                logger.error(error_msg)
+                return ExportResult(
+                    success=False,
+                    events_exported=0,
+                    batch_id=batch.batch_id,
+                    error=error_msg,
+                    duration_ms=duration_ms,
+                )
+
+            table_ref = f"{self.project_id}.{self.dataset_id}.{self.table_id}"
+            logger.info(f"Exported {batch.event_count} events to {table_ref}")
+
+            return ExportResult(
+                success=True,
+                events_exported=batch.event_count,
+                batch_id=batch.batch_id,
+                destination=table_ref,
+                duration_ms=duration_ms,
+            )
+
+        except Exception as e:
+            duration_ms = (time.time() - start_time) * 1000
+            logger.error(f"Failed to export batch {batch.batch_id}: {e}")
+
+            return ExportResult(
+                success=False,
+                events_exported=0,
+                batch_id=batch.batch_id,
+                error=str(e),
+                duration_ms=duration_ms,
+            )
+
+    def health_check(self) -> bool:
+        """
+        Check if we can connect to BigQuery.
+
+        Returns:
+            True if healthy.
+        """
+        try:
+            self._client.get_table(self._table)
+            return True
+        except Exception as e:
+            logger.warning(f"BigQuery health check failed: {e}")
+            return False