#!{{ parameter_dict.get('python-path') }}
# BEWARE: This file is operated by slapgrid
# BEWARE: It will be overwritten automatically

import hashlib
import os
import socket
import subprocess
import urllib
import gzip
import shutil
from random import shuffle
import glob
import re

import ssl

# XXX: give all of this through parameter, don't use this as template, but as module
qemu_img_path = '{{ parameter_dict.get("qemu-img-path") }}'
qemu_path = '{{ parameter_dict.get("qemu-path") }}'
disk_size = '{{ parameter_dict.get("disk-size") }}'
disk_type = '{{ parameter_dict.get("disk-type") }}'
disk_format = '{{ parameter_dict.get("disk-format", "qcow2")}}'
disk_format = disk_format \
              if disk_format in ['qcow2', 'raw', 'vdi', 'vmdk', 'cloop', 'qed'] else 'qcow2'

socket_path = '{{ parameter_dict.get("socket-path") }}'
nbd_list = (('{{ parameter_dict.get("nbd-host") }}',
              {{ parameter_dict.get("nbd-port") }}),
            ('{{ parameter_dict.get("nbd2-host") }}',
              {{ parameter_dict.get("nbd2-port") }}))
default_cdrom_iso = '{{ parameter_dict.get("default-cdrom-iso") }}'
disk_path = '{{ parameter_dict.get("disk-path") }}'

virtual_hard_drive_url = '{{ parameter_dict.get("virtual-hard-drive-url") }}'.strip()
virtual_hard_drive_md5sum = '{{ parameter_dict.get("virtual-hard-drive-md5sum") }}'.strip()
virtual_hard_drive_gzipped = '{{ parameter_dict.get("virtual-hard-drive-gzipped") }}'.strip().lower()
nat_rules = '{{ parameter_dict.get("nat-rules") }}'.strip()
use_tap = '{{ parameter_dict.get("use-tap") }}'.lower()
use_nat = '{{ parameter_dict.get("use-nat") }}'.lower()
set_nat_restrict = '{{ parameter_dict.get("nat-restrict") }}'.lower()
enable_vhost = '{{ parameter_dict.get("enable-vhost") }}'.lower()
tap_interface = '{{ parameter_dict.get("tap-interface") }}'
listen_ip = '{{ parameter_dict.get("ipv4") }}'
mac_address = '{{ parameter_dict.get("mac-address") }}'
tap_mac_address = '{{ parameter_dict.get("tap-mac-address") }}'
numa_list = '{{ parameter_dict.get("numa", "") }}'.split()
ram_size = {{ parameter_dict.get("ram-size") }}
ram_max_size = '{{ parameter_dict.get("ram-max-size") }}'
init_ram_size = {{ parameter_dict.get("init-ram-size") }}
pid_file_path = '{{ parameter_dict.get("pid-file-path") }}'
external_disk_number = {{ parameter_dict.get("external-disk-number") }}
external_disk_size = {{ parameter_dict.get("external-disk-size") }}
external_disk_format = '{{ parameter_dict.get("external-disk-format", "qcow2") }}'
external_disk_format = external_disk_format \
              if external_disk_format in ['qcow2', 'raw', 'vdi', 'vmdk', 'cloop', 'qed'] else 'qcow2'
disk_storage_dict = {}
disk_storage_list = """{{ parameter_dict.get("disk-storage-list") }}""".split('\n')
map_storage_list = []
etc_directory = '{{ parameter_dict.get("etc-directory") }}'.strip()
httpd_port = {{ parameter_dict.get("httpd-port") }}
netcat_bin = '{{ parameter_dict.get("netcat-binary") }}'.strip()
cluster_doc_host = '{{ parameter_dict.get("cluster-doc-host") }}'
cluster_doc_port = {{ parameter_dict.get("cluster-doc-port") }}
language = '{{ parameter_dict.get("language") }}'
language_list = ['ar', 'da', 'de', 'de-ch', 'en-gb', 'en-us', 'es', 'et', 'fi',
    'fo', 'fr', 'fr-be', 'fr-ca', 'fr-ch', 'hr', 'hu', 'is', 'it', 'ja', 'lt',
    'lv', 'mk', 'nl', 'nl-be', 'no', 'pl', 'pt', 'pt-br', 'ru', 'sl', 'sv',
    'th', 'tr']
