Repository: RaRe-Technologies/smart_open
Pull request #703 (2022-07-02 22:10:13)

Add cert parameter to http transport params
#### Allow client certificate authentication for http

#### Motivation

Although it seems to be a simple case of adding http transport parameters and passing them  to `requests`, there is no support for client certificate authentication for http

#### Tests

Added `test_smart_open::SmartOpenHttpTest::test_http_cert`




Patch:
diff --git a/help.txt b/help.txt
index 6d642bb..b5e250e 100644
--- a/help.txt
+++ b/help.txt
@@ -124,6 +124,8 @@ FUNCTIONS
             The username for authenticating over HTTP
         password: str, optional
             The password for authenticating over HTTP
+        cert: str/tuple, optional
+            If String, path to ssl client cert file (.pem). If Tuple, (‘cert’, ‘key’)
         headers: dict, optional
             Any headers to send in the request. If ``None``, the default headers are sent:
             ``{'Accept-Encoding': 'identity'}``. To use no headers at all,
diff --git a/smart_open/azure.py b/smart_open/azure.py
index b3a9bad..15168a9 100644
--- a/smart_open/azure.py
+++ b/smart_open/azure.py
@@ -70,6 +70,7 @@ def open(
         blob_id,
         mode,
         client=None,  # type: Union[azure.storage.blob.BlobServiceClient, azure.storage.blob.ContainerClient, azure.storage.blob.BlobClient] # noqa
+        blob_kwargs=None,
         buffer_size=DEFAULT_BUFFER_SIZE,
         min_part_size=_DEFAULT_MIN_PART_SIZE,
         max_concurrency=DEFAULT_MAX_CONCURRENCY,
@@ -86,10 +87,13 @@ def open(
         The mode for opening the object.  Must be either "rb" or "wb".
     client: azure.storage.blob.BlobServiceClient, ContainerClient, or BlobClient
         The Azure Blob Storage client to use when working with azure-storage-blob.
+    blob_kwargs: dict, optional
+        Additional parameters to pass to `BlobClient.commit_block_list`.
+        For writing only.
     buffer_size: int, optional
         The buffer size to use when performing I/O. For reading only.
     min_part_size: int, optional
-        The minimum part size for multipart uploads.  For writing only.
+        The minimum part size for multipart uploads. For writing only.
     max_concurrency: int, optional
         The number of parallel connections with which to download. For reading only.
 
@@ -111,6 +115,7 @@ def open(
             container_id,
             blob_id,
             client,
+            blob_kwargs=blob_kwargs,
             min_part_size=min_part_size
         )
     else:
@@ -387,12 +392,14 @@ class Writer(io.BufferedIOBase):
             container,
             blob,
             client,  # type: Union[azure.storage.blob.BlobServiceClient, azure.storage.blob.ContainerClient, azure.storage.blob.BlobClient]  # noqa
+            blob_kwargs=None,
             min_part_size=_DEFAULT_MIN_PART_SIZE,
     ):
         self._is_closed = False
         self._container_name = container
 
         self._blob = _get_blob_client(client, container, blob)
+        self._blob_kwargs = blob_kwargs or {}
         # type: azure.storage.blob.BlobClient
 
         self._min_part_size = min_part_size
@@ -419,7 +426,7 @@ class Writer(io.BufferedIOBase):
         if not self.closed:
             if self._current_part.tell() > 0:
                 self._upload_part()
-            self._blob.commit_block_list(self._block_list)
+            self._blob.commit_block_list(self._block_list, **self._blob_kwargs)
             self._block_list = []
             self._is_closed = True
         logger.debug("successfully closed")
diff --git a/smart_open/http.py b/smart_open/http.py
index e4439bd..e6fe2ae 100644
--- a/smart_open/http.py
+++ b/smart_open/http.py
@@ -49,7 +49,8 @@ def open_uri(uri, mode, transport_params):
     return open(uri, mode, **kwargs)
 
 
-def open(uri, mode, kerberos=False, user=None, password=None, headers=None, timeout=None):
+def open(uri, mode, kerberos=False, user=None, password=None, cert=None,
+         headers=None, timeout=None):
     """Implement streamed reader from a web site.
 
     Supports Kerberos and Basic HTTP authentication.
@@ -66,6 +67,8 @@ def open(uri, mode, kerberos=False, user=None, password=None, headers=None, time
         The username for authenticating over HTTP
     password: str, optional
         The password for authenticating over HTTP
+    cert: str/tuple, optional
+        if String, path to ssl client cert file (.pem). If Tuple, (‘cert’, ‘key’)
     headers: dict, optional
         Any headers to send in the request. If ``None``, the default headers are sent:
         ``{'Accept-Encoding': 'identity'}``. To use no headers at all,
@@ -80,7 +83,8 @@ def open(uri, mode, kerberos=False, user=None, password=None, headers=None, time
     if mode == constants.READ_BINARY:
         fobj = SeekableBufferedInputBase(
             uri, mode, kerberos=kerberos,
-            user=user, password=password, headers=headers, timeout=timeout,
+            user=user, password=password, cert=cert,
+            headers=headers, timeout=timeout,
         )
         fobj.name = os.path.basename(urllib.parse.urlparse(uri).path)
         return fobj
@@ -90,7 +94,8 @@ def open(uri, mode, kerberos=False, user=None, password=None, headers=None, time
 
 class BufferedInputBase(io.BufferedIOBase):
     def __init__(self, url, mode='r', buffer_size=DEFAULT_BUFFER_SIZE,
-                 kerberos=False, user=None, password=None, headers=None, timeout=None):
+                 kerberos=False, user=None, password=None, cert=None,
+                 headers=None, timeout=None):
         if kerberos:
             import requests_kerberos
             auth = requests_kerberos.HTTPKerberosAuth()
@@ -112,6 +117,7 @@ class BufferedInputBase(io.BufferedIOBase):
         self.response = requests.get(
             url,
             auth=auth,
+            cert=cert,
             stream=True,
             headers=self.headers,
             timeout=self.timeout,
@@ -204,13 +210,15 @@ class BufferedInputBase(io.BufferedIOBase):
 class SeekableBufferedInputBase(BufferedInputBase):
     """
     Implement seekable streamed reader from a web site.
-    Supports Kerberos and Basic HTTP authentication.
+    Supports Kerberos, client certificate and Basic HTTP authentication.
     """
 
     def __init__(self, url, mode='r', buffer_size=DEFAULT_BUFFER_SIZE,
-                 kerberos=False, user=None, password=None, headers=None, timeout=None):
+                 kerberos=False, user=None, password=None, cert=None,
+                 headers=None, timeout=None):
         """
         If Kerberos is True, will attempt to use the local Kerberos credentials.
+        If cert is set, will try to use a client certificate
         Otherwise, will try to use "basic" HTTP authentication via username/password.
 
         If none of those are set, will connect unauthenticated.
@@ -230,6 +238,7 @@ class SeekableBufferedInputBase(BufferedInputBase):
         else:
             self.headers = headers
 
+        self.cert = cert
         self.timeout = timeout
 
         self.buffer_size = buffer_size
@@ -325,6 +334,7 @@ class SeekableBufferedInputBase(BufferedInputBase):
             self.url,
             auth=self.auth,
             stream=True,
+            cert=self.cert,
             headers=self.headers,
             timeout=self.timeout,
         )


