<s> """<STR_LIT>""" <EOL> from __future__ import with_statement <EOL> from io import BytesIO <EOL> import errno <EOL> import logging <EOL> import os <EOL> import select <EOL> import shutil <EOL> import subprocess <EOL> import sys <EOL> import tarfile <EOL> import time <EOL> from . import errors <EOL> from . import colors <EOL> DONE = <NUM_LIT:0> <EOL> STDOUT = <NUM_LIT:1> <EOL> STDERR = <NUM_LIT:2> <EOL> class RemoteControl ( object ) : <EOL> def __init__ ( self , node ) : <EOL> self . node = node <EOL> self . warn_timeout = <NUM_LIT> <EOL> self . terminate_timeout = node . get_tree_property ( "<STR_LIT>" , <NUM_LIT> ) <EOL> def get_out_line ( self , color , tag , command , result ) : <EOL> desc = "<STR_LIT>" % ( color ( self . node . name , "<STR_LIT>" ) , <EOL> color ( self . node . get ( "<STR_LIT:host>" ) , "<STR_LIT:host>" ) , <EOL> color ( command , "<STR_LIT>" ) ) <EOL> if result is not None : <EOL> tag = "<STR_LIT>" % ( tag , result ) <EOL> out_line = "<STR_LIT>" % ( color ( "<STR_LIT>" , "<STR_LIT>" ) , <EOL> tag , <EOL> desc , <EOL> color ( "<STR_LIT>" , "<STR_LIT>" ) ) <EOL> return out_line <EOL> def tag_line ( self , tag , command , result = None , verbose = False , color = None , <EOL> out_file = None ) : <EOL> assert color <EOL> if out_file and not out_file . isatty ( ) : <EOL> no_color = self . get_color ( None , out_file = out_file ) <EOL> plain_out = self . get_out_line ( no_color , tag , command , result ) <EOL> out_file . write ( plain_out ) <EOL> if verbose : <EOL> color_out = self . get_out_line ( color , tag , command , result ) <EOL> sys . stdout . write ( color_out ) <EOL> def get_color ( self , color , out_file = None ) : <EOL> out_file = out_file or sys . stdout <EOL> if color : <EOL> return color <EOL> else : <EOL> return colors . Output ( out_file , color = "<STR_LIT>" ) . color <EOL> def execute ( self , command , verbose = False , color = None , output_lines = None , <EOL> output_file = None , quiet = False , exec_options = None ) : <EOL> exec_options = exec_options or { } <EOL> if output_file is not None : <EOL> stdout_file = output_file <EOL> stderr_file = output_file <EOL> elif not quiet : <EOL> stdout_file = sys . stdout <EOL> stderr_file = sys . stderr <EOL> else : <EOL> stdout_file = None <EOL> stderr_file = None <EOL> result = None <EOL> output_chunks = [ ] <EOL> color = self . get_color ( color , out_file = stdout_file ) <EOL> self . tag_line ( color ( "<STR_LIT>" , "<STR_LIT>" ) , command , verbose = verbose , <EOL> color = color , out_file = stdout_file ) <EOL> start = time . time ( ) <EOL> try : <EOL> while True : <EOL> for code , output in self . execute_command ( command , <EOL> ** exec_options ) : <EOL> if code == STDOUT : <EOL> if output_lines is not None : <EOL> output_chunks . append ( output ) <EOL> elif stdout_file : <EOL> if not verbose or stdout_file : <EOL> stdout_file . write ( output ) <EOL> else : <EOL> for line in output . splitlines ( True ) : <EOL> stdout_file . write ( <EOL> "<STR_LIT>" % ( color ( self . node . name , <EOL> "<STR_LIT>" ) , line ) ) <EOL> stdout_file . flush ( ) <EOL> elif code == STDERR : <EOL> if not verbose or stdout_file : <EOL> stderr_file . write ( output ) <EOL> elif stderr_file : <EOL> for line in output . splitlines ( True ) : <EOL> stderr_file . write ( <EOL> "<STR_LIT>" % ( color ( self . node . name , <EOL> "<STR_LIT>" ) , line ) ) <EOL> stderr_file . flush ( ) <EOL> else : <EOL> if output_lines is not None : <EOL> output_lines . extend ( <EOL> ( "<STR_LIT>" . join ( output_chunks ) ) . splitlines ( ) ) <EOL> if not output : <EOL> result = color ( "<STR_LIT:OK>" , "<STR_LIT>" ) <EOL> else : <EOL> result = color ( output , "<STR_LIT>" ) <EOL> return output <EOL> except Exception as error : <EOL> result = color ( "<STR_LIT>" % ( error . __class__ . __name__ , error ) , <EOL> "<STR_LIT>" ) <EOL> raise <EOL> finally : <EOL> elapsed = time . time ( ) - start <EOL> self . tag_line ( color ( "<STR_LIT>" % elapsed , "<STR_LIT>" ) , command , result = result , <EOL> verbose = verbose , color = color , out_file = stdout_file ) <EOL> def shell ( self , verbose = False , color = None ) : <EOL> color = self . get_color ( color ) <EOL> self . tag_line ( "<STR_LIT>" , "<STR_LIT>" , verbose = verbose , color = color ) <EOL> try : <EOL> return self . execute_shell ( ) <EOL> finally : <EOL> self . tag_line ( "<STR_LIT>" , "<STR_LIT>" , verbose = verbose , color = color ) <EOL> def close ( self ) : <EOL> pass <EOL> def stat ( self , file_path ) : <EOL> assert <NUM_LIT:0> , "<STR_LIT>" <EOL> def read_file ( self , file_path ) : <EOL> assert <NUM_LIT:0> , "<STR_LIT>" <EOL> def put_file ( self , source_path , dest_path , callback = None ) : <EOL> assert <NUM_LIT:0> , "<STR_LIT>" <EOL> def write_file ( self , file_path , contents , mode = None , owner = None , <EOL> group = None ) : <EOL> assert <NUM_LIT:0> , "<STR_LIT>" <EOL> def execute_command ( self , command , pseudo_tty = False ) : <EOL> assert <NUM_LIT:0> , "<STR_LIT>" <EOL> def execute_shell ( self ) : <EOL> assert <NUM_LIT:0> , "<STR_LIT>" <EOL> def makedirs ( self , dir_path ) : <EOL> assert <NUM_LIT:0> , "<STR_LIT>" <EOL> def convert_local_errors ( method ) : <EOL> """<STR_LIT>""" <EOL> def wrapper ( self , * args , ** kw ) : <EOL> try : <EOL> return method ( self , * args , ** kw ) <EOL> except IOError as error : <EOL> if error . errno == errno . ENOENT : <EOL> raise errors . RemoteFileDoesNotExist ( str ( error ) ) <EOL> else : <EOL> raise errors . RemoteError ( "<STR_LIT>" % ( error . __class__ . __name__ , <EOL> error ) ) <EOL> except OSError as error : <EOL> raise errors . RemoteError ( "<STR_LIT>" % ( error . __class__ . __name__ , <EOL> error ) ) <EOL> wrapper . __doc__ = method . __doc__ <EOL> wrapper . __name__ = method . __name__ <EOL> return wrapper <EOL> class LocalControl ( RemoteControl ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , node ) : <EOL> RemoteControl . __init__ ( self , node ) <EOL> @ convert_local_errors <EOL> def put_file ( self , source_path , dest_path , callback = None ) : <EOL> shutil . copy ( source_path , dest_path ) <EOL> @ convert_local_errors <EOL> def makedirs ( self , dir_path ) : <EOL> return os . makedirs ( dir_path ) <EOL> @ convert_local_errors <EOL> def read_file ( self , file_path ) : <EOL> return open ( file_path , "<STR_LIT:rb>" ) . read ( ) <EOL> @ convert_local_errors <EOL> def write_file ( self , file_path , contents , mode = None , owner = None , <EOL> group = None ) : <EOL> f = open ( file_path , "<STR_LIT:wb>" if isinstance ( contents , bytes ) else "<STR_LIT:w>" ) <EOL> if mode is not None : <EOL> os . chmod ( file_path , mode ) <EOL> if ( owner is not None ) or ( group is not None ) : <EOL> file_stat = os . stat ( file_path ) <EOL> os . chown ( file_path , <EOL> owner if ( owner is not None ) else file_stat . st_uid , <EOL> group if ( group is not None ) else file_stat . st_gid ) <EOL> f . write ( contents ) <EOL> f . close ( ) <EOL> @ convert_local_errors <EOL> def execute_command ( self , cmd , pseudo_tty = False ) : <EOL> process = subprocess . Popen ( cmd , stdout = subprocess . PIPE , <EOL> stderr = subprocess . PIPE ) <EOL> fds = [ process . stdout , process . stderr ] <EOL> CHUNK = <NUM_LIT:2> ** <NUM_LIT:20> <EOL> while process . poll ( ) is None : <EOL> r , w , e = select . select ( fds , [ ] , fds ) <EOL> if process . stdout in r : <EOL> yield STDOUT , process . stdout . read ( CHUNK ) <EOL> if process . stderr in r : <EOL> yield STDERR , process . stderr . read ( CHUNK ) <EOL> yield DONE , process . returncode <EOL> @ convert_local_errors <EOL> def execute_shell ( self ) : <EOL> shell = os . environ . get ( "<STR_LIT>" ) <EOL> if not shell : <EOL> raise errors . RemoteError ( "<STR_LIT>" ) <EOL> return self . execute ( [ "<STR_LIT>" , shell ] ) <EOL> @ convert_local_errors <EOL> def stat ( self , file_path ) : <EOL> return os . stat ( file_path ) <EOL> @ convert_local_errors <EOL> def utime ( self , file_path , times ) : <EOL> os . utime ( file_path , times ) <EOL> class LocalTarControl ( RemoteControl ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , node , tar_dir ) : <EOL> self . tar_dir = tar_dir <EOL> RemoteControl . __init__ ( self , node ) <EOL> @ convert_local_errors <EOL> def put_file ( self , source_path , dest_path , callback = None ) : <EOL> self . write_file ( dest_path , open ( source_path , "<STR_LIT:rb>" ) . read ( ) ) <EOL> @ convert_local_errors <EOL> def makedirs ( self , dir_path ) : <EOL> pass <EOL> @ convert_local_errors <EOL> def read_file ( self , file_path ) : <EOL> raise IOError ( errno . ENOENT , "<STR_LIT>" ) <EOL> @ convert_local_errors <EOL> def write_file ( self , file_path , contents , mode = None , owner = None , <EOL> group = None ) : <EOL> full_path = os . path . join ( self . tar_dir , self . node [ "<STR_LIT:host>" ] , "<STR_LIT>" ) <EOL> if not os . path . isdir ( os . path . dirname ( full_path ) ) : <EOL> os . makedirs ( os . path . dirname ( full_path ) ) <EOL> file_obj = BytesIO ( contents ) <EOL> with tarfile . open ( full_path , "<STR_LIT:a>" ) as output : <EOL> info = tarfile . TarInfo ( file_path ) <EOL> info . size = len ( contents ) <EOL> if mode is not None : <EOL> info . mode = mode <EOL> if owner is not None : <EOL> info . uid = owner <EOL> if group is not None : <EOL> info . gid = group <EOL> info . mtime = time . time ( ) <EOL> output . addfile ( info , file_obj ) <EOL> @ convert_local_errors <EOL> def execute_command ( self , cmd , pseudo_tty = False ) : <EOL> full_path = os . path . join ( self . tar_dir , self . node [ "<STR_LIT:host>" ] , "<STR_LIT>" ) <EOL> if not os . path . isdir ( os . path . dirname ( full_path ) ) : <EOL> os . makedirs ( os . path . dirname ( full_path ) ) <EOL> with open ( full_path , "<STR_LIT:a>" ) as output : <EOL> output . write ( cmd + "<STR_LIT:\n>" ) <EOL> yield DONE , <NUM_LIT:0> <EOL> @ convert_local_errors <EOL> def execute_shell ( self ) : <EOL> assert <NUM_LIT:0> , "<STR_LIT>" <EOL> @ convert_local_errors <EOL> def stat ( self , file_path ) : <EOL> raise IOError ( errno . ENOENT , "<STR_LIT>" ) <EOL> @ convert_local_errors <EOL> def utime ( self , file_path , times ) : <EOL> pass <EOL> class SshRemoteControl ( RemoteControl ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , node ) : <EOL> self . log = logging . getLogger ( "<STR_LIT>" ) <EOL> RemoteControl . __init__ ( self , node ) <EOL> self . key_filename = None <EOL> cloud_prop = node . get ( "<STR_LIT>" , { } ) <EOL> cloud_key = cloud_prop . get ( "<STR_LIT>" , cloud_prop . get ( "<STR_LIT>" ) ) <EOL> if cloud_key : <EOL> self . key_filename = "<STR_LIT>" % cloud_key <EOL> else : <EOL> self . key_filename = node . get_tree_property ( "<STR_LIT>" ) <EOL> self . connect_timeout = node . get_tree_property ( "<STR_LIT>" , <NUM_LIT> ) </s>