url_check_certificate = '{{ parameter_dict.get("hard-drive-url-check-certificate", "true") }}'.lower()
auto_ballooning = '{{ parameter_dict.get("auto-ballooning") }}' in ('true', 'True', '1')
vm_name = '{{ parameter_dict.get("name") }}'
disk_cache = '{{ parameter_dict.get("disk-cache", "writeback") }}'.strip()
disk_cache = disk_cache if disk_cache in ["none", "writeback", "unsafe", 
              "directsync", "writethrough"] else "writeback"
disk_aio = '{{ parameter_dict.get("disk-aio", "threads") }}'.strip()
disk_aio = disk_aio if disk_aio in ["threads", "native"] else "threads"

# If a device (ie.: /dev/sdb) is provided, use it instead 
# the disk_path with disk_format
disk_device_path = '{{ parameter_dict.get("disk-device-path", "") }}'
if disk_device_path.startswith("/dev/"):
  disk_path = disk_device_path 
  disk_format = "raw" 
  disk_aio = "native"
  disk_cache = "none"

smp_count = {{ parameter_dict.get("smp-count") }}
smp_max_count = {{ parameter_dict.get("smp-max-count") }}
machine_options = '{{ parameter_dict.get("machine-options", "") }}'.strip()
cpu_model = '{{ parameter_dict.get("cpu-model") }}'.strip()

enable_device_hotplug = '{{ parameter_dict.get("enable-device-hotplug") }}'.lower()

logfile = '{{ parameter_dict.get("log-file") }}'

if hasattr(ssl, '_create_unverified_context') and url_check_certificate == 'false':
  opener = urllib.FancyURLopener(context=ssl._create_unverified_context())
else:
  opener = urllib.FancyURLopener({})

def md5Checksum(file_path):
    with open(file_path, 'rb') as fh:
        m = hashlib.md5()
        while True:
            data = fh.read(8192)
            if not data:
                break
            m.update(data)
        return m.hexdigest()

def getSocketStatus(host, port):
  s = None
  for res in socket.getaddrinfo(host, port,
      socket.AF_UNSPEC, socket.SOCK_STREAM):
    af, socktype, proto, canonname, sa = res
    try:
      s = socket.socket(af, socktype, proto)
    except socket.error, msg:
      s = None
      continue
    try:
      s.connect(sa)
    except socket.error, msg:
      s.close()
      s = None
      continue
    break
  return s

def getMapStorageList(disk_storage_dict, external_disk_number):
  map_disk_file = os.path.join(etc_directory, '.data-disk-ids')
  last_disk_num_f = os.path.join(etc_directory, '.data-disk-amount')
  id_list = []
  last_amount = 0
  map_f_exist = os.path.exists(map_disk_file)
  if os.path.exists(last_disk_num_f):
    with open(last_disk_num_f, 'r') as lf:
      last_amount = int(lf.readline())
  if map_f_exist:
    with open(map_disk_file, 'r') as mf:
      # ID are writen in one line: data1 data3 data2 ...
      content = mf.readline()
      for id in content.split(' '):
        if disk_storage_dict.has_key(id):
          id_list.append(id)
        else:
          # Mean that this disk path has been removed (disk unmounted)
          last_amount -= 1
  for key in disk_storage_dict:
    if not key in id_list:
      id_list.append(key)

  if id_list:
    if not map_f_exist:
      # shuffle the list to not write disk in data1, data2, ... everytime 
      shuffle(id_list)
    if external_disk_number < last_amount:
      # Drop created disk is not allowed
      external_disk_number = last_amount
    with open(map_disk_file, 'w') as mf:
      mf.write(' '.join(id_list))
    with open(last_disk_num_f, 'w') as lf:
      lf.write('%s' % external_disk_number)
  return id_list, external_disk_number

