298
298
from typing import Optional
299
299
from typing import NoReturn
300
300
from typing import Tuple
301
+ from typing import ByteString
301
302
302
303
try :
303
304
import boto3
@@ -339,7 +340,7 @@ def wrapped(self, *args, **kwargs):
339
340
for attempt in range (remaining_tries ):
340
341
try :
341
342
return_tuple = func (self , * args , ** kwargs )
342
- self ._vvvv (f"ssm_retry: (success ) { to_text (return_tuple )} " )
343
+ self ._vvvv (f"ssm_retry: (end of process ) { to_text (return_tuple )} " )
343
344
break
344
345
345
346
except (AnsibleConnectionFailure , Exception ) as e :
@@ -662,7 +663,33 @@ def poll(self, label: str, cmd: str) -> NoReturn:
662
663
raise AnsibleConnectionFailure (f"{ label } command '{ cmd } ' timeout on host: { self .instance_id } " )
663
664
yield self .poll_stdout ()
664
665
665
- def exec_communicate (self , cmd : str , mark_start : str , mark_begin : str , mark_end : str ) -> Tuple [int , str , str ]:
666
+ def _check_become_success (self , cmd : str ) -> ByteString :
667
+ become_output = b''
668
+ for result in self .poll ("ENSURE_BECOME" , cmd ):
669
+ if not result :
670
+ continue
671
+ become_output += self ._stdout .readline ()
672
+ if not self .become .check_success (become_output ) and not self .become .check_password_prompt (become_output ):
673
+ return ByteString
674
+
675
+ def ensure_become (self , cmd : str , sudoable : bool ) -> ByteString :
676
+ """When become is activated we ensure that the become step has succeed and if not the process may be waiting for
677
+ the password to be provided. In this case we will send it to the stdin
678
+ """
679
+ become_output = b''
680
+ if self .become and self .become .expect_prompt () and sudoable :
681
+ try :
682
+ become_output = self ._check_become_success (cmd )
683
+ except AnsibleConnectionFailure :
684
+ # Inject password into command
685
+ if not self .become .check_success (become_output ):
686
+ become_pass = self .become .get_option ('become_pass' , playcontext = self ._play_context )
687
+ self ._session .stdin .write (to_bytes (become_pass , errors = 'surrogate_or_strict' ) + b'\n ' )
688
+ become_output = self ._check_become_success (cmd )
689
+ return become_output
690
+
691
+
692
+ def exec_communicate (self , cmd : str , mark_start : str , mark_begin : str , mark_end : str , sudoable : bool ) -> Tuple [int , str , str ]:
666
693
"""Interact with session.
667
694
Read stdout between the markers until 'mark_end' is reached.
668
695
@@ -673,7 +700,7 @@ def exec_communicate(self, cmd: str, mark_start: str, mark_begin: str, mark_end:
673
700
:returns: A tuple with the return code, the stdout and the stderr content.
674
701
"""
675
702
# Read stdout between the markers
676
- stdout = ""
703
+ stdout = to_text ( self . ensure_become ( cmd , sudoable ))
677
704
win_line = ""
678
705
begin = False
679
706
returncode = None
@@ -708,6 +735,7 @@ def exec_communicate(self, cmd: str, mark_start: str, mark_begin: str, mark_end:
708
735
def exec_command (self , cmd : str , in_data : bool = None , sudoable : bool = True ) -> Tuple [int , str , str ]:
709
736
"""run a command on the ssm host"""
710
737
738
+ self ._v (f"[_exec_command] => { cmd } - sudoable = { sudoable } " )
711
739
super ().exec_command (cmd , in_data = in_data , sudoable = sudoable )
712
740
713
741
self ._vvv (f"EXEC: { to_text (cmd )} " )
@@ -727,7 +755,7 @@ def exec_command(self, cmd: str, in_data: bool = None, sudoable: bool = True) ->
727
755
for chunk in chunks (cmd , 1024 ):
728
756
self ._session .stdin .write (to_bytes (chunk , errors = "surrogate_or_strict" ))
729
757
730
- return self .exec_communicate (cmd , mark_start , mark_begin , mark_end )
758
+ return self .exec_communicate (cmd , mark_start , mark_begin , mark_end , sudoable )
731
759
732
760
def _prepare_terminal (self ):
733
761
"""perform any one-time terminal settings"""
@@ -1016,6 +1044,7 @@ def _file_transport_command(self, in_path, out_path, ssm_action):
1016
1044
def put_file (self , in_path , out_path ):
1017
1045
"""transfer a file from local to remote"""
1018
1046
1047
+ self ._v ("put_file" )
1019
1048
super ().put_file (in_path , out_path )
1020
1049
1021
1050
self ._vvv (f"PUT { in_path } TO { out_path } " )
@@ -1027,13 +1056,15 @@ def put_file(self, in_path, out_path):
1027
1056
def fetch_file (self , in_path , out_path ):
1028
1057
"""fetch a file from remote to local"""
1029
1058
1059
+ self ._v ("fetch_file" )
1030
1060
super ().fetch_file (in_path , out_path )
1031
1061
1032
1062
self ._vvv (f"FETCH { in_path } TO { out_path } " )
1033
1063
return self ._file_transport_command (in_path , out_path , "get" )
1034
1064
1035
1065
def close (self ):
1036
1066
"""terminate the connection"""
1067
+ self ._v ("close" )
1037
1068
if self ._session_id :
1038
1069
self ._vvv (f"CLOSING SSM CONNECTION TO: { self .instance_id } " )
1039
1070
if self ._has_timeout :
0 commit comments