vm-manager/VMHelper.py

import json
import subprocess
import os, errno
import threading
import time
import sys
import pwd
from qmp import QEMUMonitorProtocol

class VMHelper:
	def __init__(self, filename):
		self.config = json.load(open(filename))

	def getVmIds(self):
		if ('VMs' in self.config):
			return [x for x , y in self.config['VMs'].items()]
		else:
			return []

	def humanMonitorCommand(self, vmid: str, command: str, args: list) -> str:
		proto = QEMUMonitorProtocol(self.config['kvm']['qmpsocket'].replace("$VMID", vmid))
		proto.connect()
		commandline = command + " " + " ".join(args)
		result = proto.cmd("human-monitor-command", { "command-line" : commandline })

		proto.close()
		return result["return"]

	def getPid(self,vmid) -> "int or None":
		result = None
		try:
			with open(self.config['kvm']['pidfile'].replace("$VMID", vmid)) as pidfd:
				firstline = pidfd.readline().strip()
				result = int(firstline) if firstline.isdigit() else None
		except IOError:
			pass
		return result
				

	def process_running(self,vmid) -> bool:

		pid = self.getPid(vmid)
		
		if pid is None:	
			return False

		try:
			os.kill(pid,0)
		except OSError as e:
			return e.errno != errno.ESRCH
		else:
			return True

	def getQMPStatus(self,vmid):
		# cmd doku http://git.qemu.org/?p=qemu.git;a=blob;f=qmp-commands.hx;h=1e0e11ee32571209e2dfce41b5c18f01d6ad3880;hb=HEAD
		proto = QEMUMonitorProtocol(self.config['kvm']['qmpsocket'].replace("$VMID", vmid))
		proto.connect()
	
		result = {}
		
		queryKvm = proto.command("query-kvm")
		result["kvm-enabled"] = queryKvm["enabled"]
		result["kvm-present"] = queryKvm["present"]

		result["status"] = proto.command("query-status")["status"]

		proto.close()
		return result

	def autostartVMs(self,managerpath):
		if ('VMs' in self.config):
			for vmid, vmcfg in self.config['VMs'].items():
				if "autostart" in vmcfg and vmcfg["autostart"] and not self.process_running(vmid):
					self.startVM(vmid,managerpath)
		else:
			raise Exception("Missing VMs config section!")

	def startVM(self, vmid, managerpath):
		self.setupNetwork(vmid)
		cmd = []

		fullpath_helper = os.path.realpath(__file__)
		pathname_helper = os.path.dirname(fullpath_helper)

		cmd.append(pathname_helper + "/qemu-wrapper.sh")
		cmd.append(sys.executable)
		cmd.append(managerpath)
		cmd.append(self.config['kvm']['executable'])
		cmd.append(vmid)
		
		cmd.append("-pidfile")
		cmd.append(self.config['kvm']['pidfile'].replace("$VMID", vmid))

		cmd.append("-qmp")
		cmd.append("unix:" + self.config['kvm']['qmpsocket'].replace("$VMID", vmid) + ",server,nowait")

		cmd += ["-name", vmid]

		default_args = self.config['kvm']['default_args'].replace("$VMID", vmid)
		cmd += default_args.split()
		cmd += self.createArguments(vmid).split()
		#print(" ".join(cmd))
		subprocess.Popen(cmd, stdout=open("/dev/null"), stderr=open("/dev/null"))	
		#subprocess.call(cmd

	def shutdownVMs(self,timeout,parallel=True, statusCallback=lambda vmid, st : None):
		threads = []
		for vmid in self.getVmIds():
			if self.process_running(vmid):
				# i=vmid: strange workaround for strange problem the it uses the wrong vmid...
				def stopCallback(st, i=vmid): statusCallback(i, st)

				thread = threading.Thread(target=lambda : self.stopVM(vmid, timeout, stopCallback))
				thread.start()
				if parallel:
					threads.append(thread)
				else:
					thread.join()

		for thread in threads:
			thread.join()
	
	def stopVM(self, vmid, timeout=None, statusCallback=lambda st : None, wait=False):
		proto = QEMUMonitorProtocol(self.config['kvm']['qmpsocket'].replace("$VMID", vmid))
		proto.connect()
		
		statusCallback("send_powerdown")
		proto.cmd("system_powerdown")
		
		if timeout is None and not wait:
			proto.close()
			return

		timeoutEvent = threading.Event()

		def waitForShutdown():
			shutDown = False
			while not timeoutEvent.is_set() and not shutDown :
				event = proto.pull_event(wait=True)
				shutDown = event is not None and event["event"] == 'SHUTDOWN'
			
		eventWaitThread = threading.Thread(target=waitForShutdown)
		eventWaitThread.start()
	
		if(timeout is None and wait):
			eventWaitThread.join()
			proto.close()
			return

		eventWaitThread.join(timeout)
		timeoutEvent.set()
		if eventWaitThread.is_alive():
			statusCallback("send_quit")
			proto.cmd("quit")
			time.sleep(1)
			if(self.process_running(vmid)):
				statusCallback("kill_vm")
				self.killVm(vmid, 9) #kill it with fire!

		proto.close()

	def killVm(self, vmid, signal=15):
		pid = self.getPid(vmid)
		
		if pid is not None:	
			os.kill(pid,signal)

	
	def createArguments(self, vmid):
		if ('VMs' in self.config) and (vmid in self.config['VMs']):
			config = self.config['VMs'][vmid]
		else:
			raise Exception("No such VM configuration")

		args = ""
		if ('cpu' in config):
			args += " -cpu " + config['cpu']
		if ('smp' in config):
			args += " -smp " + str(config['smp'])
		if ('memory' in config):
			args += " -m " + str(config['memory'])
		if ('disk' in config):
			disk = config['disk']
			if ('file' in disk):
				args += " -drive file=" + disk['file']
				if ('hw' in disk):
					args += ",if=" + disk['hw']
		if ('cdrom' in config):
			args += " -cdrom " + config['cdrom']
		if ('network' in config):
			net = config['network']
			args += " -net nic,vlan=1"
			if ('mac' in net):
				args += ",macaddr=" + net['mac']
			if ('hw' in net):
				args += ",model=" + net['hw']
			args += " -net tap,vlan=1"
			if ('dev' in net):
				args += ",ifname=" + net['dev']
			args += ",script=no,downscript=no"
		if ('vnc' in config):
			vnc = config['vnc']
			if ('display' in vnc):
				args += " -vnc 127.0.0.1:" + str(vnc['display'])
		if ('keyboard' in config):
			args += " -k " + config['keyboard']
		if ('kernel' in config):
			args += " -kernel " + config['kernel']
		if ('append' in config):
			args += " -append \"" + config['append'] + "\""
		return args

	def setupNetwork(self, vmid):
		if ('VMs' in self.config) and (vmid in self.config['VMs']):
			config = self.config['VMs'][vmid]
		else:
			raise Exception("No such VM configuration")
		
		commands = []
		if ('network' in config):
			net = config['network']
			if ('dev' in net):
				commands.append(["tunctl", "-t", net['dev']])
				commands.append(["ip", "link", "set", "dev", net['dev'], "up"])
				commands.append(["ip", "addr", "add", "dev", net['dev'], "0.0.0.0"])
				commands.append(["ip", "-6" , "addr", "add", "dev", net['dev'], "fe80::1/64"]) 

				chain = "VMNET-FWD-" + vmid.upper()

				commands.append(["iptables", "--new-chain", chain])
				if ('ip' in net):
					for ip in net['ip']:
						commands.append(["ip", "route", "add", ip, "dev", net['dev']])
						commands.append(["iptables", "-A", chain, "-i", net['dev'], "-s", ip, "-j", "ACCEPT"])
				commands.append(["iptables", "-A", chain, "-i", net['dev'], "-j", "REJECT", "--reject-with", "icmp-admin-prohibited"])
				commands.append(["iptables", "-A", "FORWARD", "-j", chain])
				
				commands.append(["ip6tables", "--new-chain", chain])
				if ('ipv6' in net):
					for ipv6 in net['ipv6']:
						commands.append( ["ip", "-6", "route", "add", ipv6, "dev", net['dev']])
						commands.append(["ip6tables", "-A", chain, "-i", net['dev'], "-s", ipv6, "-j", "ACCEPT"])
				commands.append(["ip6tables", "-A", chain, "-i", net['dev'], "-j", "REJECT", "--reject-with", "icmp6-adm-prohibited"])
				commands.append(["ip6tables", "-A", "FORWARD", "-j", chain])
		
		for cmd in commands:
			subprocess.call(cmd, stdout=open("/dev/null"))
	
	def teardownNetwork(self, vmid):
		if ('VMs' in self.config) and (vmid in self.config['VMs']):
			config = self.config['VMs'][vmid]
		else:
			raise Exception("No such VM configuration")
		
		commands = []
		if ('network' in config):
			net = config['network']
			if ('dev' in net):
				commands.append(["tunctl", "-d", net['dev']])

				chain = "VMNET-FWD-" + vmid.upper()

				commands.append(["iptables", "-F", chain])
				commands.append(["iptables", "-D", "FORWARD", "-j", chain])
				commands.append(["iptables", "--delete-chain", chain])
				
				commands.append(["ip6tables", "-F", chain])
				commands.append(["ip6tables", "-D", "FORWARD", "-j", chain])
				commands.append(["ip6tables", "--delete-chain", chain])
		
		for cmd in commands:
			subprocess.call(cmd, stdout=open("/dev/null"))

	def generateAuthorizedKeys(self):
		userkeys = {}
		keydir = os.path.join(self.config["ssh"]["homedir"], self.config["ssh"]["keydir"])
		for filename in os.listdir(keydir):
			fnsplit = filename.split("@")
			if len(fnsplit) == 2:
				user = fnsplit[0]
				with open(os.path.join(keydir,filename)) as f:
					userkeys[user] = userkeys[user].append(f.readline().rstrip('\n')) if user in userkeys else [f.readline().rstrip('\n')]
		authorized_keys = ""
		for user, keys in userkeys.items():
			prepend = 'no-agent-forwarding,no-user-rc,no-X11-forwarding,command="read",'
			for vm, vals in self.config["VMs"].items():
				if vals["owner"] in userkeys:
					prepend += 'permitopen="localhost:{0}",'.format(vals["vnc"]["display"] + 5900)
					prepend += 'permitopen="127.0.0.1:{0}",'.format(vals["vnc"]["display"] + 5900)
					prepend += 'permitopen="[::1]:{0}",'.format(vals["vnc"]["display"] + 5900)
			prepend = prepend[:-1]

			for key in keys:
				authorized_keys += prepend + " " + key + "\n"

		authorized_key_file = os.path.join(self.config["ssh"]["homedir"], ".ssh/authorized_keys")
		with open(authorized_key_file, mode="w") as f:
			f.write(authorized_keys)
		
		os.chmod(authorized_key_file ,0o600)
		os.chown(authorized_key_file, pwd.getpwnam(self.config["ssh"]["user"]).pw_uid, pwd.getpwnam(self.config["ssh"]["user"]).pw_gid)