Test patch:
diff --git a/smart_open/tests/test_azure.py b/smart_open/tests/test_azure.py
index d249748..3c668aa 100644
--- a/smart_open/tests/test_azure.py
+++ b/smart_open/tests/test_azure.py
@@ -51,10 +51,12 @@ class FakeBlobClient(object):
         self.__contents = io.BytesIO()
         self._staged_contents = {}
 
-    def commit_block_list(self, block_list):
+    def commit_block_list(self, block_list, metadata=None):
         data = b''.join([self._staged_contents[block_blob['id']] for block_blob in block_list])
         self.__contents = io.BytesIO(data)
-        self.set_blob_metadata(dict(size=len(data)))
+        metadata = metadata or {}
+        metadata.update({"size": len(data)})
+        self.set_blob_metadata(metadata)
         self._container_client.register_blob_client(self)
 
     def delete_blob(self):
@@ -590,6 +592,32 @@ class WriterTest(unittest.TestCase):
         ))
         assert output == [test_string]
 
+    def test_write_blob_client(self):
+        """Does writing into Azure Blob Storage work correctly?"""
+        test_string = u"žluťoučký koníček".encode('utf8')
+        blob_name = "test_write_blob_client_%s" % BLOB_NAME
+
+        container_client = CLIENT.get_container_client(CONTAINER_NAME)
+        blob_client = container_client.get_blob_client(blob_name)
+
+        with smart_open.open(
+            "azure://%s/%s" % (CONTAINER_NAME, blob_name),
+            "wb",
+            transport_params={
+                "client": blob_client, "blob_kwargs": {"metadata": {"name": blob_name}}
+            },
+        ) as fout:
+            fout.write(test_string)
+
+        self.assertEqual(blob_client.get_blob_properties()["name"], blob_name)
+
+        output = list(smart_open.open(
+            "azure://%s/%s" % (CONTAINER_NAME, blob_name),
+            "rb",
+            transport_params=dict(client=CLIENT),
+        ))
+        self.assertEqual(output, [test_string])
+
     def test_incorrect_input(self):
         """Does azure write fail on incorrect input?"""
         blob_name = "test_incorrect_input_%s" % BLOB_NAME
diff --git a/smart_open/tests/test_smart_open.py b/smart_open/tests/test_smart_open.py
index 5434c16..367e994 100644
--- a/smart_open/tests/test_smart_open.py
+++ b/smart_open/tests/test_smart_open.py
@@ -448,6 +448,18 @@ class SmartOpenHttpTest(unittest.TestCase):
         self.assertTrue('Authorization' in actual_request.headers)
         self.assertTrue(actual_request.headers['Authorization'].startswith('Basic '))
 
+    @responses.activate
+    def test_http_cert(self):
+        """Does cert parameter get passed to requests"""
+        responses.add(responses.GET, "http://127.0.0.1/index.html",
+                      body='line1\nline2', stream=True)
+        cert_path = '/path/to/my/cert.pem'
+        tp = dict(cert=cert_path)
+        smart_open.open("http://127.0.0.1/index.html", transport_params=tp)
+        self.assertEqual(len(responses.calls), 1)
+        actual_request = responses.calls[0].request
+        self.assertEqual(cert_path, actual_request.req_kwargs['cert'])
+
     @responses.activate
     def _test_compressed_http(self, suffix, query):
         """Can open <suffix> via http?"""