# Download existing hard drive if needed at first boot
if not os.path.exists(disk_path) and virtual_hard_drive_url != '':
  print('Downloading virtual hard drive...')
  try:
    downloaded_disk = disk_path
    if virtual_hard_drive_gzipped == 'true':
      downloaded_disk = '%s.gz' % disk_path
    opener.retrieve(virtual_hard_drive_url, downloaded_disk)
  except:
    if os.path.exists(downloaded_disk):
      os.remove(downloaded_disk)
    raise
  md5sum = virtual_hard_drive_md5sum.strip()
  if md5sum:
    print('Checking MD5 checksum...')
    local_md5sum = md5Checksum(downloaded_disk)
    if local_md5sum != md5sum:
      os.remove(downloaded_disk)
      raise Exception('MD5 mismatch. MD5 of local file is %s, Specified MD5 is %s.' % (
          local_md5sum, md5sum))
    print('MD5sum check passed.')
  else:
    print('Warning: not checksum specified.')
  if downloaded_disk.endswith('.gz'):
    try:
      with open(disk_path, 'w') as disk:
        with gzip.open(downloaded_disk, 'rb') as disk_gz:
          shutil.copyfileobj(disk_gz, disk)
    except Exception:
      if os.path.exists(disk_path):
        os.remove(disk_path)
      raise
    os.remove(downloaded_disk)

# Create disk if doesn't exist
# XXX: move to Buildout profile
if not disk_device_path.startswith("/dev/") and not os.path.exists(disk_path):
  print('Creating virtual hard drive...')
  subprocess.check_call([qemu_img_path, 'create' ,'-f', disk_format,
      disk_path, '%sG' % disk_size])
  print('Done.')

# Check and create external disk
additional_disk_list = []
for storage in disk_storage_list:
  if storage:
    key, val = storage.split(' ')
    disk_storage_dict[key.strip()] = val.strip()

map_storage_list, external_disk_number = getMapStorageList(disk_storage_dict,
                                                      int(external_disk_number))

assert len(map_storage_list) == len(disk_storage_dict)
if disk_storage_dict:
  if external_disk_number > 0:
    index = 0
    while (index < len(disk_storage_dict)) and (index < external_disk_number):
      path = disk_storage_dict[map_storage_list[index]]
      if os.path.exists(path):
        disk_filepath = os.path.join(path,
                                  'kvm_virtual_disk.%s' % external_disk_format)
        disk_list = glob.glob('%s.*' % os.path.join(path, 'kvm_virtual_disk'))
        if disk_list == []:
          print('Creating one additional virtual hard drive...')
          process = subprocess.check_call([qemu_img_path, 'create' ,'-f', external_disk_format,
              disk_filepath, '%sG' % external_disk_size])
        else:
          # Cannot change or recreate if disk is exists
          disk_filepath = disk_list[0]
        additional_disk_list.append(disk_filepath)
      else:
        print('Data folder %s was not used to create external disk %r' % (index +1))
      index += 1

additional_disk_options = ''

if disk_aio == 'native':
  additional_disk_options += ',cache.direct=on'


# Generate network parameters
# XXX: use_tap should be a boolean
tap_network_parameter = []
nat_network_parameter = []
numa_parameter = []
number = -1
if use_nat == 'true':
  number += 1
  rules = 'user,id=lan%s' % number
  for rule in nat_rules.split():
    proto = 'tcp'
    rule = rule.split(':')
    if len(rule) == 1:
      port = int(rule[0])
    elif len(rule) == 2:
      proto = rule[0]
      port = int(rule[1])

    rules += ',hostfwd={proto}:{hostaddr}:{hostport}-:{guestport}'.format(
      proto=proto,
      hostaddr=listen_ip,
      hostport=port + 10000,
      guestport=port
    )

  if httpd_port > 0:
    rules += ',guestfwd=tcp:10.0.2.100:80-cmd:%s %s %s' % (netcat_bin,
                                                        listen_ip, httpd_port)
  if cluster_doc_host and cluster_doc_port > 0:
    rules += ',guestfwd=tcp:10.0.2.101:443-cmd:%s %s %s' % (netcat_bin,
                                           cluster_doc_host, cluster_doc_port)
  if set_nat_restrict == 'true':
    rules += ',restrict=on'
  nat_network_parameter = ['-netdev', rules,
          '-device', 'virtio-net-pci,netdev=lan%s,mac=%s' % (number, mac_address)]
