Merge pull request #26 from nightfallai/dan/plat-1364-python-sdk-missing-requestmetadata-param

add request metadata

Author: Dan Hertz

nightfall/api.py

@@ -4,14 +4,11 @@
 ~~~~~~~~~~~~~
     This module provides a class which abstracts the Nightfall REST API.
 """
-import time
 from datetime import datetime, timedelta
 import hmac
 import hashlib
-import json
 import logging
 import os
-from functools import wraps
 from typing import List, Tuple, Optional
 
 import requests
@@ -104,7 +101,7 @@ def scan_text(self, texts: List[str], detection_rules: Optional[List[DetectionRu
         return findings, parsed_response.get("redactedPayload")
 
     def _scan_text_v3(self, data: dict):
-        response = self.session.post(url=self.TEXT_SCAN_ENDPOINT_V3, data=json.dumps(data))
+        response = self.session.post(url=self.TEXT_SCAN_ENDPOINT_V3, json=data)
 
         self.logger.debug(f"HTTP Request URL: {response.request.url}")
         self.logger.debug(f"HTTP Request Body: {response.request.body}")
@@ -119,7 +116,8 @@ def _scan_text_v3(self, data: dict):
 
     def scan_file(self, location: str, webhook_url: Optional[str] = None, policy_uuid: Optional[str] = None,
                   detection_rules: Optional[List[DetectionRule]] = None,
-                  detection_rule_uuids: Optional[List[str]] = None) -> Tuple[str, str]:
+                  detection_rule_uuids: Optional[List[str]] = None,
+                  request_metadata: Optional[str] = None) -> Tuple[str, str]:
         """Scan file with Nightfall.
         At least one of policy_uuid, detection_rule_uuids or detection_rules is required.
 
