<s> import os <EOL> import stat <EOL> import logging <EOL> import subprocess <EOL> import re <EOL> import threading <EOL> import tempfile <EOL> import shutil <EOL> from pexpect import EOF , TIMEOUT , spawn , pxssh <EOL> from wlauto . exceptions import HostError , DeviceError , TimeoutError , ConfigError <EOL> from wlauto . utils . misc import which , strip_bash_colors , escape_single_quotes , check_output <EOL> ssh = None <EOL> scp = None <EOL> sshpass = None <EOL> logger = logging . getLogger ( '<STR_LIT>' ) <EOL> def ssh_get_shell ( host , username , password = None , keyfile = None , port = None , timeout = <NUM_LIT:10> , telnet = False ) : <EOL> _check_env ( ) <EOL> if telnet : <EOL> if keyfile : <EOL> raise ConfigError ( '<STR_LIT>' ) <EOL> conn = TelnetConnection ( ) <EOL> else : <EOL> conn = pxssh . pxssh ( ) <EOL> try : <EOL> if keyfile : <EOL> conn . login ( host , username , ssh_key = keyfile , port = port , login_timeout = timeout ) <EOL> else : <EOL> conn . login ( host , username , password , port = port , login_timeout = timeout ) <EOL> except EOF : <EOL> raise DeviceError ( '<STR_LIT>' . format ( host ) ) <EOL> return conn <EOL> class TelnetConnection ( pxssh . pxssh ) : <EOL> def login ( self , server , username , password = '<STR_LIT>' , original_prompt = r'<STR_LIT>' , login_timeout = <NUM_LIT:10> , <EOL> auto_prompt_reset = True , sync_multiplier = <NUM_LIT:1> , port = <NUM_LIT> ) : <EOL> cmd = '<STR_LIT>' . format ( username , server , port ) <EOL> spawn . _spawn ( self , cmd ) <EOL> try : <EOL> i = self . expect ( '<STR_LIT>' , timeout = login_timeout ) <EOL> if i == <NUM_LIT:0> : <EOL> self . sendline ( password ) <EOL> i = self . expect ( [ original_prompt , '<STR_LIT>' ] , timeout = login_timeout ) <EOL> if i : <EOL> raise pxssh . ExceptionPxssh ( '<STR_LIT>' ) <EOL> except TIMEOUT : <EOL> if not password : <EOL> pass <EOL> else : <EOL> raise pxssh . ExceptionPxssh ( '<STR_LIT>' ) <EOL> if not self . sync_original_prompt ( sync_multiplier ) : <EOL> self . close ( ) <EOL> raise pxssh . ExceptionPxssh ( '<STR_LIT>' ) <EOL> if auto_prompt_reset : <EOL> if not self . set_unique_prompt ( ) : <EOL> self . close ( ) <EOL> message = '<STR_LIT>' <EOL> raise pxssh . ExceptionPxssh ( message . format ( self . before , self . PROMPT ) ) <EOL> return True <EOL> def check_keyfile ( keyfile ) : <EOL> """<STR_LIT>""" <EOL> desired_mask = stat . S_IWUSR | stat . S_IRUSR <EOL> actual_mask = os . stat ( keyfile ) . st_mode & <NUM_LIT> <EOL> if actual_mask != desired_mask : <EOL> tmp_file = os . path . join ( tempfile . gettempdir ( ) , os . path . basename ( keyfile ) ) <EOL> shutil . copy ( keyfile , tmp_file ) <EOL> os . chmod ( tmp_file , desired_mask ) <EOL> return tmp_file <EOL> else : <EOL> return keyfile <EOL> class SshShell ( object ) : <EOL> default_password_prompt = '<STR_LIT>' <EOL> max_cancel_attempts = <NUM_LIT:5> <EOL> def __init__ ( self , password_prompt = None , timeout = <NUM_LIT:10> , telnet = False ) : <EOL> self . password_prompt = password_prompt if password_prompt is not None else self . default_password_prompt <EOL> self . timeout = timeout <EOL> self . telnet = telnet <EOL> self . conn = None <EOL> self . lock = threading . Lock ( ) <EOL> self . connection_lost = False <EOL> def login ( self , host , username , password = None , keyfile = None , port = None , timeout = None ) : <EOL> logger . debug ( '<STR_LIT>' . format ( username , host ) ) <EOL> self . host = host <EOL> self . username = username <EOL> self . password = password <EOL> self . keyfile = check_keyfile ( keyfile ) if keyfile else keyfile <EOL> self . port = port <EOL> timeout = self . timeout if timeout is None else timeout <EOL> self . conn = ssh_get_shell ( host , username , password , self . keyfile , port , timeout , self . telnet ) <EOL> def push_file ( self , source , dest , timeout = <NUM_LIT:30> ) : <EOL> dest = '<STR_LIT>' . format ( self . username , self . host , dest ) <EOL> return self . _scp ( source , dest , timeout ) <EOL> def pull_file ( self , source , dest , timeout = <NUM_LIT:30> ) : <EOL> source = '<STR_LIT>' . format ( self . username , self . host , source ) <EOL> return self . _scp ( source , dest , timeout ) <EOL> def background ( self , command , stdout = subprocess . PIPE , stderr = subprocess . PIPE ) : <EOL> port_string = '<STR_LIT>' . format ( self . port ) if self . port else '<STR_LIT>' <EOL> keyfile_string = '<STR_LIT>' . format ( self . keyfile ) if self . keyfile else '<STR_LIT>' <EOL> command = '<STR_LIT>' . format ( ssh , keyfile_string , port_string , self . username , self . host , command ) <EOL> logger . debug ( command ) <EOL> if self . password : <EOL> command = _give_password ( self . password , command ) <EOL> return subprocess . Popen ( command , stdout = stdout , stderr = stderr , shell = True ) <EOL> def reconnect ( self ) : <EOL> self . conn = ssh_get_shell ( self . host , self . username , self . password , <EOL> self . keyfile , self . port , self . timeout , self . telnet ) <EOL> def execute ( self , command , timeout = None , check_exit_code = True , as_root = False , strip_colors = True ) : <EOL> try : <EOL> with self . lock : <EOL> if self . connection_lost : <EOL> logger . debug ( '<STR_LIT>' ) <EOL> self . reconnect ( ) <EOL> self . connection_lost = False <EOL> output = self . _execute_and_wait_for_prompt ( command , timeout , as_root , strip_colors ) <EOL> if check_exit_code : <EOL> exit_code_text = self . _execute_and_wait_for_prompt ( '<STR_LIT>' , strip_colors = strip_colors , log = False ) <EOL> try : <EOL> exit_code = int ( exit_code_text . split ( ) [ <NUM_LIT:0> ] ) <EOL> if exit_code : <EOL> message = '<STR_LIT>' <EOL> raise DeviceError ( message . format ( exit_code , command , output ) ) <EOL> except ( ValueError , IndexError ) : <EOL> logger . warning ( '<STR_LIT>' . format ( command , exit_code_text ) ) <EOL> return output <EOL> except EOF : <EOL> self . connection_lost = True <EOL> raise DeviceError ( '<STR_LIT>' ) <EOL> def logout ( self ) : <EOL> logger . debug ( '<STR_LIT>' . format ( self . username , self . host ) ) <EOL> self . conn . logout ( ) <EOL> def cancel_running_command ( self ) : <EOL> logger . debug ( '<STR_LIT>' ) <EOL> for _ in xrange ( self . max_cancel_attempts ) : <EOL> self . conn . sendline ( chr ( <NUM_LIT:3> ) ) <EOL> if self . conn . prompt ( <NUM_LIT:0.1> ) : <EOL> return True <EOL> return False <EOL> def _execute_and_wait_for_prompt ( self , command , timeout = None , as_root = False , strip_colors = True , log = True ) : <EOL> self . conn . prompt ( <NUM_LIT:0.1> ) <EOL> if as_root : <EOL> command = "<STR_LIT>" . format ( escape_single_quotes ( command ) ) <EOL> if log : <EOL> logger . debug ( command ) <EOL> self . conn . sendline ( command ) <EOL> index = self . conn . expect_exact ( [ self . password_prompt , TIMEOUT ] , timeout = <NUM_LIT:0.5> ) <EOL> if index == <NUM_LIT:0> : <EOL> self . conn . sendline ( self . password ) <EOL> timed_out = self . _wait_for_prompt ( timeout ) <EOL> output = re . sub ( r'<STR_LIT>' , r'<STR_LIT>' , self . conn . before ) <EOL> output = process_backspaces ( output ) <EOL> output = re . sub ( r'<STR_LIT>' . format ( re . escape ( command ) ) , '<STR_LIT>' , output , <NUM_LIT:1> ) . strip ( ) <EOL> else : <EOL> if log : <EOL> logger . debug ( command ) <EOL> self . conn . sendline ( command ) <EOL> timed_out = self . _wait_for_prompt ( timeout ) <EOL> output = re . sub ( r'<STR_LIT>' , r'<STR_LIT>' , self . conn . before ) <EOL> output = process_backspaces ( output ) <EOL> command_index = output . find ( command ) <EOL> output = output [ command_index + len ( command ) : ] . strip ( ) <EOL> if timed_out : <EOL> self . cancel_running_command ( ) <EOL> raise TimeoutError ( command , output ) <EOL> if strip_colors : <EOL> output = strip_bash_colors ( output ) <EOL> return output <EOL> def _wait_for_prompt ( self , timeout = None ) : <EOL> if timeout : <EOL> return not self . conn . prompt ( timeout ) <EOL> else : <EOL> while not self . conn . prompt ( self . timeout ) : <EOL> pass <EOL> return False <EOL> def _scp ( self , source , dest , timeout = <NUM_LIT:30> ) : <EOL> port_string = '<STR_LIT>' . format ( self . port ) if ( self . port and self . port != <NUM_LIT> ) else '<STR_LIT>' <EOL> keyfile_string = '<STR_LIT>' . format ( self . keyfile ) if self . keyfile else '<STR_LIT>' <EOL> command = '<STR_LIT>' . format ( scp , keyfile_string , port_string , source , dest ) <EOL> pass_string = '<STR_LIT>' <EOL> logger . debug ( command ) <EOL> if self . password : <EOL> command = _give_password ( self . password , command ) <EOL> try : <EOL> check_output ( command , timeout = timeout , shell = True ) <EOL> except subprocess . CalledProcessError as e : <EOL> raise subprocess . CalledProcessError ( e . returncode , e . cmd . replace ( pass_string , '<STR_LIT>' ) , e . output ) <EOL> except TimeoutError as e : <EOL> raise TimeoutError ( e . command . replace ( pass_string , '<STR_LIT>' ) , e . output ) <EOL> def _give_password ( password , command ) : <EOL> if not sshpass : <EOL> raise HostError ( '<STR_LIT>' ) <EOL> pass_string = "<STR_LIT>" . format ( password ) <EOL> return pass_string + command <EOL> def _check_env ( ) : <EOL> global ssh , scp , sshpass <EOL> if not ssh : <EOL> ssh = which ( '<STR_LIT>' ) <EOL> scp = which ( '<STR_LIT>' ) <EOL> sshpass = which ( '<STR_LIT>' ) <EOL> if not ( ssh and scp ) : <EOL> raise HostError ( '<STR_LIT>' ) <EOL> def process_backspaces ( text ) : <EOL> chars = [ ] <EOL> for c in text : <EOL> if c == chr ( <NUM_LIT:8> ) and chars : <EOL> chars . pop ( ) <EOL> else : <EOL> chars . append ( c ) <EOL> return '<STR_LIT>' . join ( chars ) </s>
