feat(pygitguardian): add multiscan

Author: João Guerreiro

.github/workflows/test-lint.yml

@@ -30,7 +30,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: [3.5, 3.6, 3.7, 3.8]
+        python-version: [3.6, 3.7, 3.8]
 
     steps:
       - uses: actions/checkout@v2
@@ -42,8 +42,7 @@ jobs:
         run: |
           python -m pip install --upgrade pip
           python -m pip install --upgrade pipenv
-          pipenv install --system --skip-lock
-          pip install nose coverage
+          pipenv install --system --dev --skip-lock
       - name: Test with pytest
         run: |
           pipenv run coverage run --source pygitguardian -m nose tests

Pipfile

@@ -4,16 +4,21 @@ verify_ssl = true
 name = "pypi"
 
 [packages]
-pygitguardian = {editable = true,path = "."}
 marshmallow = "~=3.5"
+pygitguardian = {editable = true,path = "."}
 requests = ">=2.21.0"
 
 [dev-packages]
 black = "==19.10b0"
-flake8 = "*"
-ipython = "*"
-pre-commit = "*"
 coverage = "*"
+flake8 = "*"
 flake8-isort = "*"
+ipython = "*"
 isort = "*"
 nose = "*"
+pre-commit = "*"
+vcrpy = "*"
+vcrpy-unittest = "*"
+
+[pipenv]
+allow_prereleases = true

examples/content_scan.py

@@ -2,6 +2,8 @@
 import sys
 import traceback
 
+from requests import codes
+
 from pygitguardian import GGClient
 
 
@@ -17,9 +19,9 @@
 client = GGClient(token=API_KEY)
 
 # Check the health of the API and the token used.
-health_obj = client.health_check()
+health_obj, status = client.health_check()
 
-if health_obj.success:
+if status == codes[r"\o/"]:  # this is 200 but cooler
     try:
         scan_result = client.content_scan(filename=FILENAME, document=DOCUMENT)
     except Exception as exc:

pygitguardian/__init__.py

@@ -4,7 +4,8 @@
 from .schemas import DetailSchema, DocumentSchema, ScanResultSchema
 
 
-__version__ = "1.0.0"
+__version__ = "1.0.1"
+GGClient._version = __version__
 
 __all__ = [
     "Detail",

pygitguardian/client.py

@@ -1,5 +1,6 @@
+import platform
 import urllib.parse
-from typing import Any, Optional, Tuple, Union
+from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
 
 import requests
 from marshmallow import Schema
@@ -18,26 +19,22 @@ class GGClient:
     DETAIL_SCHEMA = DetailSchema()
     DOCUMENT_SCHEMA = DocumentSchema()
     SCAN_RESULT_SCHEMA = ScanResultSchema()
+    _version = "undefined"
 
     def __init__(
         self,
         token: str,
-        base_uri: str = None,
-        session: requests.Session = None,
-        user_agent: str = "",
-        timeout: float = _DEFAULT_TIMEOUT,
-    ):
+        base_uri: Optional[str] = None,
+        session: Optional[requests.Session] = None,
+        user_agent: Optional[str] = None,
+        timeout: Optional[float] = _DEFAULT_TIMEOUT,
+    ) -> "GGClient":
         """
         :param token: APIKey to be added to requests
-        :type token: str
         :param base_uri: Base URI for the API, defaults to "https://api.gitguardian.com"
-        :type base_uri: str, optional
         :param session: custom requests session, defaults to requests.Session()
-        :type session: requests.Session, optional
         :param user_agent: user agent to identify requests, defaults to ""
-        :type user_agent: str, optional
         :param timeout: request timeout, defaults to 20s
-        :type timeout: float, optional
 
         :raises ValueError: if the protocol is invalid
         """
@@ -57,12 +54,15 @@ def __init__(
             session if session is isinstance(session, Session) else requests.Session()
         )
         self.timeout = timeout