@@ -131,6 +129,8 @@ def scan_file(self, location: str, webhook_url: Optional[str] = None, policy_uui
         :type detection_rules: List[DetectionRule] or None
         :param detection_rule_uuids: list of detection rule UUIDs.
         :type detection_rule_uuids: List[str] or None
+        :param request_metadata: additional metadata that will be returned with the webhook response
+        :type request_metadata: str or None
         :returns: (scan_id, message)
         """
 
@@ -153,7 +153,8 @@ def scan_file(self, location: str, webhook_url: Optional[str] = None, policy_uui
         response = self._file_scan_scan(session_id,
                                         detection_rules=detection_rules,
                                         detection_rule_uuids=detection_rule_uuids,
-                                        webhook_url=webhook_url, policy_uuid=policy_uuid)
+                                        webhook_url=webhook_url, policy_uuid=policy_uuid,
+                                        request_metadata=request_metadata)
         _validate_response(response, 200)
         parsed_response = response.json()
 
@@ -163,7 +164,7 @@ def _file_scan_initialize(self, location: str):
         data = {
             "fileSizeBytes": os.path.getsize(location)
         }
-        response = self.session.post(url=self.FILE_SCAN_INITIALIZE_ENDPOINT, data=json.dumps(data))
+        response = self.session.post(url=self.FILE_SCAN_INITIALIZE_ENDPOINT, json=data)
 
         return response
 
@@ -200,7 +201,7 @@ def _file_scan_finalize(self, session_id: str):
 
     def _file_scan_scan(self, session_id: str, detection_rules: Optional[List[DetectionRule]] = None,
                         detection_rule_uuids: Optional[List[str]] = None, webhook_url: Optional[str] = None,
-                        policy_uuid: Optional[str] = None) -> requests.Response:
+                        policy_uuid: Optional[str] = None, request_metadata: Optional[str] = None) -> requests.Response:
         if policy_uuid:
             data = {"policyUUID": policy_uuid}
         else:
@@ -210,7 +211,10 @@ def _file_scan_scan(self, session_id: str, detection_rules: Optional[List[Detect
             if detection_rules:
                 data["policy"]["detectionRules"] = [d.as_dict() for d in detection_rules]
 
-        response = self.session.post(url=self.FILE_SCAN_SCAN_ENDPOINT.format(session_id), data=json.dumps(data))
+        if request_metadata:
+            data["requestMetadata"] = request_metadata
+
+        response = self.session.post(url=self.FILE_SCAN_SCAN_ENDPOINT.format(session_id), json=data)
         return response
 
     def validate_webhook(self, request_signature: str, request_timestamp: str, request_data: str) -> bool:

tests/test_api.py

@@ -263,24 +263,58 @@ def test_scan_file(tmpdir):
     responses.add(responses.POST, 'https://api.nightfall.ai/v3/upload/1/scan', status=200,
                   json={"id": 1, "message": "scan_started"})
 
-    id, message = nightfall.scan_file(file, "https://my-website.example/callback", detection_rule_uuids=["a_uuid"])
+    id, message = nightfall.scan_file(file, "https://my-website.example/callback", detection_rule_uuids=["a_uuid"], request_metadata="some test data")
 
     assert len(responses.calls) == 5
     for call in responses.calls:
         assert call.request.headers.get("Authorization") == "Bearer NF-NOT_REAL"
 
-    assert responses.calls[0].request.body == '{"fileSizeBytes": 44}'
-    assert responses.calls[1].request.body == "4916-6734-7572-5015 is"
+    assert responses.calls[0].request.body == b'{"fileSizeBytes": 44}'
+    assert responses.calls[1].request.body == b"4916-6734-7572-5015 is"
     assert responses.calls[1].request.headers.get("X-UPLOAD-OFFSET") == '0'
-    assert responses.calls[2].request.body == " my credit card number"
+    assert responses.calls[2].request.body == b" my credit card number"
     assert responses.calls[2].request.headers.get("X-UPLOAD-OFFSET") == '22'
-    assert responses.calls[4].request.body == '{"policy": {"webhookURL": "https://my-website.example/callback", ' \
-                                              '"detectionRuleUUIDs": ["a_uuid"]}}'
-
+    assert responses.calls[4].request.body == b'{"policy": {"webhookURL": "https://my-website.example/callback", ' \
+                                              b'"detectionRuleUUIDs": ["a_uuid"]}, "requestMetadata": "some test data"}'
     assert id == 1
     assert message == "scan_started"
 
 
+@responses.activate
+def test_file_scan_upload_short(tmpdir):
+    file = tmpdir.mkdir("test_data").join("file.txt")
+
+    file.write("4916-6734-7572-5015 is my credit card number")
+
+    nightfall = Nightfall("NF-NOT_REAL")
+
+    responses.add(responses.PATCH, 'https://api.nightfall.ai/v3/upload/1', status=204)
+
+    assert nightfall._file_scan_upload(1, file, 200)
+    assert len(responses.calls) == 1
+    assert responses.calls[0].request.headers.get("Authorization") == "Bearer NF-NOT_REAL"
+    assert responses.calls[0].request.body == b"4916-6734-7572-5015 is my credit card number"
+    assert responses.calls[0].request.headers.get("X-UPLOAD-OFFSET") == "0"
+
+
+@responses.activate
+def test_file_scan_upload_long(tmpdir):
+    file = tmpdir.mkdir("test_data").join("file.txt")
+    test_str = b"4916-6734-7572-5015 is my credit card number"
+    file.write_binary(test_str)
+
+    responses.add(responses.PATCH, 'https://api.nightfall.ai/v3/upload/1', status=204)
+
+    nightfall = Nightfall("NF-NOT_REAL")
+
+    assert nightfall._file_scan_upload(1, file, 1)
+    assert len(responses.calls) == 44
+    for i, call in enumerate(responses.calls):
+        assert call.request.headers.get("Authorization") == "Bearer NF-NOT_REAL"
+        assert call.request.body.decode('utf-8') == test_str.decode('utf-8')[i]
+        assert call.request.headers.get("X-UPLOAD-OFFSET") == str(i)
+
+
 @freeze_time("2021-10-04T17:30:50Z")
 def test_validate_webhook(nightfall):
     nightfall.signing_secret = "super-secret-shhhh"