<s> import ast <EOL> import collections <EOL> import json <EOL> import subprocess <EOL> import select <EOL> import time <EOL> import os <EOL> import base64 <EOL> import zlib <EOL> import logging <EOL> from threading import Thread <EOL> from BroControl import py3bro <EOL> Queue = py3bro . Queue <EOL> Empty = py3bro . Empty <EOL> def get_muxer ( shell ) : <EOL> pythonpath = "<STR_LIT>" <EOL> muxer = r"""<STR_LIT>""" <EOL> if shell : <EOL> muxer = muxer . replace ( "<STR_LIT>" , "<STR_LIT>" ) <EOL> else : <EOL> muxer = muxer . replace ( "<STR_LIT>" , "<STR_LIT>" ) <EOL> if py3bro . using_py3 : <EOL> muxer = muxer . encode ( ) <EOL> muxer = base64 . b64encode ( zlib . compress ( muxer ) ) <EOL> if py3bro . using_py3 : <EOL> muxer = muxer . decode ( ) <EOL> muxer = "<STR_LIT>" % ( pythonpath , muxer ) <EOL> if py3bro . using_py3 : <EOL> muxer = muxer . encode ( ) <EOL> return muxer <EOL> CmdResult = collections . namedtuple ( "<STR_LIT>" , "<STR_LIT>" ) <EOL> class SSHMaster : <EOL> def __init__ ( self , host , localaddrs ) : <EOL> self . host = host <EOL> self . base_cmd = [ <EOL> "<STR_LIT>" , <EOL> "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" , <EOL> host , <EOL> ] <EOL> self . need_connect = True <EOL> self . master = None <EOL> self . localaddrs = localaddrs <EOL> self . run_mux = get_muxer ( False ) <EOL> self . run_mux_shell = get_muxer ( True ) <EOL> def connect ( self ) : <EOL> if self . need_connect : <EOL> if self . host in self . localaddrs : <EOL> cmd = [ "<STR_LIT>" ] <EOL> else : <EOL> cmd = self . base_cmd + [ "<STR_LIT>" ] <EOL> self . master = subprocess . Popen ( cmd , bufsize = <NUM_LIT:0> , stdout = subprocess . PIPE , stdin = subprocess . PIPE , close_fds = True , preexec_fn = os . setsid ) <EOL> self . need_connect = False <EOL> def readline_with_timeout ( self , timeout ) : <EOL> readable , _ , _ = select . select ( [ self . master . stdout ] , [ ] , [ ] , timeout ) <EOL> if not readable : <EOL> return False <EOL> jtxt = self . master . stdout . readline ( ) <EOL> if py3bro . using_py3 : <EOL> jtxt = jtxt . decode ( ) <EOL> return jtxt <EOL> def exec_command ( self , cmd , shell = False , timeout = <NUM_LIT> ) : <EOL> return self . exec_commands ( [ cmd ] , shell , timeout ) [ <NUM_LIT:0> ] <EOL> def exec_commands ( self , cmds , shell = False , timeout = <NUM_LIT> ) : <EOL> self . send_commands ( cmds , timeout , shell ) <EOL> return self . collect_results ( timeout ) <EOL> def send_commands ( self , cmds , timeout , shell = False ) : <EOL> self . connect ( ) <EOL> if shell : <EOL> self . master . stdin . write ( self . run_mux_shell ) <EOL> else : <EOL> self . master . stdin . write ( self . run_mux ) <EOL> self . master . stdin . flush ( ) <EOL> self . readline_with_timeout ( timeout ) <EOL> for cmd in cmds : <EOL> jcmd = "<STR_LIT>" % json . dumps ( cmd ) <EOL> if py3bro . using_py3 : <EOL> jcmd = jcmd . encode ( ) <EOL> self . master . stdin . write ( jcmd ) <EOL> self . master . stdin . write ( b"<STR_LIT>" ) <EOL> self . master . stdin . flush ( ) <EOL> self . sent_commands = len ( cmds ) <EOL> def collect_results ( self , timeout ) : <EOL> outputs = [ Exception ( "<STR_LIT>" ) ] * self . sent_commands <EOL> while True : <EOL> line = self . readline_with_timeout ( timeout ) <EOL> if not line : <EOL> logging . debug ( "<STR_LIT>" , self . host ) <EOL> self . close ( ) <EOL> break <EOL> resp = ast . literal_eval ( line ) <EOL> if resp == "<STR_LIT>" : <EOL> break <EOL> idx , result = resp <EOL> status , out , err = result <EOL> if py3bro . using_py3 : <EOL> out = out . decode ( errors = "<STR_LIT:replace>" ) <EOL> err = err . decode ( errors = "<STR_LIT:replace>" ) <EOL> outputs [ idx ] = CmdResult ( status , out , err ) <EOL> return outputs <EOL> def ping ( self , timeout = <NUM_LIT:10> ) : <EOL> output = self . exec_command ( [ "<STR_LIT>" , "<STR_LIT>" ] , timeout = timeout ) <EOL> return output and output . stdout . strip ( ) == "<STR_LIT>" <EOL> def close ( self ) : <EOL> if not self . master : <EOL> return <EOL> self . master . stdin . close ( ) <EOL> try : <EOL> self . master . kill ( ) <EOL> except OSError : <EOL> pass <EOL> self . master . wait ( ) <EOL> self . master = None <EOL> self . need_connect = True <EOL> __del__ = close <EOL> STOP_RUNNING = object ( ) <EOL> class HostHandler ( Thread ) : <EOL> def __init__ ( self , host , localaddrs , timeout ) : <EOL> self . host = host <EOL> self . localaddrs = localaddrs <EOL> self . timeout = timeout <EOL> self . q = Queue ( ) <EOL> self . alive = "<STR_LIT>" <EOL> self . master = None <EOL> Thread . __init__ ( self ) <EOL> def shutdown ( self ) : <EOL> self . q . put ( ( STOP_RUNNING , None , None ) ) <EOL> def connect ( self ) : <EOL> if self . master : <EOL> self . master . close ( ) <EOL> self . master = SSHMaster ( self . host , self . localaddrs ) <EOL> def ping ( self ) : <EOL> try : <EOL> return self . master . ping ( ) <EOL> except Exception as e : <EOL> logging . debug ( "<STR_LIT>" , self . host ) <EOL> return False <EOL> def connect_and_ping ( self ) : <EOL> if self . alive != True : <EOL> self . connect ( ) <EOL> self . alive = self . ping ( ) <EOL> def run ( self ) : <EOL> while True : <EOL> if self . iteration ( ) : <EOL> return <EOL> def iteration ( self ) : <EOL> try : <EOL> item , shell , rq = self . q . get ( timeout = <NUM_LIT:30> ) <EOL> except Empty : <EOL> self . connect_and_ping ( ) <EOL> return False <EOL> if item is STOP_RUNNING : <EOL> return True <EOL> self . connect_and_ping ( ) <EOL> if not self . alive : <EOL> resp = [ Exception ( "<STR_LIT>" % self . host ) ] * len ( item ) <EOL> rq . put ( resp ) <EOL> return False <EOL> try : <EOL> resp = self . master . exec_commands ( item , shell , self . timeout ) <EOL> except Exception as e : <EOL> self . alive = False <EOL> msg = "<STR_LIT>" % ( self . host , e ) <EOL> logging . debug ( msg ) <EOL> resp = [ Exception ( msg ) ] * len ( item ) <EOL> time . sleep ( <NUM_LIT:2> ) <EOL> rq . put ( resp ) <EOL> return False <EOL> def send_commands ( self , commands , shell , rq ) : <EOL> self . q . put ( ( commands , shell , rq ) ) <EOL> class MultiMasterManager : <EOL> def __init__ ( self , localaddrs = [ ] ) : <EOL> self . masters = { } <EOL> self . response_queues = { } <EOL> self . localaddrs = localaddrs <EOL> def setup ( self , host , timeout ) : <EOL> if host not in self . masters : <EOL> self . masters [ host ] = HostHandler ( host , self . localaddrs , timeout ) <EOL> self . masters [ host ] . start ( ) <EOL> def send_commands ( self , host , commands , timeout , shell = False ) : <EOL> self . setup ( host , timeout ) <EOL> rq = Queue ( ) <EOL> self . response_queues [ host ] = rq <EOL> self . masters [ host ] . send_commands ( commands , shell , rq ) <EOL> def get_result ( self , host , timeout ) : <EOL> rq = self . response_queues [ host ] <EOL> try : <EOL> return rq . get ( timeout = timeout ) <EOL> except Empty : <EOL> logging . debug ( "<STR_LIT>" , host ) <EOL> self . shutdown ( host ) <EOL> return [ Exception ( "<STR_LIT>" ) ] <EOL> def exec_command ( self , host , command , timeout = <NUM_LIT:30> ) : <EOL> return self . exec_commands ( host , [ command ] , timeout ) [ <NUM_LIT:0> ] <EOL> def exec_commands ( self , host , commands , timeout = <NUM_LIT> ) : <EOL> self . send_commands ( host , commands , timeout ) <EOL> return self . get_result ( host , timeout ) <EOL> def exec_multihost_commands ( self , cmds , shell = False , timeout = <NUM_LIT> ) : <EOL> hosts = collections . defaultdict ( list ) <EOL> for host , cmd in cmds : <EOL> hosts [ host ] . append ( cmd ) <EOL> for host , cmds in hosts . items ( ) : <EOL> self . send_commands ( host , cmds , timeout , shell ) <EOL> for host in hosts : <EOL> for res in self . get_result ( host , timeout ) : <EOL> yield host , res <EOL> def host_status ( self ) : <EOL> for h , o in self . masters . items ( ) : <EOL> if h not in self . localaddrs : <EOL> yield h , o . alive <EOL> def shutdown ( self , host ) : <EOL> self . masters [ host ] . shutdown ( ) <EOL> del self . masters [ host ] <EOL> def shutdown_all ( self ) : <EOL> for handler in self . masters . values ( ) : <EOL> handler . shutdown ( ) <EOL> self . masters = { } <EOL> __del__ = shutdown_all </s>
