#!/usr/bin/env python
import argparse
import collections
import datetime
import os
import signal
import subprocess
import sys
import fileinput

Instance = collections.namedtuple("Instance", "path name")

if sys.platform == "linux" or sys.platform == "linux2":
    vmrun_cmd = "vmrun"
    ovftool_cmd = "ovftool"
else:
    vmrun_cmd = "/Applications/VMware Fusion.app/Contents/Library/vmrun"
    ovftool_cmd = "/Applications/VMware OVF Tool/ovftool"
instances = []

# Argparse setup
parser = argparse.ArgumentParser(description="IncludeOS VMmware runner")
parser.add_argument("-n", "--instances", dest="instances",
	type=int, help="Number of instances to launch")
parser.add_argument("vm_filename", type=str, help="Filename of VM")
args = parser.parse_args()
if args.instances == None:
	args.instances = 1

def handler(signal, frame):
	stop_instances()
	sys.exit(0)

def stop_instances():
	for instance in instances:
		stop_instance(instance)

def start_instance(instance):
	serial_path = os.path.join(instance.path, "includeos-serial.txt")
	vmxpath = os.path.join(instance.path, instance.name)
	return subprocess.call([vmrun_cmd, "start", vmxpath, "nogui"])

def stop_instance(instance):
	full_path = os.path.join(instance.path, instance.name)
	print "Stopping vm {path}".format(path=full_path)
	subprocess.call([vmrun_cmd, "stop", full_path, "hard"])

def convert_image(servicename, destination_path):
	vmdkfile = "{filename}.vmdk".format(filename=servicename)
	vmdk_path = os.path.join(destination_path, vmdkfile)
	convert_cmd = "qemu-img convert -O vmdk {filename}.img {dest}".format(dest=vmdk_path, filename=servicename)
	try:
		os.makedirs(destination_path)
	except:
		pass
	return subprocess.call(convert_cmd, shell=True)

def write_vmx(instance_path, servicename):
	vmdkfile = "{filename}.vmdk".format(filename=servicename)
	vmdk_full_path = os.path.join(instance_path, vmdkfile)
	vmxfile = "{filename}.vmx".format(filename=servicename)
	vmx_full_path = os.path.join(instance_path, vmxfile)
	vmx = """.encoding = "UTF-8"
config.version = "8"
virtualHW.version = "11"
displayName = "IncludeOS"
guestOS = "other-64"
nvram = "IncludeOS.nvram"
memsize = "1024"
ide0:0.present = "TRUE"
ide0:0.fileName = "{path}"
ide1:0.present = "FALSE"
ethernet0.present = "TRUE"
ethernet0.connectionType = "nat"
ethernet0.wakeOnPcktRcv = "FALSE"
ethernet0.addressType = "generated"
ethernet0.linkStatePropagation.enable = "TRUE"
ethernet0.virtualDev = "vmxnet3"
sound.present = "FALSE"
pciBridge0.present = "TRUE"
pciBridge4.present = "TRUE"
pciBridge4.virtualDev = "pcieRootPort"
pciBridge4.functions = "8"
pciBridge5.present = "TRUE"
pciBridge5.virtualDev = "pcieRootPort"
pciBridge5.functions = "8"
pciBridge6.present = "TRUE"
pciBridge6.virtualDev = "pcieRootPort"
pciBridge6.functions = "8"
pciBridge7.present = "TRUE"
pciBridge7.virtualDev = "pcieRootPort"
pciBridge7.functions = "8"
vmci0.present = "FALSE"
vprobe.enable = "FALSE"
hpet0.present = "FALSE"
tools.syncTime = "TRUE"
virtualHW.productCompatibility = "hosted"
tools.upgrade.policy = "upgradeAtPowerCycle"
powerType.powerOff = "hard"
powerType.powerOn = "hard"
powerType.suspend = "hard"
powerType.reset = "hard"
extendedConfigFile = "IncludeOS.vmxf"
floppy0.present = "FALSE"
serial0.present = "TRUE"
serial0.fileType = "file"
serial0.fileName = "includeos-serial.txt"
serial0.tryNoRxLoss = "FALSE"
""".format(path=vmdk_full_path)
	try:
		os.makedirs(vmx_output_path)
	except:
		pass
	with open(vmx_full_path, "w") as f:
		f.write(vmx)
	instance = Instance(path=instance_path, name=vmxfile)
	return instance

def write_ova(instance_path, servicename):
    vmxfile = "{filename}.vmx".format(filename=servicename)
    vmx_full_path = os.path.join(instance_path, vmxfile)
    vmx_full_path_new = vmx_full_path + ".new"
    f1 = open(vmx_full_path, 'r')
    f2 = open(vmx_full_path_new, 'w')
    for line in f1:
        f2.write(line.replace('nat', 'IOS-Network'))
    f1.close()
    f2.close()
    os.rename(vmx_full_path_new, vmx_full_path)
    ovafile = "{filename}.ova".format(filename=servicename)
    ova_full_path = os.path.join(instance_path, ovafile)
    subprocess.call([ovftool_cmd, vmx_full_path, ova_full_path])
    return

if __name__ == "__main__":
	timestamp = "{:%Y%m%d_%H%M%S}".format(datetime.datetime.now())
	servicename = args.vm_filename
	path, _ = os.path.split(os.path.abspath(args.vm_filename))
	root_path = os.path.join(path, timestamp)
	for i in range(args.instances):
		instance_path = os.path.join(root_path, str(i))
		# VMWare wants an exclusive lock on the disk file, so every instance
		# needs their own copy :(
		res = convert_image(servicename, instance_path)
		if res == 0:
			instance = write_vmx(instance_path, servicename)
                        if os.path.isfile(ovftool_cmd):
                            write_ova(instance_path, servicename)
			res = start_instance(instance)
			if res == 0:
				instances.append(instance)
	if len(instances) != args.instances:
		print "One or more instances failed to launch, stopping all"
		stop_instances()
	else:
		signal.signal(signal.SIGINT, handler)
		tail_cmd = ["tail", "-f"]
		# Feel free to substitute your preferred tail variant
		#tail_cmd = ["multitail", "-s", "2"]
		for instance in instances:
			log_path = os.path.join(instance.path, "includeos-serial.txt")
			tail_cmd.append(log_path)
		subprocess.call(tail_cmd)
