1+ import io
12import os
23import tempfile
34from contextlib import contextmanager
@@ -65,30 +66,41 @@ def connect(self):
6566return ssh
6667
6768# Command execution
68- def exec_command (
69- self , cmd , wait_exit = False ,verbose = False , expect_error = False , encoding = "utf-8"
70- ):
69+ def exec_command (self , cmd , wait_exit = False , verbose = False ,
70+ expect_error = False ,encoding = None , shell = True , text = False ,
71+ input = None , stdout = None , stderr = None , proc = None ):
7172if isinstance (cmd ,list ):
7273cmd = " " .join (cmd )
7374log .debug (f"os_ops.exec_command: `{ cmd } `; remote={ self .remote } " )
7475# Source global profile file + execute command
7576try :
7677cmd = f"source /etc/profile.d/custom.sh;{ cmd } "
7778with self .ssh_connect ()as ssh :
78- stdin ,stdout ,stderr = ssh .exec_command (cmd )
79+ if input :
80+ # encode input and feed it to stdin
81+ stdin ,stdout ,stderr = ssh .exec_command (cmd )
82+ stdin .write (input )
83+ stdin .flush ()
84+ else :
85+ stdin ,stdout ,stderr = ssh .exec_command (cmd )
7986exit_status = 0
8087if wait_exit :
8188exit_status = stdout .channel .recv_exit_status ()
82- result = stdout .read ().decode (encoding )
83- error = stderr .read ().decode (encoding )
89+ if encoding :
90+ result = stdout .read ().decode (encoding )
91+ error = stderr .read ().decode (encoding )
92+ else :
93+ # Save as binary string
94+ result = io .BytesIO (stdout .read ()).getvalue ()
95+ error = io .BytesIO (stderr .read ()).getvalue ()
96+ error_str = stderr .read ()
8497
8598if expect_error :
8699raise Exception (result ,error )
87- if exit_status != 0 or " error" in error . lower () :
100+ if exit_status != 0 or ' error' in error_str :
88101log .error (
89102f"Problem in executing command: `{ cmd } `\n error:{ error } \n exit_code:{ exit_status } "
90103 )
91- exit (1 )
92104
93105if verbose :
94106return exit_status ,result ,error
@@ -203,9 +215,9 @@ def write(self, filename, data, truncate=False, binary=False, read_and_write=Fal
203215 """
204216mode = "wb" if binary else "w"
205217if not truncate :
206- mode = "a" + mode
218+ mode = "ab" if binary else "a"
207219if read_and_write :
208- mode = "r+" + mode
220+ mode = "r+b" if binary else "r+"
209221
210222with tempfile .NamedTemporaryFile (mode = mode )as tmp_file :
211223if isinstance (data ,list ):
@@ -229,17 +241,28 @@ def touch(self, filename):
229241 """
230242self .exec_command (f"touch{ filename } " )
231243
232- def read (self ,filename ,encoding = "utf-8" ):
244+ def read (self ,filename ,binary = False , encoding = None ):
233245cmd = f"cat{ filename } "
234- return self .exec_command (cmd ,encoding = encoding )
246+ result = self .exec_command (cmd ,encoding = encoding )
247+
248+ if not binary and result :
249+ result = result .decode (encoding or 'utf-8' )
250+
251+ return result
235252
236- def readlines (self ,filename ,num_lines = 0 ,encoding = None ):
237- encoding = encoding or "utf-8"
253+ def readlines (self ,filename ,num_lines = 0 ,binary = False ,encoding = None ):
238254if num_lines > 0 :
239255cmd = f"tail -n{ num_lines } { filename } "
240- lines = self .exec_command (cmd ,encoding )
241256else :
242- lines = self .read (filename ,encoding = encoding ).splitlines ()
257+ cmd = f"cat{ filename } "
258+
259+ result = self .exec_command (cmd ,encoding = encoding )
260+
261+ if not binary and result :
262+ lines = result .decode (encoding or 'utf-8' ).splitlines ()
263+ else :
264+ lines = result .splitlines ()
265+
243266return lines
244267
245268def isfile (self ,remote_file ):