Merge pull request #62 from GitGuardian/ppetit/-/jwt-creation

feat: add create_jwt() endpoint

Author: agateau-gg

changelog.d/20230613_163511_paul.petit.ext_jwt_creation.md

@@ -0,0 +1,3 @@
+### Added
+
+- Added `GGClient.create_jwt()` method. Only used to interact with HasMySecretLeaked for now.

pygitguardian/client.py

@@ -19,6 +19,8 @@
     DocumentSchema,
     HealthCheckResponse,
     HoneytokenResponse,
+    JWTResponse,
+    JWTService,
     MultiScanResult,
     QuotaResponse,
     ScanResult,
@@ -512,3 +514,27 @@ def read_metadata(self) -> Optional[Detail]:
 
         self.secret_scan_preferences = metadata.secret_scan_preferences
         return None
+
+    def create_jwt(
+        self,
+        jwt_audience: str,
+        service: JWTService,
+        extra_headers: Optional[Dict[str, str]] = None,
+    ) -> Union[Detail, JWTResponse]:
+        """
+        Create a JWT token for other GitGuardian services.
+        :return: Detail or JWT response and status code
+        """
+
+        resp = self.post(
+            endpoint="auth/jwt",
+            data={"audience": jwt_audience, "audience_type": service.value},
+            extra_headers=extra_headers,
+        )
+        obj: Union[Detail, JWTResponse]
+        if is_ok(resp):
+            obj = JWTResponse.from_dict(resp.json())
+        else:
+            obj = load_detail(resp)
+        obj.status_code = resp.status_code
+        return obj

pygitguardian/models.py

@@ -1,5 +1,6 @@
 from dataclasses import dataclass, field
 from datetime import date, datetime
+from enum import Enum
 from typing import Any, ClassVar, Dict, List, Optional, cast
 from uuid import UUID
 
@@ -668,3 +669,34 @@ class ServerMetadata(Base, FromDictMixin):
     BaseSchema,
     marshmallow_dataclass.class_schema(ServerMetadata, base_schema=BaseSchema)(),
 )
+
+
+class JWTResponseSchema(BaseSchema):
+    token = fields.String(required=True)
+
+    @post_load
+    def make_response(self, data: Dict[str, str], **kwargs: Any) -> "JWTResponse":
+        return JWTResponse(**data)
+
+
+class JWTResponse(Base, FromDictMixin):
+    """Token to use the HasMySecretLeaked service.
+
+    Attributes:
+        token (str): the JWT token
+    """
+
+    SCHEMA = JWTResponseSchema()
+
+    def __init__(self, token: str, **kwargs: Any) -> None:
+        super().__init__()
+        self.token = token
+
+    def __repr__(self) -> str:
+        return self.token
+
+
+class JWTService(Enum):
+    """Enum for the different services GIM can generate a JWT for."""
+
+    HMSL = "hmsl"

tests/test_client.py

@@ -20,6 +20,8 @@
 from pygitguardian.models import (
     Detail,
     HoneytokenResponse,
+    JWTResponse,
+    JWTService,
     MultiScanResult,
     QuotaResponse,
     ScanResult,
@@ -736,3 +738,26 @@ def test_create_honeytoken_error(
 
     assert mock_response.call_count == 1
     assert isinstance(result, Detail)
+
+
+@responses.activate
+def test_get_wt(
+    client: GGClient,
+):
+    """
+    GIVEN a ggclient
+    WHEN calling create_jwt
+    THEN we receive a token
+    """
+    mock_response = responses.post(
+        url=client._url_from_endpoint("auth/jwt", "v1"),
+        content_type="application/json",
+        status=200,
+        json={"token": "dummy_token"},
+    )
+
+    result = client.create_jwt("dummy_audience", JWTService.HMSL)
+
+    assert mock_response.call_count == 1
+    assert isinstance(result, JWTResponse)
+    assert result.token == "dummy_token"