<s> from __future__ import division , unicode_literals , print_function <EOL> from os import path <EOL> from socket import ( <EOL> gaierror , <EOL> error as socket_error , timeout as timeout_error <EOL> ) <EOL> import six <EOL> import gevent <EOL> from termcolor import colored <EOL> from paramiko . agent import AgentRequestHandler <EOL> from paramiko import ( <EOL> SSHClient , SFTPClient , RSAKey , <EOL> MissingHostKeyPolicy , SSHException , AuthenticationException <EOL> ) <EOL> from pyinfra import logger <EOL> from pyinfra . api . util import read_buffer , make_command <EOL> def connect ( host , ** kwargs ) : <EOL> '''<STR_LIT>''' <EOL> logger . debug ( '<STR_LIT>' . format ( host . name , kwargs ) ) <EOL> name = host . name <EOL> hostname = host . data . ssh_hostname or name <EOL> try : <EOL> client = SSHClient ( ) <EOL> client . set_missing_host_key_policy ( MissingHostKeyPolicy ( ) ) <EOL> client . connect ( hostname , ** kwargs ) <EOL> session = client . get_transport ( ) . open_session ( ) <EOL> AgentRequestHandler ( session ) <EOL> logger . info ( '<STR_LIT>' . format ( <EOL> colored ( name , attrs = [ '<STR_LIT>' ] ) , <EOL> colored ( '<STR_LIT>' , '<STR_LIT>' ) <EOL> ) ) <EOL> return client <EOL> except AuthenticationException as e : <EOL> logger . error ( '<STR_LIT>' . format ( name , e ) ) <EOL> except SSHException as e : <EOL> logger . error ( '<STR_LIT>' . format ( name , e ) ) <EOL> except gaierror : <EOL> if hostname != name : <EOL> logger . error ( '<STR_LIT>' . format ( name , hostname ) ) <EOL> else : <EOL> logger . error ( '<STR_LIT>' . format ( name ) ) <EOL> except socket_error as e : <EOL> logger . error ( '<STR_LIT>' . format ( <EOL> name , kwargs . get ( '<STR_LIT:port>' , <NUM_LIT> ) , e ) <EOL> ) <EOL> except EOFError as e : <EOL> logger . error ( '<STR_LIT>' . format ( name , e ) ) <EOL> def connect_all ( state ) : <EOL> '''<STR_LIT>''' <EOL> greenlets = { } <EOL> for host in state . inventory : <EOL> kwargs = { <EOL> '<STR_LIT:username>' : host . data . ssh_user , <EOL> '<STR_LIT:port>' : host . data . ssh_port or <NUM_LIT> , <EOL> '<STR_LIT>' : state . config . TIMEOUT , <EOL> '<STR_LIT>' : False , <EOL> '<STR_LIT>' : False <EOL> } <EOL> if host . data . ssh_password : <EOL> kwargs [ '<STR_LIT:password>' ] = host . data . ssh_password <EOL> elif host . data . ssh_key : <EOL> ssh_key_filenames = [ <EOL> path . expanduser ( host . data . ssh_key ) <EOL> ] <EOL> if state . deploy_dir : <EOL> ssh_key_filenames . append ( <EOL> path . join ( state . deploy_dir , host . data . ssh_key ) <EOL> ) <EOL> for filename in ssh_key_filenames : <EOL> if path . isfile ( filename ) : <EOL> kwargs [ '<STR_LIT>' ] = RSAKey . from_private_key_file ( <EOL> filename = filename , <EOL> password = host . data . ssh_key_password <EOL> ) <EOL> break <EOL> else : <EOL> kwargs [ '<STR_LIT>' ] = True <EOL> kwargs [ '<STR_LIT>' ] = True <EOL> greenlets [ host . name ] = state . pool . spawn ( connect , host , ** kwargs ) <EOL> gevent . wait ( greenlets . values ( ) ) <EOL> failed_hosts = set ( ) <EOL> connected_hosts = set ( ) <EOL> for name , greenlet in six . iteritems ( greenlets ) : <EOL> client = greenlet . get ( ) <EOL> if not client : <EOL> failed_hosts . add ( name ) <EOL> else : <EOL> state . ssh_connections [ name ] = client <EOL> connected_hosts . add ( name ) <EOL> state . inventory . connected_hosts = connected_hosts <EOL> state . inventory . active_hosts = set ( greenlets . keys ( ) ) <EOL> state . fail_hosts ( failed_hosts ) <EOL> def run_shell_command ( <EOL> state , hostname , command , <EOL> sudo = False , sudo_user = None , env = None , timeout = None , print_output = False <EOL> ) : <EOL> '''<STR_LIT>''' <EOL> print_prefix = '<STR_LIT>' . format ( colored ( hostname , attrs = [ '<STR_LIT>' ] ) ) <EOL> if env is None : <EOL> env = { } <EOL> logger . debug ( '<STR_LIT>' . format ( hostname , command ) ) <EOL> logger . debug ( '<STR_LIT>' . format ( <EOL> sudo , sudo_user , env <EOL> ) ) <EOL> command = make_command ( command , env = env , sudo = sudo , sudo_user = sudo_user ) <EOL> if print_output : <EOL> print ( '<STR_LIT>' . format ( print_prefix , command ) ) <EOL> connection = state . ssh_connections [ hostname ] <EOL> _ , stdout_buffer , stderr_buffer = connection . exec_command ( command ) <EOL> channel = stdout_buffer . channel <EOL> stdout_reader = gevent . spawn ( <EOL> read_buffer , stdout_buffer , <EOL> print_output = print_output , <EOL> print_func = lambda line : '<STR_LIT>' . format ( print_prefix , line ) <EOL> ) <EOL> stderr_reader = gevent . spawn ( <EOL> read_buffer , stderr_buffer , <EOL> print_output = print_output , <EOL> print_func = lambda line : '<STR_LIT>' . format ( print_prefix , colored ( line , '<STR_LIT>' ) ) <EOL> ) <EOL> greenlets = gevent . wait ( ( stdout_reader , stderr_reader ) , timeout = timeout ) <EOL> if len ( greenlets ) != <NUM_LIT:2> : <EOL> stdout_reader . kill ( ) <EOL> stderr_reader . kill ( ) <EOL> raise timeout_error ( ) <EOL> stdout = stdout_reader . get ( ) <EOL> stderr = stderr_reader . get ( ) <EOL> return channel , stdout , stderr <EOL> def _get_sftp_connection ( state , hostname ) : <EOL> if hostname in state . sftp_connections : <EOL> return state . sftp_connections [ hostname ] <EOL> ssh_connection = state . ssh_connections [ hostname ] <EOL> transport = ssh_connection . get_transport ( ) <EOL> client = SFTPClient . from_transport ( transport ) <EOL> state . sftp_connections [ hostname ] = client <EOL> return client <EOL> def _put_file ( state , hostname , file_io , remote_location ) : <EOL> sftp = _get_sftp_connection ( state , hostname ) <EOL> sftp . putfo ( file_io , remote_location ) <EOL> def put_file ( <EOL> state , hostname , file_io , remote_file , <EOL> sudo = False , sudo_user = None , print_output = False <EOL> ) : <EOL> '''<STR_LIT>''' <EOL> print_prefix = '<STR_LIT>' . format ( colored ( hostname , attrs = [ '<STR_LIT>' ] ) ) <EOL> if not sudo : <EOL> _put_file ( state , hostname , file_io , remote_file ) <EOL> else : <EOL> temp_file = state . get_temp_filename ( remote_file ) <EOL> _put_file ( state , hostname , file_io , temp_file ) <EOL> command = '<STR_LIT>' . format ( temp_file , remote_file ) <EOL> if sudo_user : <EOL> command = '<STR_LIT>' . format ( command , sudo_user , remote_file ) <EOL> channel , _ , stderr = run_shell_command ( <EOL> state , hostname , command , <EOL> sudo = sudo , sudo_user = sudo_user , <EOL> print_output = print_output <EOL> ) <EOL> if channel . exit_status > <NUM_LIT:0> : <EOL> logger . error ( '<STR_LIT>' . format ( '<STR_LIT:\n>' . join ( stderr ) ) ) <EOL> return False <EOL> if print_output : <EOL> print ( '<STR_LIT>' . format ( print_prefix , remote_file ) ) </s>
