made custom headers be available to async aws signer (#863)

* made custom headers be available to async aws signer Signed-off-by: Bruno Murino <[email protected]> * updated changelog Signed-off-by: Bruno Murino <[email protected]> * added tests for using host header for AWS request signature on both sync and async clients Signed-off-by: Bruno Murino <[email protected]> * added documentation guide about aws auth when accessing via tunnel Signed-off-by: Bruno Murino <[email protected]> --------- Signed-off-by: Bruno Murino <[email protected]>
opensearch-project · Dec 1, 2024 · 6f761ab · 6f761ab
1 parent 090e11e
commit 6f761ab
Show file tree

Hide file tree

Showing 6 changed files with 105 additions and 6 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -3,6 +3,7 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 
 ## [Unreleased]
 ### Added
+- Added option to pass custom headers to 'AWSV4SignerAsyncAuth' ([863](https://github.com/opensearch-project/opensearch-py/pull/863))
 ### Updated APIs
 ### Changed
 ### Deprecated

diff --git a/guides/auth.md b/guides/auth.md
@@ -2,6 +2,7 @@
   - [IAM Authentication](#iam-authentication)
   - [IAM Authentication with a Synchronous Client](#iam-authentication-with-a-synchronous-client)
   - [IAM Authentication with an Async Client](#iam-authentication-with-an-async-client)
+  - [IAM Authentication via Tunnel](#iam-authentication-via-tunnel)
   - [Kerberos](#kerberos)
 
 # Authentication
@@ -104,6 +105,43 @@ async def search():
 search()
 ```
 
+## IAM Authentication via Tunnel
+
+If you're accessing OpenSearch via SSH or SSM tunnel, then you need to specify the Host to be used for signing the AWS requests by passing a "Host" header, like so:
+
+
+```python
+from opensearchpy import OpenSearch, RequestsHttpConnection, RequestsAWSV4SignerAuth, AsyncOpenSearch, AsyncHttpConnection, AWSV4SignerAsyncAuth
+import boto3
+
+host = 'localhost' # local endpoint used by the SSH/SSM tunnel
+port = 8443
+signature_host = 'my-test-domain.eu-west-1.es.amazonaws.com:443' # this needs to be the real host provided by AWS
+region = 'eu-west-1'
+service = 'es' # 'aoss' for OpenSearch Serverless
+credentials = boto3.Session().get_credentials()
+
+# Sync
+client = OpenSearch(
+    hosts = [{'host': host, 'port': port, 'headers': {'host': signature_host}}],
+    http_auth = RequestsAWSV4SignerAuth(credentials, region, service),
+    use_ssl = True,
+    verify_certs = True,
+    connection_class = RequestsHttpConnection,
+    pool_maxsize = 20
+)
+
+# Async
+async_client = AsyncOpenSearch(
+    hosts = [{'host': host, 'port': port, 'headers': {'host': signature_host}}],
+    http_auth = AWSV4SignerAsyncAuth(credentials, region, service),
+    use_ssl = True,
+    verify_certs = True,
+    connection_class = AsyncHttpConnection
+)
+
+```
+
 ## Kerberos
 
 There are several python packages that provide Kerberos support over HTTP, such as [requests-kerberos](http://pypi.org/project/requests-kerberos) and [requests-gssapi](https://pypi.org/project/requests-gssapi). The following example shows how to setup Kerberos authentication. 

diff --git a/opensearchpy/helpers/asyncsigner.py b/opensearchpy/helpers/asyncsigner.py
@@ -8,6 +8,7 @@
 # GitHub history for details.
 
 from typing import Any, Dict, Optional, Union
+from urllib.parse import parse_qs, urlencode, urlparse
 
 
 class AWSV4SignerAsyncAuth:
@@ -34,15 +35,17 @@ def __call__(
         url: str,
         query_string: Optional[str] = None,
         body: Optional[Union[str, bytes]] = None,
+        headers: Optional[Dict[str, str]] = None,
     ) -> Dict[str, str]:
-        return self._sign_request(method, url, query_string, body)
+        return self._sign_request(method, url, query_string, body, headers)
 
     def _sign_request(
         self,
         method: str,
         url: str,
         query_string: Optional[str],
         body: Optional[Union[str, bytes]],
+        headers: Optional[Dict[str, str]],
     ) -> Dict[str, str]:
         """
         This method helps in signing the request by injecting the required headers.
@@ -53,10 +56,12 @@ def _sign_request(
         from botocore.auth import SigV4Auth
         from botocore.awsrequest import AWSRequest
 
+        signature_host = self._fetch_url(url, headers or dict())
+
         # create an AWS request object and sign it using SigV4Auth
         aws_request = AWSRequest(
             method=method,
-            url=url,
+            url=signature_host,
             data=body,
         )
 
@@ -80,3 +85,26 @@ def _sign_request(
 
         # copy the headers from AWS request object into the prepared_request
         return dict(aws_request.headers.items())
+
+    def _fetch_url(self, url: str, headers: Optional[Dict[str, str]]) -> str:
+        """
+        This is a util method that helps in reconstructing the request url.
+        :param prepared_request: unsigned request
+        :return: reconstructed url
+        """
+        parsed_url = urlparse(url)
+        path = parsed_url.path or "/"
+
+        # fetch the query string if present in the request
+        querystring = ""
+        if parsed_url.query:
+            querystring = "?" + urlencode(
+                parse_qs(parsed_url.query, keep_blank_values=True), doseq=True
+            )
+
+        # fetch the host information from headers
+        headers = {key.lower(): value for key, value in (headers or dict()).items()}
+        location = headers.get("host") or parsed_url.netloc
+
+        # construct the url and return
+        return parsed_url.scheme + "://" + location + path + querystring
diff --git a/opensearchpy/helpers/signer.py b/opensearchpy/helpers/signer.py
@@ -92,14 +92,14 @@ def _sign_request(self, prepared_request):  # type: ignore
         prepared_request.headers.update(
             self.signer.sign(
                 prepared_request.method,
-                self._fetch_url(prepared_request),  # type: ignore
+                self._fetch_url(prepared_request),
                 prepared_request.body,
             )
         )
 
         return prepared_request
 
-    def _fetch_url(self, prepared_request):  # type: ignore
+    def _fetch_url(self, prepared_request: requests.PreparedRequest) -> str:
         """
         This is a util method that helps in reconstructing the request url.
         :param prepared_request: unsigned request
@@ -112,7 +112,7 @@ def _fetch_url(self, prepared_request):  # type: ignore
         querystring = ""
         if url.query:
             querystring = "?" + urlencode(
-                parse_qs(url.query, keep_blank_values=True), doseq=True
+                parse_qs(url.query, keep_blank_values=True), doseq=True  # type: ignore
             )
 
         # fetch the host information from headers
@@ -122,7 +122,7 @@ def _fetch_url(self, prepared_request):  # type: ignore
         location = headers.get("host") or url.netloc
 
         # construct the url and return
-        return url.scheme + "://" + location + path + querystring
+        return url.scheme + "://" + location + path + querystring  # type: ignore
 
 
 # Deprecated: use RequestsAWSV4SignerAuth

diff --git a/test_opensearchpy/test_async/test_signer.py b/test_opensearchpy/test_async/test_signer.py
@@ -77,6 +77,20 @@ async def test_aws_signer_async_when_service_is_specified(self) -> None:
         assert "X-Amz-Security-Token" in headers
         assert "X-Amz-Content-SHA256" in headers
 
+    async def test_aws_signer_async_fetch_url_with_querystring(self) -> None:
+        region = "us-west-2"
+        service = "aoss"
+
+        from opensearchpy.helpers.asyncsigner import AWSV4SignerAsyncAuth
+
+        auth = AWSV4SignerAsyncAuth(self.mock_session(), region, service)
+
+        signature_host = auth._fetch_url(
+            "http://localhost/?foo=bar", headers={"host": "otherhost"}
+        )
+
+        assert signature_host == "http://otherhost/?foo=bar"
+
 
 class TestAsyncSignerWithFrozenCredentials(TestAsyncSigner):
     def mock_session(self, disable_get_frozen: bool = True) -> Mock:
@@ -143,6 +157,7 @@ def _sign_request(
                 url: str,
                 query_string: Optional[str] = None,
                 body: Optional[Union[str, bytes]] = None,
+                headers: Optional[Dict[str, str]] = None,
             ) -> Dict[str, str]:
                 nonlocal signed_url
                 signed_url = url

diff --git a/test_opensearchpy/test_connection/test_requests_http_connection.py b/test_opensearchpy/test_connection/test_requests_http_connection.py
@@ -457,6 +457,23 @@ def mock_session(self) -> Any:
 
         return dummy_session
 
+    def test_aws_signer_fetch_url_with_querystring(self) -> None:
+        region = "us-west-2"
+
+        import requests
+
+        from opensearchpy.helpers.signer import RequestsAWSV4SignerAuth
+
+        auth = RequestsAWSV4SignerAuth(self.mock_session(), region)
+
+        prepared_request = requests.Request(
+            "GET", "http://localhost/?foo=bar", headers={"host": "otherhost:443"}
+        ).prepare()
+
+        signature_host = auth._fetch_url(prepared_request)
+
+        assert signature_host == "http://otherhost:443/?foo=bar"
+
     def test_aws_signer_as_http_auth(self) -> None:
         region = "us-west-2"