feat: added iac models and modified GGClient

Author: amascia-gg

Pipfile

@@ -7,6 +7,7 @@ verify_ssl = true
 marshmallow = ">=3.5"
 pygitguardian = { editable = true, path = "." }
 requests = ">=2"
+marshmallow-dataclass = ">=8.5.8,<8.6.0"
 
 [dev-packages]
 black = "==22.3.0"

pygitguardian/client.py

@@ -1,7 +1,13 @@
+import os
 import platform
+import tarfile
 import urllib.parse
+from io import BytesIO
+from pathlib import Path
 from typing import Any, Dict, List, Optional, Union, cast
 
+import click
+import requests
 from requests import Response, Session, codes
 
 from .config import (
@@ -10,6 +16,12 @@
     DEFAULT_TIMEOUT,
     MULTI_DOCUMENT_LIMIT,
 )
+from .iac_models import (
+    IaCScanParameters,
+    IaCScanParametersSchema,
+    IaCScanResult,
+    IaCScanResultSchema,
+)
 from .models import (
     Detail,
     Document,
@@ -20,6 +32,10 @@
 )
 
 
+# max files size to create a tar from
+MAX_TAR_CONTENT_SIZE = 30 * 1024 * 1024
+
+
 class Versions:
     app_version: Optional[str] = None
     secrets_engine_version: Optional[str] = None
@@ -57,6 +73,27 @@ def is_ok(resp: Response) -> bool:
     )
 
 
+def _create_tar(root_path: Path, filenames: List[str]) -> bytes:
+    """
+    :param root_path: the root_path from which the tar is created
+    :param files: the files which need to be added to the tar, filenames should be the paths relative to the root_path
+    :return: a bytes object containing the tar.gz created from the files, with paths relative to root_path
+    """
+    tar_stream = BytesIO()
+    current_dir_size = 0
+    with tarfile.open(fileobj=tar_stream, mode="w:gz") as tar:
+        for filename in filenames:
+            full_path = root_path / filename
+            current_dir_size += os.path.getsize(full_path)
+            if current_dir_size > MAX_TAR_CONTENT_SIZE:
+                raise click.ClickException(
+                    f"The total size of the files processed exceeds {MAX_TAR_CONTENT_SIZE / (1024 * 1024):.0f}MB, "
+                    f"please try again with less files"
+                )
+            tar.add(full_path, filename)
+    return tar_stream.getvalue()
+
+
 class GGClient:
     _version = "undefined"
     session: Session
