#!%(python-path)s
# 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

# XXX: give all of this through parameter, don't use this as template, but as module
qemu_img_path = '%(qemu-img-path)s'
qemu_path = '%(qemu-path)s'
disk_size = '%(disk-size)s'
disk_type = '%(disk-type)s'
socket_path = '%(socket-path)s'
nbd_list = (('%(nbd-host)s', %(nbd-port)s), ('%(nbd2-host)s', %(nbd2-port)s))
default_disk_image = '%(default-disk-image)s'
disk_path = '%(disk-path)s'
virtual_hard_drive_url = '%(virtual-hard-drive-url)s'.strip()
virtual_hard_drive_md5sum = '%(virtual-hard-drive-md5sum)s'.strip()
virtual_hard_drive_gzipped = '%(virtual-hard-drive-gzipped)s'.strip().lower()
nat_rules = '%(nat-rules)s'.strip()
use_tap = '%(use-tap)s'.lower()
use_nat = '%(use-nat)s'.lower()
tap_interface = '%(tap-interface)s'
listen_ip = '%(ipv4)s'
mac_address = '%(mac-address)s'
tap_mac_address = '%(tap-mac-address)s'
smp_count = '%(smp-count)s'
smp_options = '%(smp-options)s'.strip()
numa_list = '%(numa)s'.split()
ram_size = '%(ram-size)s'
pid_file_path = '%(pid-file-path)s'
external_disk_number = %(external-disk-number)s
external_disk_size = '%(external-disk-size)s'
external_disk_format = '%(external-disk-format)s'
disk_storage_dict = {}
disk_storage_list = """%(disk-storage-list)s""".split('\n')
map_storage_list = []
etc_directory = '%(etc-directory)s'.strip()
httpd_port = %(httpd-port)s
netcat_bin = '%(netcat-binary)s'.strip()
cluster_doc_host = '%(cluster-doc-host)s'
cluster_doc_port = %(cluster-doc-port)s

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
    urllib.urlretrieve(virtual_hard_drive_url, downloaded_disk)
  except:
    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 os.path.exists(disk_path):
  print('Creating virtual hard drive...')
  subprocess.Popen([qemu_img_path, 'create' ,'-f', 'qcow2',
      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()

if not external_disk_format in ['qcow2', 'raw', 'vdi', 'vmdk', 'cloop', 'qed']:
  external_disk_format = 'qcow2'

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...')
          subprocess.Popen([qemu_img_path, 'create' ,'-f', '%%s' %% 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


# 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 + ','.join('hostfwd=tcp:%%s:%%s-:%%s' %% (listen_ip,
                        int(port) + 10000, port) for port in nat_rules.split())
  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)
  nat_network_parameter = ['-netdev', rules,
          '-device', 'e1000,netdev=lan%%s,mac=%%s' %% (number, mac_address)]
if use_tap == 'true':
  number += 1
  tap_network_parameter = ['-netdev',
          'tap,id=lan%%s,ifname=%%s,script=no,downscript=no' %% (number,
            tap_interface),
          '-device', 'e1000,netdev=lan%%s,mac=%%s' %% (number, tap_mac_address)]

smp = smp_count
if smp_options:
  for option in smp_options.split(','):
    key, val = option.split('=')
    if key in ('cores', 'threads', 'sockets', 'maxcpus') and val.isdigit():
      smp += ',%%s=%%s' %% (key, val)
kvm_argument_list = [qemu_path,
  '-enable-kvm', '-smp', smp,
  '-m', ram_size, '-vga', 'std',
  '-drive', 'file=%%s,if=%%s' %% (disk_path, disk_type),
  '-vnc', '%%s:1,ipv4,password' %% listen_ip,
  '-boot', 'order=cd,menu=on',
  '-qmp', 'unix:%%s,server' %% socket_path,
  '-pidfile', pid_file_path,
]

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

for disk in additional_disk_list:
  kvm_argument_list.extend([
            '-drive', 'file=%%s,if=%%s' %% (disk, disk_type)])
# 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
      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_disk_image
  ])

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