if use_tap == 'true':
  number += 1
  vhost = ''
  if enable_vhost == 'true':
    vhost = ',vhost=on'
  tap_network_parameter = ['-netdev',
          'tap,id=lan%s,ifname=%s,script=no,downscript=no%s' % (number,
            tap_interface, vhost),
          '-device', 'virtio-net-pci,netdev=lan%s,mac=%s' % (number, tap_mac_address)]

if enable_device_hotplug != 'true':
  smp = '%s,maxcpus=%s' % (smp_count, smp_max_count)
  ram = '%sM,slots=128,maxmem=%sM' % (ram_size, ram_max_size)
else:
  smp = '1,maxcpus=%s' % smp_max_count
  ram = '%sM,slots=128,maxmem=%sM' % (init_ram_size, ram_max_size)

kvm_argument_list = [qemu_path,
  '-enable-kvm', '-smp', smp, '-name', vm_name, '-m', ram, '-vga', 'std',
  '-drive', 'file=%s,if=%s,cache=%s,aio=%s%s' % (disk_path, disk_type, disk_cache,
                                                 disk_aio, additional_disk_options),
  '-vnc', '%s:1,ipv4,password' % listen_ip,
  '-boot', 'order=cd,menu=on',
  '-qmp', 'unix:%s,server' % socket_path,
  '-pidfile', pid_file_path, '-msg', 'timestamp=on',
  '-D', logfile,
]

rgx = re.compile('^[\w*\,][\=\d+\-\,\w]*$')
for numa in numa_list:
  if rgx.match(numa):
    numa_parameter.extend(['-numa', numa])
kvm_argument_list += numa_parameter

if tap_network_parameter == [] and nat_network_parameter == []:
  print 'Warning : No network interface defined.'
else:
  kvm_argument_list += nat_network_parameter + tap_network_parameter
if language in language_list:
  kvm_argument_list.extend(['-k', language])

for disk in additional_disk_list:
  kvm_argument_list.extend([
            '-drive', 'file=%s,if=%s' % (disk, disk_type)])

if auto_ballooning:
  kvm_argument_list.extend(['-device', 'virtio-balloon-pci,id=balloon0'])

machine_option_list = machine_options.split(',')
if machine_options and len(machine_option_list) > 0:
  name = 'type'
  if '=' in machine_option_list[0]:
    name, val = machine_option_list[0].split('=')
  else:
    val = machine_option_list[0]
  machine_option_list[0] = 'type=%s' % val
  if name == 'type':
    machine = ''
    for option in machine_option_list:
      key, val = option.split('=')
      machine += ',%s=%s' % (key, val)

    kvm_argument_list.extend(['-machine', machine])

if cpu_model:
  rgx = re.compile('^[\w*\,-_][\=\d+\-\,\w]*$')
  if rgx.match(cpu_model):
    kvm_argument_list.extend(['-cpu', cpu_model])

# Try to connect to NBD server (and second nbd if defined).
# If not available, don't even specify it in qemu command line parameters.
# Reason: if qemu starts with unavailable NBD drive, it will just crash.
for nbd_ip, nbd_port in nbd_list:
  if nbd_ip and nbd_port:
    s = getSocketStatus(nbd_ip, nbd_port)
    if s is None:
      # NBD is not available : launch kvm without it
      print 'Warning : Nbd is not available.'
    else:
      # NBD is available
      # We close the NBD socket else qemu won't be able to use it apparently
      s.close()
      kvm_argument_list.extend([
          '-drive',
          'file=nbd:[%s]:%s,media=cdrom' % (nbd_ip, nbd_port)])
# If no NBD is specified/available: use internal disk image
else:
  kvm_argument_list.extend([
      '-drive', 'file=%s,media=cdrom' % default_cdrom_iso
  ])

print 'Starting KVM: \n %s' % ' '.join(kvm_argument_list)
os.execv(qemu_path, kvm_argument_list)