@@ -197,7 +234,7 @@ def get(
     def post(
         self,
         endpoint: str,
-        data: Optional[str] = None,
+        data: Optional[Dict[str, str]] = None,
         version: str = DEFAULT_API_VERSION,
         extra_headers: Optional[Dict[str, str]] = None,
         **kwargs: Any,
@@ -347,3 +384,38 @@ def quota_overview(
         obj.status_code = resp.status_code
 
         return obj
+
+    # For IaC Scans
+    def iac_directory_scan(
+        self,
+        directory: Path,
+        filenames: List[str],
+        scan_parameters: IaCScanParameters,
+        extra_headers: Optional[Dict[str, str]] = None,
+    ) -> Union[Detail, IaCScanResult]:
+
+        tar = _create_tar(directory, filenames)
+        result: Union[Detail, IaCScanResult]
+        try:
+            resp = self.post(
+                endpoint="iac_scan",
+                extra_headers=extra_headers,
+                files={
+                    "directory": tar,
+                },
+                data={
+                    "scan_parameters": IaCScanParametersSchema().dumps(scan_parameters),
+                },
+            )
+        except requests.exceptions.ReadTimeout:
+            result = Detail("The request timed out.")
+            result.status_code = 504
+        else:
+            if is_ok(resp):
+                result = IaCScanResultSchema().load(resp.json())
+            else:
+                result = load_detail(resp)
+
+            result.status_code = resp.status_code
+
+        return result

pygitguardian/iac_models.py

@@ -0,0 +1,54 @@
+from dataclasses import dataclass, field
+from typing import List, Optional
+
+import marshmallow_dataclass
+
+from pygitguardian.models import Base, BaseSchema
+
+
+@dataclass
+class IaCVulnerability(Base):
+    policy: str
+    policy_id: str
+    line_end: int
+    line_start: int
+    description: str
+    documentation_url: str
+    component: str = ""
+    severity: str = ""
+
+
+IaCVulnerabilitySchema = marshmallow_dataclass.class_schema(
+    IaCVulnerability, BaseSchema
+)
+
+
+@dataclass
+class IaCFileResult(Base):
+    filename: str
+    incidents: List[IaCVulnerability]
+
+
+IaCFileResultSchema = marshmallow_dataclass.class_schema(IaCFileResult, BaseSchema)
+
+
+@dataclass
+class IaCScanParameters(Base):
+    ignored_policies: List[str] = field(default_factory=list)
+    minimum_severity: Optional[str] = None
+
+
+IaCScanParametersSchema = marshmallow_dataclass.class_schema(
+    IaCScanParameters, BaseSchema
+)
+
+
+@dataclass
+class IaCScanResult(Base):
+    id: str = ""
+    type: str = ""
+    iac_engine_version: str = ""
+    entities_with_incidents: List[IaCFileResult] = field(default_factory=list)
+
+
+IaCScanResultSchema = marshmallow_dataclass.class_schema(IaCScanResult, BaseSchema)

tests/test_iac_models.py

@@ -0,0 +1,100 @@
+import pytest
+
+from pygitguardian.iac_models import (
+    IaCFileResult,
+    IaCFileResultSchema,
+    IaCScanParameters,
+    IaCScanParametersSchema,
+    IaCScanResult,
+    IaCScanResultSchema,
+    IaCVulnerability,
+    IaCVulnerabilitySchema,
+)
+
+
+class TestModel:
+    @pytest.mark.parametrize(
+        "schema_klass, expected_klass, instance_data",
+        [
+            (
+                IaCScanResultSchema,
+                IaCScanResult,
+                {
+                    "id": "myid",
+                    "type": "type",
+                    "iac_engine_version": "version",
+                    "entities_with_incidents": [
+                        {
+                            "filename": "filename",
+                            "incidents": [
+                                {
+                                    "policy": "mypolicy,",
+                                    "policy_id": "mypolicyid",
+                                    "line_end": 0,
+                                    "line_start": 0,
+                                    "description": "mydescription",
+                                    "documentation_url": "mydoc",
+                                    "component": "mycomponent",
+                                    "severity": "myseverity",
+                                    "some_extra_field": "extra",
+                                }
+                            ],
+                        }
+                    ],
+                },
+            ),
+            (
+                IaCScanParametersSchema,
+                IaCScanParameters,
+                {"ignored_policies": ["pol1", "pol2"], "minimum_severity": "LOW"},
+            ),
+            (
+                IaCVulnerabilitySchema,
+                IaCVulnerability,
+                {
+                    "policy": "mypolicy,",
+                    "policy_id": "mypolicyid",
+                    "line_end": 0,
+                    "line_start": 0,
+                    "description": "mydescription",
+                    "documentation_url": "mydoc",
+                    "component": "mycomponent",
+                    "severity": "myseverity",
+                    "some_extra_field": "extra",
+                },
+            ),
+            (
+                IaCFileResultSchema,
+                IaCFileResult,
+                {
+                    "filename": "filename",
+                    "incidents": [
+                        {
+                            "policy": "mypolicy,",
+                            "policy_id": "mypolicyid",
+                            "line_end": 0,
+                            "line_start": 0,
+                            "description": "mydescription",
+                            "documentation_url": "mydoc",
+                            "component": "mycomponent",
+                            "severity": "myseverity",
+                            "some_extra_field": "extra",
+                        }
+                    ],
+                },
+            ),
+        ],
+    )
+    def test_schema_loads(self, schema_klass, expected_klass, instance_data):
+        """
+        GIVEN the right kwargs and an extra field in dict format
+        WHEN loading using the schema
+        THEN the extra field is not taken into account
+            AND the result should be an instance of the expected class
+        """
+        schema = schema_klass()
+
+        data = {**instance_data, "field": "extra"}
+
+        obj = schema.load(data)
+        assert isinstance(obj, expected_klass)

tests/test_utils.py

@@ -0,0 +1,50 @@
+import tarfile
+from io import BytesIO
+from unittest import mock
+
+import click
+import pytest
+
+from pygitguardian.client import _create_tar
+
+
+def test_create_tar(tmp_path):
+    """
+    GIVEN a list of filenames, representing paths relative to the tmp directory
+    WHEN the _create_tar method is called
+    THEN a bytes object is outputted, representing a tar of the files, with paths relative to the tmp directory
+    """
+    file1_name = "file1.txt"
+    dir_path = "my_test_dir"
+    file2_name = f"{dir_path}/file2.txt"
+    file1_content = "My first document"
+    file2_content = "My second document"
+
+    (tmp_path / file1_name).write_text(file1_content)
+    (tmp_path / dir_path).mkdir(parents=True, exist_ok=True)
+    (tmp_path / file2_name).write_text(file2_content)
+
+    tar_stream = _create_tar(tmp_path, [file1_name, file2_name])
+
+    # Create tar archive from bytes stream and write it in the tmp_path/output directory
+    (tmp_path / "output").mkdir(exist_ok=True)
+    with tarfile.open(fileobj=BytesIO(tar_stream), mode="r:gz") as tar:
+        tar.extractall(tmp_path / "output")
+
+    assert file1_content == (tmp_path / f"output/{file1_name}").read_text()
+    assert file2_content == (tmp_path / f"output/{file2_name}").read_text()
+
+
+def test_create_tar_cannot_exceed_max_tar_content_size(tmp_path):
+    with mock.patch("os.path.getsize", return_value=16 * 1024 * 1024):
+        file1_name = "file1.txt"
+        file2_name = "file2.txt"
+
+        (tmp_path / file1_name).write_text("")
+        (tmp_path / file2_name).write_text("")
+
+        with pytest.raises(
+            click.ClickException,
+            match=r"The total size of the files processed exceeds \d+MB, please try again with less files",
+        ):
+            _create_tar(tmp_path, [file1_name, file2_name])