feat: use server scan preferences when available

Author: agateau-gg

changelog.d/20230523_130204_aurelien.gateau_read_metadata.md

@@ -0,0 +1,3 @@
+### Added
+
+- Added `GGClient.read_metadata()` to read metadata from the server. The metadata is then used by further secret scan calls and is available in a new `GGClient.secret_scan_preferences` attribute.

pygitguardian/client.py

@@ -9,12 +9,7 @@
 import requests
 from requests import Response, Session, codes
 
-from .config import (
-    DEFAULT_API_VERSION,
-    DEFAULT_BASE_URI,
-    DEFAULT_TIMEOUT,
-    MULTI_DOCUMENT_LIMIT,
-)
+from .config import DEFAULT_API_VERSION, DEFAULT_BASE_URI, DEFAULT_TIMEOUT
 from .iac_models import (
     IaCScanParameters,
     IaCScanParametersSchema,
@@ -29,6 +24,8 @@
     MultiScanResult,
     QuotaResponse,
     ScanResult,
+    SecretScanPreferences,
+    ServerMetadata,
 )
 
 
@@ -121,6 +118,7 @@ class GGClient:
     timeout: Optional[float]
     user_agent: str
     extra_headers: Dict
+    secret_scan_preferences: SecretScanPreferences
 
     def __init__(
         self,
@@ -178,6 +176,7 @@ def __init__(
                 "Authorization": f"Token {api_key}",
             },
         )
+        self.secret_scan_preferences = SecretScanPreferences()
 
     def request(
         self,
@@ -308,6 +307,9 @@ def content_scan(
             doc_dict["filename"] = filename
 
         request_obj = Document.SCHEMA.load(doc_dict)
+        Document.SCHEMA.validate_size(
+            request_obj, self.secret_scan_preferences.maximum_document_size
+        )
 
         resp = self.post(
             endpoint="scan",
@@ -344,16 +346,22 @@ def multi_content_scan(
         :param ignore_known_secrets: indicates whether known secrets should be ignored
         :return: Detail or ScanResult response and status code
         """
-        if len(documents) > MULTI_DOCUMENT_LIMIT:
+        max_documents = self.secret_scan_preferences.maximum_documents_per_scan
+        if len(documents) > max_documents:
             raise ValueError(
-                f"too many documents submitted for scan (max={MULTI_DOCUMENT_LIMIT})"
+                f"too many documents submitted for scan (max={max_documents})"
             )
 
         if all(isinstance(doc, dict) for doc in documents):
             request_obj = Document.SCHEMA.load(documents, many=True)
         else:
             raise TypeError("each document must be a dict")
 
+        for document in request_obj:
+            Document.SCHEMA.validate_size(
+                document, self.secret_scan_preferences.maximum_document_size
+            )
+
         params = (
             {"ignore_known_secrets": ignore_known_secrets}
             if ignore_known_secrets
@@ -472,3 +480,23 @@ def iac_directory_scan(
             result.status_code = resp.status_code
 
         return result
+
+    def read_metadata(self) -> Optional[Detail]:
+        """
+        Fetch server preferences and store them in `self.secret_scan_preferences`.
+        These preferences are then used by all future secret scans.
+
+        Note that the call fails if the API key is not valid.
+
+        :return: a Detail instance in case of error, None otherwise
+        """
+        resp = self.get("metadata")
+
+        if not is_ok(resp):
+            result = load_detail(resp)
+            result.status_code = resp.status_code
+            return result
+        metadata = ServerMetadata.SCHEMA.load(resp.json())
+
+        self.secret_scan_preferences = metadata.secret_scan_preferences
+        return None

pygitguardian/models.py

@@ -1,7 +1,9 @@
+from dataclasses import dataclass, field
 from datetime import date, datetime
 from typing import Any, ClassVar, Dict, List, Optional, cast
 from uuid import UUID
 
+import marshmallow_dataclass
 from marshmallow import (
     EXCLUDE,
     Schema,
@@ -10,10 +12,9 @@
     post_load,
     pre_load,
     validate,
-    validates,
 )
 
-from .config import DOCUMENT_SIZE_THRESHOLD_BYTES
+from .config import DOCUMENT_SIZE_THRESHOLD_BYTES, MULTI_DOCUMENT_LIMIT
 
 
 class BaseSchema(Schema):
@@ -52,17 +53,18 @@ class DocumentSchema(BaseSchema):
     filename = fields.String(validate=validate.Length(max=256), allow_none=True)
     document = fields.String(required=True)
 
-    @validates("document")
-    def validate_document(self, document: str) -> None:
-        """
-        validate that document is smaller than scan limit
+    @staticmethod
+    def validate_size(document: Dict[str, Any], maximum_size: int) -> None:
+        """Raises a ValidationError if the content of the document is longer than
+        `maximum_size`.
+
+        This is not implemented as a Marshmallow validator because the maximum size can
+        vary.
         """
-        encoded = document.encode("utf-8", errors="replace")
-        if len(encoded) > DOCUMENT_SIZE_THRESHOLD_BYTES:
+        encoded = document["document"].encode("utf-8", errors="replace")
+        if len(encoded) > maximum_size:
             raise ValidationError(
-                "file exceeds the maximum allowed size of {}B".format(
-                    DOCUMENT_SIZE_THRESHOLD_BYTES
-                )
+                f"file exceeds the maximum allowed size of {maximum_size}B"
             )
 
     @post_load
@@ -620,3 +622,23 @@ def __repr__(self) -> str:
                 self.secrets_engine_version or "",
             )
         )
+
+
+@dataclass
+class SecretScanPreferences:
+    maximum_document_size: int = DOCUMENT_SIZE_THRESHOLD_BYTES
+    maximum_documents_per_scan: int = MULTI_DOCUMENT_LIMIT
+
+
+@dataclass
+class ServerMetadata(Base):
+    version: str
+    preferences: Dict[str, Any]
+    secret_scan_preferences: SecretScanPreferences = field(
+        default_factory=SecretScanPreferences
+    )
+
+
+ServerMetadata.SCHEMA = marshmallow_dataclass.class_schema(
+    ServerMetadata, base_schema=BaseSchema
+)()