Sanity test ok

Author: abikouo

Diff for: plugins/connection/aws_ssm.py

@@ -86,6 +86,27 @@
     vars:
     - name: ansible_aws_ssm_bucket_endpoint_url
     version_added: 5.3.0
+  host_port_number:
+    description:
+    - The Port number of the server on the instance when using Port Forwarding Using AWS System Manager Session Manager
+      to transfer files from/to local host to/from remote host.
+    - The port V(80) is used if not provided.
+    - This is not supported for Windows hosts for now.
+    type: integer
+    default: 80
+    vars:
+    - name: ansible_aws_ssm_host_port_number
+    version_added: 9.3.0
+  local_port_number:
+    description:
+    - Port number on local machine to forward traffic to when using Port Forwarding Using AWS System Manager Session Manager
+      to transfer files from/to local host to/from remote host.
+    - An open port is chosen at run-time if not provided.
+    - This is not supported for Windows hosts for now.
+    type: integer
+    vars:
+    - name: ansible_aws_ssm_local_port_number
+    version_added: 9.3.0
   plugin:
     description:
     - This defines the location of the session-manager-plugin binary.
@@ -359,6 +380,7 @@
 from ansible.utils.display import Display
 
 from ansible_collections.amazon.aws.plugins.module_utils.botocore import HAS_BOTO3
+from ansible_collections.community.aws.plugins.plugin_utils.ssm_file_transfer import PortForwardingFileTransferManager
 
 display = Display()
 
@@ -469,6 +491,7 @@ class Connection(ConnectionBase):
     _stdout = None
     _session_id = ""
     _timeout = False
+    _filetransfer_mgr = None
     MARK_LENGTH = 26
 
     def __init__(self, *args, **kwargs):
@@ -517,8 +540,29 @@ def _init_clients(self) -> None:
         # Initialize SSM client
         self._initialize_ssm_client(region_name, profile_name)
 
-        # Initialize S3 client
-        self._initialize_s3_client(profile_name)
+        if self._should_use_port_forwarding_for_file_transfer():
+            # Initialize S3 client
+            self._initialize_s3_client(profile_name)
+        else:
+            self._initialize_file_transfer_manager()
+
+    def _initialize_file_transfer_manager(self) -> None:
+        ssm_timeout = self.get_option("ssm_timeout")
+        region_name = self.get_option("region")
+        profile_name = self.get_option("profile") or ""
+        host_port = self.get_option("host_port_number")
+        local_port = self.get_option("local_port_number")
+        self._filetransfer_mgr = PortForwardingFileTransferManager(
+            self.host,
+            ssm_client=self._client,
+            instance_id=self.instance_id,
+            executable=self.get_executable(),
+            ssm_timeout=ssm_timeout,
+            region_name=region_name,
+            profile_name=profile_name,
+            host_port=host_port,
+            local_port=local_port,
+        )
 
     def _initialize_ssm_client(self, region_name: Optional[str], profile_name: str) -> None:
         """
@@ -621,6 +665,10 @@ def reset(self):
         self.close()
         return self.start_session()
 
+    def _should_use_port_forwarding_for_file_transfer(self) -> bool:
+        """return true if the user has defined a bucket_name to be used for transport"""
+        return (not self.is_windows and self.get_option("bucket_name") is not None)
+
     @property
     def instance_id(self) -> str:
         if not self._instance_id:
@@ -1159,15 +1207,21 @@ def put_file(self, in_path, out_path):
         if not os.path.exists(to_bytes(in_path, errors="surrogate_or_strict")):
             raise AnsibleFileNotFound(f"file or module does not exist: {in_path}")
 
-        return self._file_transport_command(in_path, out_path, "put")
+        if self._should_use_port_forwarding_for_file_transfer():
+            return self._file_transport_command(in_path, out_path, "put")
+        else:
+            return self._filetransfer_mgr.put_file(in_path, out_path)
 
     def fetch_file(self, in_path, out_path):
         """fetch a file from remote to local"""
 
         super().fetch_file(in_path, out_path)
 
         self._vvv(f"FETCH {in_path} TO {out_path}")
-        return self._file_transport_command(in_path, out_path, "get")
+        if self._should_use_port_forwarding_for_file_transfer():
+            return self._file_transport_command(in_path, out_path, "get")
+        else:
+            return self._filetransfer_mgr.fetch_file(in_path, out_path)
 
     def close(self):
         """terminate the connection"""