From a6324e847f272c332738c8c1c9d3c5d61a202261 Mon Sep 17 00:00:00 2001
From: Bruno Murino <bruno.murino@landbay.co.uk>
Date: Sat, 30 Nov 2024 17:20:51 +0000
Subject: [PATCH] feat: allow to set a signature host for tunnel usage with
 async client

Signed-off-by: Bruno Murino <bfsmurino@gmail.com>
---
 opensearchpy/helpers/asyncsigner.py         | 32 +++++++++++++++++++--
 test_opensearchpy/test_async/test_signer.py |  1 +
 2 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/opensearchpy/helpers/asyncsigner.py b/opensearchpy/helpers/asyncsigner.py
index c045f1384..52e0b092e 100644
--- 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,8 +35,9 @@ 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,
@@ -43,6 +45,7 @@ def _sign_request(
         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())  # type: ignore
+
         # 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, headers):  # type: ignore
+        """
+        This is a util method that helps in reconstructing the request url.
+        :param prepared_request: unsigned request
+        :return: reconstructed url
+        """
+        url = urlparse(url)
+        path = url.path or "/"
+
+        # fetch the query string if present in the request
+        querystring = ""
+        if url.query:
+            querystring = "?" + urlencode(
+                parse_qs(url.query, keep_blank_values=True), doseq=True
+            )
+
+        # fetch the host information from headers
+        headers = {key.lower(): value for key, value in headers.items()}
+        location = headers.get("host") or url.netloc
+
+        # construct the url and return
+        return url.scheme + "://" + location + path + querystring
diff --git a/test_opensearchpy/test_async/test_signer.py b/test_opensearchpy/test_async/test_signer.py
index 473faa76e..7995797fa 100644
--- a/test_opensearchpy/test_async/test_signer.py
+++ b/test_opensearchpy/test_async/test_signer.py
@@ -143,6 +143,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