#!/usr/bin/env python import os,time,base64,json,socket,select,errno,sys # FIXME: this script is adapted from # https://wiki.mikrotik.com/wiki/Manual:CHR#Provisioning # I don't know if it is freely usable/redistributable class GuestAgent(object): ''' Qemu guest agent interface runScript and runFile commands are tailored for ROS agent implementation Transport provided by derived classes (transact method) ''' def __init__(self,**kwargs): # Due to file contents being passed as base64 inside json: # - large chunk sizes may slow down guest-side parsing. # - small chunk sizes result in additional message fragmentation overhead. # Default value is a guestimate. self.__chunkSize = kwargs.get('chunkSize', 4096) def _qmpError(self,cls,msg): ''' Generic callback to log qmp errors before (optionally) raising an exception ''' print(cls) for line in msg.split('\n'): print(line) # raise RuntimeError() def _error(self,msg,*a): ''' Generic callback to misc errors before (optionally) raising an exception ''' print(msg.format(*a)) # raise RuntimeError() def _info(self,msg,*a): ''' Generic callback to log info ''' print(msg.format(*a)) def _monitorJob(self,pid): ''' Block untill script job completes, echo output. Returns None on failure ''' ret = self.transact('guest-exec-status',{'pid':pid}) if ret is None: return None while not bool(ret['exited']): time.sleep(1) ret = self.transact('guest-exec-status',{'pid':pid}) if ret is None: return None # err-data is never sent out = [] if 'out-data' in ret.keys(): out = base64.b64decode(ret['out-data']).decode('utf-8').split('\n') if not out[-1]: out = out[:-1] exitcode = int(ret['exitcode']) return exitcode, out def putFile(self,src,dst): ''' Upload file ''' src = os.path.expanduser(src) if not os.path.exists(src) or not os.path.isfile(src): self._error('File does not exist: \'{}\'', src) return None ret = self.transact('guest-file-open', {'path':dst,'mode':'w'}) if ret is None: return None handle = int(ret) file = open(src, 'rb') for chunk in iter(lambda: file.read(self.__chunkSize), b''): count = len(chunk) chunk = base64.b64encode(chunk).decode('ascii') ret = self.transact('guest-file-write',{'handle':handle,'buf-b64':chunk,'count':count}) if ret is None: return None self.transact('guest-file-flush',{'handle':handle}) ret = self.transact('guest-file-close',{'handle':handle}) return True def getFile(self,src,dst): ''' Download file ''' dst = os.path.expanduser(dst) ret = self.transact('guest-file-open',{'path':src,'mode':'rb'}) if ret is None: return None handle = int(ret) data = '' size = 0 while True: ret = self.transact('guest-file-read',{'handle':handle,'count':self.__chunkSize}) if ret is None: return None data += ret['buf-b64'] size += int(ret['count']) if bool(ret['eof']): break ret = self.transact('guest-file-close',{'handle':handle}) data = base64.b64decode(data.encode('ascii')) with open(dst,'wb') as f: f.write(data) return True def ping(self): ret = self.transact('guest-ping',{}) if ret is None: return None return ret def runFile(self,fileName): ''' Execute file (on guest) as script ''' ret = self.transact('guest-exec',{'path':fileName, 'capture-output':True}) if ret is None: return None pid = ret['pid'] return self._monitorJob(pid) def runSource(self,cmd): ''' Execute script ''' if isinstance(cmd,list): cmd = '\n'.join(cmd) cmd += '\n' cmd = base64.b64encode(cmd.encode('utf-8')).decode('ascii') ret = self.transact('guest-exec',{'input-data':cmd, 'capture-output':True}) if ret is None: return None pid = ret['pid'] return self._monitorJob(pid) def shutdown(self,mode='powerdown'): ''' Execut shutdown command mode == 'reboot' - reboot guest mode == 'shutdown' or mode == 'halt' - shutdown guest ''' ret = self.transact('guest-shutdown',{'mode':mode}) return ret class SocketAgent(GuestAgent): ''' GuestAgent using unix/tcp sockets for communication. ''' def __init__(self): GuestAgent.__init__(self,chunkSize= 32 * 65536) @staticmethod def unix(dev): ''' Connect using unix socket ''' self = SocketAgent() self.__af = socket.AF_UNIX self.__args = dev self.__wait = True return self @staticmethod def tcp(ip,port,wait = True): ''' Connect using tcp socket ''' self = SocketAgent() self.__af = socket.AF_INET self.__args = (ip,port) self.__wait = wait return self def __enter__(self): self._sock = socket.socket(self.__af, socket.SOCK_STREAM) if self.__wait: self._info('Waiting for guest ...') # Wait for hyper to create channel while True: try: self._sock.connect(self.__args) break except socket.error as e: print("error connecting", e) if e.errno == errno.EHOSTUNREACH or e.errno == errno.ECONNREFUSED: time.sleep(1) else: self._sock.close() raise #Wait for guest agent to initialize and sync while True: import random key = random.randint(0, 0xffffffff) msg = json.dumps({'execute':'guest-sync-delimited','arguments':{'id':key}},separators=(',',':'),sort_keys=True) self._sock.send(msg.encode('ascii')) self._sock.setblocking(0) response = b'' if (select.select([self._sock],[],[])[0]): response += self._sock.recv(65536) else: raise RuntimeError() self._sock.setblocking(1) sentinel = b'\xff' response = response.split(sentinel)[-1] if not response: time.sleep(3) continue response = json.loads(response.decode('utf-8').strip()) if 'return' in response.keys(): if int(response['return']) == key: break time.sleep(3) else: self._sock.connect(self.__args) return self def __exit__(self,*a): self._sock.close() def transact(self,cmd,args={}): ''' Exchange a single command with guest agent ''' timeout = 2 msg = json.dumps({'execute':cmd,'arguments':args},separators=(',',':'),sort_keys=True) self._sock.send(msg.encode('ascii')) self._sock.setblocking(0) response = b'' if (select.select([self._sock],[],[],timeout)[0]): response += self._sock.recv(65536) self._sock.setblocking(1) if not response: response = None else: if response[0] == 255: # sync response = response[1:] print(response.decode('utf-8').strip()) response = json.loads(response.decode('utf-8').strip()) if 'error' in response.keys(): self._qmpError(response['error']['class'],response['error']['desc']) response = None elif 'return' in response: response = response['return'] return response #------------------------------------------------------------------------------- if __name__ == '__main__': socketpath,filename=sys.argv[1:] script = open(filename,"r").readlines() with SocketAgent.unix(socketpath) as agent: ret,out = agent.runSource(script) print('ret = {}'.format(ret)) for line in out: print(line)