Skip to content

Commit

Permalink
update ssm_document module
Browse files Browse the repository at this point in the history
  • Loading branch information
abikouo committed Feb 6, 2025
1 parent dbf0b8f commit c12f665
Show file tree
Hide file tree
Showing 2 changed files with 40 additions and 6 deletions.
39 changes: 35 additions & 4 deletions plugins/connection/aws_ssm.py
Original file line number Diff line number Diff line change
Expand Up @@ -298,6 +298,7 @@
from typing import Optional
from typing import NoReturn
from typing import Tuple
from typing import ByteString

try:
import boto3
Expand Down Expand Up @@ -339,7 +340,7 @@ def wrapped(self, *args, **kwargs):
for attempt in range(remaining_tries):
try:
return_tuple = func(self, *args, **kwargs)
self._vvvv(f"ssm_retry: (success) {to_text(return_tuple)}")
self._vvvv(f"ssm_retry: (end of process) {to_text(return_tuple)}")
break

except (AnsibleConnectionFailure, Exception) as e:
Expand Down Expand Up @@ -662,7 +663,33 @@ def poll(self, label: str, cmd: str) -> NoReturn:
raise AnsibleConnectionFailure(f"{label} command '{cmd}' timeout on host: {self.instance_id}")
yield self.poll_stdout()

def exec_communicate(self, cmd: str, mark_start: str, mark_begin: str, mark_end: str) -> Tuple[int, str, str]:
def _check_become_success(self, cmd: str) -> ByteString:
become_output = b''
for result in self.poll("ENSURE_BECOME", cmd):
if not result:
continue
become_output += self._stdout.readline()
if not self.become.check_success(become_output) and not self.become.check_password_prompt(become_output):
return ByteString

def ensure_become(self, cmd: str, sudoable: bool) -> ByteString:
"""When become is activated we ensure that the become step has succeed and if not the process may be waiting for
the password to be provided. In this case we will send it to the stdin
"""
become_output = b''
if self.become and self.become.expect_prompt() and sudoable:
try:
become_output = self._check_become_success(cmd)
except AnsibleConnectionFailure:
# Inject password into command
if not self.become.check_success(become_output):
become_pass = self.become.get_option('become_pass', playcontext=self._play_context)
self._session.stdin.write(to_bytes(become_pass, errors='surrogate_or_strict') + b'\n')
become_output = self._check_become_success(cmd)
return become_output


def exec_communicate(self, cmd: str, mark_start: str, mark_begin: str, mark_end: str, sudoable: bool) -> Tuple[int, str, str]:
"""Interact with session.
Read stdout between the markers until 'mark_end' is reached.
Expand All @@ -673,7 +700,7 @@ def exec_communicate(self, cmd: str, mark_start: str, mark_begin: str, mark_end:
:returns: A tuple with the return code, the stdout and the stderr content.
"""
# Read stdout between the markers
stdout = ""
stdout = to_text(self.ensure_become(cmd, sudoable))
win_line = ""
begin = False
returncode = None
Expand Down Expand Up @@ -708,6 +735,7 @@ def exec_communicate(self, cmd: str, mark_start: str, mark_begin: str, mark_end:
def exec_command(self, cmd: str, in_data: bool = None, sudoable: bool = True) -> Tuple[int, str, str]:
"""run a command on the ssm host"""

self._v(f"[_exec_command] => {cmd} - sudoable = {sudoable}")
super().exec_command(cmd, in_data=in_data, sudoable=sudoable)

self._vvv(f"EXEC: {to_text(cmd)}")
Expand All @@ -727,7 +755,7 @@ def exec_command(self, cmd: str, in_data: bool = None, sudoable: bool = True) ->
for chunk in chunks(cmd, 1024):
self._session.stdin.write(to_bytes(chunk, errors="surrogate_or_strict"))

return self.exec_communicate(cmd, mark_start, mark_begin, mark_end)
return self.exec_communicate(cmd, mark_start, mark_begin, mark_end, sudoable)

def _prepare_terminal(self):
"""perform any one-time terminal settings"""
Expand Down Expand Up @@ -1016,6 +1044,7 @@ def _file_transport_command(self, in_path, out_path, ssm_action):
def put_file(self, in_path, out_path):
"""transfer a file from local to remote"""

self._v("put_file")
super().put_file(in_path, out_path)

self._vvv(f"PUT {in_path} TO {out_path}")
Expand All @@ -1027,13 +1056,15 @@ def put_file(self, in_path, out_path):
def fetch_file(self, in_path, out_path):
"""fetch a file from remote to local"""

self._v("fetch_file")
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")

def close(self):
"""terminate the connection"""
self._v("close")
if self._session_id:
self._vvv(f"CLOSING SSM CONNECTION TO: {self.instance_id}")
if self._has_timeout:
Expand Down
7 changes: 5 additions & 2 deletions plugins/modules/ssm_document.py
Original file line number Diff line number Diff line change
Expand Up @@ -436,8 +436,11 @@ def update_document(
) != existing_document.get(aws_key):
params[aws_key] = document_parameters.get(ansible_key)

if params and document_parameters.get("attachments"):
params["Attachments"] = document_parameters.get("attachments")
if document_parameters.get("attachments"):
ansible_attachments = sorted([x.get("name") for x in document_parameters.get("attachments")])
aws_attachments = sorted([x.get("Name") for x in existing_document.get("AttachmentsInformation")])
if aws_attachments != ansible_attachments:
params["Attachments"] = document_parameters.get("attachments")

changed = False
ssm_document = None
Expand Down

0 comments on commit c12f665

Please sign in to comment.