ros-exec-script.py (8659B)
1 #!/usr/bin/env python 2 3 import os,time,base64,json,socket,select,errno,sys 4 # FIXME: this script is adapted from 5 # https://wiki.mikrotik.com/wiki/Manual:CHR#Provisioning 6 7 # I don't know if it is freely usable/redistributable 8 9 class GuestAgent(object): 10 ''' 11 Qemu guest agent interface 12 runScript and runFile commands are tailored for ROS agent implementation 13 Transport provided by derived classes (transact method) 14 ''' 15 16 def __init__(self,**kwargs): 17 # Due to file contents being passed as base64 inside json: 18 # - large chunk sizes may slow down guest-side parsing. 19 # - small chunk sizes result in additional message fragmentation overhead. 20 # Default value is a guestimate. 21 self.__chunkSize = kwargs.get('chunkSize', 4096) 22 23 def _qmpError(self,cls,msg): 24 ''' Generic callback to log qmp errors before (optionally) raising an exception ''' 25 print(cls) 26 for line in msg.split('\n'): 27 print(line) 28 # raise RuntimeError() 29 30 def _error(self,msg,*a): 31 ''' Generic callback to misc errors before (optionally) raising an exception ''' 32 print(msg.format(*a)) 33 # raise RuntimeError() 34 35 def _info(self,msg,*a): 36 ''' Generic callback to log info ''' 37 print(msg.format(*a)) 38 39 def _monitorJob(self,pid): 40 ''' Block untill script job completes, echo output. Returns None on failure ''' 41 ret = self.transact('guest-exec-status',{'pid':pid}) 42 if ret is None: 43 return None 44 45 while not bool(ret['exited']): 46 time.sleep(1) 47 ret = self.transact('guest-exec-status',{'pid':pid}) 48 if ret is None: 49 return None 50 51 # err-data is never sent 52 out = [] 53 if 'out-data' in ret.keys(): 54 out = base64.b64decode(ret['out-data']).decode('utf-8').split('\n') 55 if not out[-1]: 56 out = out[:-1] 57 58 exitcode = int(ret['exitcode']) 59 return exitcode, out 60 61 def putFile(self,src,dst): 62 ''' Upload file ''' 63 src = os.path.expanduser(src) 64 if not os.path.exists(src) or not os.path.isfile(src): 65 self._error('File does not exist: \'{}\'', src) 66 return None 67 68 ret = self.transact('guest-file-open', {'path':dst,'mode':'w'}) 69 if ret is None: 70 return None 71 72 handle = int(ret) 73 74 file = open(src, 'rb') 75 for chunk in iter(lambda: file.read(self.__chunkSize), b''): 76 count = len(chunk) 77 chunk = base64.b64encode(chunk).decode('ascii') 78 79 ret = self.transact('guest-file-write',{'handle':handle,'buf-b64':chunk,'count':count}) 80 if ret is None: 81 return None 82 self.transact('guest-file-flush',{'handle':handle}) 83 ret = self.transact('guest-file-close',{'handle':handle}) 84 return True 85 86 def getFile(self,src,dst): 87 ''' Download file ''' 88 dst = os.path.expanduser(dst) 89 90 ret = self.transact('guest-file-open',{'path':src,'mode':'rb'}) 91 if ret is None: 92 return None 93 94 handle = int(ret) 95 data = '' 96 size = 0 97 98 while True: 99 ret = self.transact('guest-file-read',{'handle':handle,'count':self.__chunkSize}) 100 if ret is None: 101 return None 102 data += ret['buf-b64'] 103 size += int(ret['count']) 104 if bool(ret['eof']): 105 break 106 107 ret = self.transact('guest-file-close',{'handle':handle}) 108 data = base64.b64decode(data.encode('ascii')) 109 with open(dst,'wb') as f: 110 f.write(data) 111 return True 112 113 def ping(self): 114 ret = self.transact('guest-ping',{}) 115 if ret is None: 116 return None 117 return ret 118 119 def runFile(self,fileName): 120 ''' Execute file (on guest) as script ''' 121 ret = self.transact('guest-exec',{'path':fileName, 'capture-output':True}) 122 if ret is None: 123 return None 124 125 pid = ret['pid'] 126 return self._monitorJob(pid) 127 128 def runSource(self,cmd): 129 ''' Execute script ''' 130 if isinstance(cmd,list): 131 cmd = '\n'.join(cmd) 132 cmd += '\n' 133 cmd = base64.b64encode(cmd.encode('utf-8')).decode('ascii') 134 135 ret = self.transact('guest-exec',{'input-data':cmd, 'capture-output':True}) 136 if ret is None: 137 return None 138 139 pid = ret['pid'] 140 return self._monitorJob(pid) 141 142 def shutdown(self,mode='powerdown'): 143 ''' 144 Execut shutdown command 145 mode == 'reboot' - reboot guest 146 mode == 'shutdown' or mode == 'halt' - shutdown guest 147 ''' 148 ret = self.transact('guest-shutdown',{'mode':mode}) 149 return ret 150 151 class SocketAgent(GuestAgent): 152 ''' 153 GuestAgent using unix/tcp sockets for communication. 154 ''' 155 def __init__(self): 156 GuestAgent.__init__(self,chunkSize= 32 * 65536) 157 158 @staticmethod 159 def unix(dev): 160 ''' Connect using unix socket ''' 161 self = SocketAgent() 162 self.__af = socket.AF_UNIX 163 self.__args = dev 164 self.__wait = True 165 return self 166 167 @staticmethod 168 def tcp(ip,port,wait = True): 169 ''' Connect using tcp socket ''' 170 self = SocketAgent() 171 self.__af = socket.AF_INET 172 self.__args = (ip,port) 173 self.__wait = wait 174 return self 175 176 def __enter__(self): 177 self._sock = socket.socket(self.__af, socket.SOCK_STREAM) 178 if self.__wait: 179 self._info('Waiting for guest ...') 180 # Wait for hyper to create channel 181 while True: 182 try: 183 self._sock.connect(self.__args) 184 break 185 except socket.error as e: 186 print("error connecting", e) 187 if e.errno == errno.EHOSTUNREACH or e.errno == errno.ECONNREFUSED: 188 time.sleep(1) 189 else: 190 self._sock.close() 191 raise 192 193 #Wait for guest agent to initialize and sync 194 while True: 195 import random 196 key = random.randint(0, 0xffffffff) 197 msg = json.dumps({'execute':'guest-sync-delimited','arguments':{'id':key}},separators=(',',':'),sort_keys=True) 198 self._sock.send(msg.encode('ascii')) 199 200 self._sock.setblocking(0) 201 response = b'' 202 if (select.select([self._sock],[],[])[0]): 203 response += self._sock.recv(65536) 204 else: 205 raise RuntimeError() 206 self._sock.setblocking(1) 207 208 sentinel = b'\xff' 209 response = response.split(sentinel)[-1] 210 if not response: 211 time.sleep(3) 212 continue 213 response = json.loads(response.decode('utf-8').strip()) 214 if 'return' in response.keys(): 215 if int(response['return']) == key: 216 break 217 time.sleep(3) 218 else: 219 self._sock.connect(self.__args) 220 221 return self 222 223 def __exit__(self,*a): 224 self._sock.close() 225 226 def transact(self,cmd,args={}): 227 ''' Exchange a single command with guest agent ''' 228 timeout = 2 229 msg = json.dumps({'execute':cmd,'arguments':args},separators=(',',':'),sort_keys=True) 230 self._sock.send(msg.encode('ascii')) 231 self._sock.setblocking(0) 232 response = b'' 233 if (select.select([self._sock],[],[],timeout)[0]): 234 response += self._sock.recv(65536) 235 self._sock.setblocking(1) 236 if not response: 237 response = None 238 else: 239 if response[0] == 255: # sync 240 response = response[1:] 241 print(response.decode('utf-8').strip()) 242 response = json.loads(response.decode('utf-8').strip()) 243 if 'error' in response.keys(): 244 self._qmpError(response['error']['class'],response['error']['desc']) 245 response = None 246 elif 'return' in response: 247 response = response['return'] 248 return response 249 250 #------------------------------------------------------------------------------- 251 252 if __name__ == '__main__': 253 socketpath,filename=sys.argv[1:] 254 script = open(filename,"r").readlines() 255 256 with SocketAgent.unix(socketpath) as agent: 257 ret,out = agent.runSource(script) 258 print('ret = {}'.format(ret)) 259 for line in out: 260 print(line)