+        self.user_agent = "pygitguardian/{0} ({1};py{2})".format(
+            self._version, platform.system(), platform.python_version()
+        )
+
+        if user_agent:
+            self.user_agent = " ".join([self.user_agent, user_agent])
 
         self.session.headers.update(
-            {
-                "User-Agent": " ".join(["pygitguardian", user_agent]),
-                "Authorization": "Token {0}".format(token),
-            }
+            {"User-Agent": self.user_agent, "Authorization": "Token {0}".format(token)}
         )
 
     def request(
@@ -71,6 +71,7 @@ def request(
         endpoint: str,
         schema: Schema = None,
         version: str = _API_VERSION,
+        many: bool = False,
         **kwargs
     ) -> Tuple[Any, Response]:
         if version:
@@ -86,11 +87,15 @@ def request(
             raise TypeError("Response is not JSON")
 
         if response.status_code == codes.ok and schema:
-            obj = schema.load(response.json())
+            obj = schema.load(response.json(), many=many)
+            if many:
+                for element in obj:
+                    element.status_code = response.status_code
+            else:
+                obj.status_code = response.status_code
         else:
             obj = self.DETAIL_SCHEMA.load(response.json())
-
-        obj.status_code = response.status_code
+            obj.status_code = response.status_code
 
         return obj, response
 
@@ -100,6 +105,7 @@ def post(
         data: str = None,
         schema: Schema = None,
         version: str = _API_VERSION,
+        many: bool = False,
         **kwargs
     ) -> Tuple[Any, Response]:
         return self.request(
@@ -108,6 +114,7 @@ def post(
             schema=schema,
             json=data,
             version=version,
+            many=many,
             **kwargs,
         )
 
@@ -116,6 +123,7 @@ def get(
         endpoint: str,
         schema: Schema = None,
         version: str = _API_VERSION,
+        many: bool = False,
         **kwargs
     ) -> Tuple[Any, Response]:
         return self.request(
@@ -124,37 +132,58 @@ def get(
 
     def content_scan(
         self, document: str, filename: Optional[str] = None
-    ) -> Union[Detail, ScanResult]:
-        """content_scan handles the /scan endpoint of the API
-
-        use filename=dummy to avoid evalutation of filename and file extension policies
+    ) -> Tuple[Union[Detail, ScanResult], int]:
+        """
+        content_scan handles the /scan endpoint of the API
 
         :param filename: name of file, example: "intro.py"
-        :type filename: str
         :param document: content of file
-        :type document: str
-        :return: Detail or ScanResult response
-        :rtype: Union[Detail, ScanResult]
+        :return: Detail or ScanResult response and status code
         """
 
         doc_dict = {"document": document}
         if filename:
             doc_dict["filename"] = filename
 
         request_obj = self.DOCUMENT_SCHEMA.load(doc_dict)
-        obj, _ = self.post(
+        obj, resp = self.post(
             endpoint="scan", data=request_obj, schema=self.SCAN_RESULT_SCHEMA
         )
-        return obj
+        return obj, resp.status_code
+
+    def multi_content_scan(
+        self, documents: Iterable[Dict[str, str]],
+    ) -> Tuple[Union[Detail, List[ScanResult]], int]:
+        """
+        multi_content_scan handles the /multiscan endpoint of the API
+
+        :param documents: List of dictionaries containing the keys document
+        and, optionaly, filename.
+            example: [{"document":"example content","filename":"intro.py"}]
+        :return: Detail or ScanResult response and status code
+        """
+
+        if all(isinstance(doc, dict) for doc in documents):
+            request_obj = self.DOCUMENT_SCHEMA.load(documents, many=True)
+        else:
+            raise TypeError("documents must be a dict")
+
+        obj, resp = self.post(
+            endpoint="multiscan",
+            data=request_obj,
+            schema=self.SCAN_RESULT_SCHEMA,
+            many=True,
+        )
+        return obj, resp.status_code
 
-    def health_check(self) -> Detail:
-        """health_check handles the /health endpoint of the API
+    def health_check(self) -> Tuple[Detail, int]:
+        """
+        health_check handles the /health endpoint of the API
 
         use Detail.status_code to check the response status code of the API
 
         200 if server is online and token is valid
-        :return: Detail response,
-        :rtype: Detail
+        :return: Detail response and status code
         """
-        obj, _ = self.get(endpoint="health")
-        return obj
+        obj, resp = self.get(endpoint="health")
+        return obj, resp.status_code

pygitguardian/models.py

@@ -1,23 +1,7 @@
 from typing import List, Optional
 
 
-class BaseObject:
-    UNSET_STATUS_CODE = 600
-
-    def __init__(self):
-        self.status_code = 600
-
-    @property
-    def success(self):
-        """success returns True if call returned 200
-
-        :return: call status
-        :rtype: bool
-        """
-        return self.status_code == 200
-
-
-class Detail(BaseObject):
+class Detail:
     """Detail is a response object mostly returned on error or when the
     api output is a simple string
 
@@ -87,7 +71,7 @@ def __repr__(self):
         )
 
 
-class ScanResult(BaseObject):
+class ScanResult:
     """ScanResult is a response object returned on a Content Scan
 
     Attributes:

pygitguardian/schemas.py

@@ -3,14 +3,44 @@
 This module contains marshmallow schemas responsible for
 serializing/deserializing request and response objects
 """
-from marshmallow import Schema, fields, post_load, validate
+from marshmallow import (
+    EXCLUDE,
+    Schema,
+    ValidationError,
+    fields,
+    post_load,
+    validate,
+    validates,
+)
 
 from .models import Detail, Match, PolicyBreak, ScanResult
 
 
+MB = 1048576
+
+
 class DocumentSchema(Schema):
-    filename = fields.Str(validate=validate.Length(max=256))
-    document = fields.Str(validate=validate.Length(max=1000000), required=True)
+    class Meta:
+        unknown = EXCLUDE
+
+    filename = fields.String(validate=validate.Length(max=256))
+    document = fields.String(required=True)
+
+    @validates("document")
+    def validate_document(self, document: str) -> str:
+        """
+        validate that document is smaller than scan limit
+        """
+        encoded = document.encode("utf-8")
+        if len(encoded) > MB:
+            raise ValidationError(
+                "This file exceeds the maximum allowed size of {}B".format(MB)
+            )
+
+        if "\x00" in document:
+            raise ValidationError("Document has null characters")
+
+        return document
 
 
 class MatchSchema(Schema):

setup.cfg

@@ -7,7 +7,7 @@ max-line-length = 120
 line_length=88
 lines_after_imports=2
 multi_line_output=3
-known_third_party=marshmallow,requests,setuptools
+known_third_party=marshmallow,requests,setuptools,vcr_unittest
 include_trailing_comma=true
 
 [metadata]