#!/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)