#!/usr/bin/python
# Migrates the machine passed as arguments from the dev cluster.
# To be run on the prod cluster.

from invirt import remctl as r
from lib import database
import subprocess
import sys
import time
import os

sys.stdout = os.fdopen(sys.stdout.fileno(), 'w', 0)

kinit = dict(principal='host/aperture-science.mit.edu', keytab='/etc/krb5.keytab')

dev_db_uri = 'postgres://sipb-xen@sipb-xen-dev.mit.edu/sipb_xen'
database.connect(dev_db_uri)
dev_sess = database.session

database.connect()
prod_sess = database.session

## dump from dev db
def take_data(machine_name):
  dev_sess.begin()
  machine = dev_sess.query(database.Machine).filter_by(name=machine_name).one()
  
  # Clean out the ACL just so we don't have to think about it
  machine.acl = []
  dev_sess.update(machine)
  
  print 'VM Info:'
  print '  name: %s' % machine.name
  print '  description: %s' % machine.description
  print '  cpus: %s' % machine.cpus
  print '  memory: %s' % machine.memory
  print '  owner: %s' % machine.owner
  print '  contact: %s' % machine.contact
  print '  administrator: %s' % machine.administrator
  print '  uuid: %s' % machine.uuid
  print '  type: %s' % machine.type.type_id
  print '  autorestart: %s' % machine.autorestart
  print '  adminable: %s' % machine.adminable
  print '  Disks:'
  for disk in machine.disks:
    print '  - %s (%s)' % (disk.guest_device_name, disk.size)
  print '  NICs:'
  for nic in machine.nics:
    print '  - %s, %s, %s' % (nic.mac_addr, nic.ip, nic.hostname)
  print '==============================================='
  print
  
  disks = machine.disks
  nics = machine.nics
  for r in disks + nics + [machine]:
    dev_sess.delete(r)
  
  dev_sess.flush()
  
  for r in disks + nics + [machine]:
    dev_sess.expunge(r)
    del r._instance_key
  
  return machine

## add to prod db
def restore_data(machine, session):
  # The machine's type is still the one attached to the dev database;
  # get the right one
  machine.type = session.query(database.Type).filter_by(type_id=machine.type.type_id).one()
  session.begin()
  session.save(machine)
  session.commit()
  
def migrate_vm(machine_name):
  # Power off the VM on dev
  #
  # This has to be done first, because once the machine is deleted
  # from the database, we can't remctl for it anymore
  out, err = r.remctl('xvm-remote.mit.edu', 'control', machine_name, 'destroy', err=True, **kinit)
  print out
  
  machine = take_data(machine_name)
  subprocess.call(['zwrite', '-d', '-c', 'xvm-auto', '-i', 'migration', '-s', 'XVM Migration Script', '-m', 
                   'Migrating %s, disk size %0.2fG...' % (machine.name, sum(disk.size for disk in machine.disks) / 1024.0)])
  
  success = True
  ## copy disk image... copy, copy...
  for disk in machine.disks:
    lvname='d_%s_%s' % (machine.name, disk.guest_device_name)
    
    if 0 != subprocess.call(['lvcreate', '-L%sM' % str(disk.size), '-n', lvname, 'xenvg']):
      success = False
      break
    
    ssh = subprocess.Popen(['rsh',
                '10.6.0.165',
                'dd', 'if=/dev/xenvg/%s' % lvname, 'bs=1M'],
                 stdout=subprocess.PIPE)
    dd = subprocess.Popen(['dd', 'of=/dev/xenvg/%s' % lvname, 'bs=1M'],
                stdin=ssh.stdout)
    if 0 != dd.wait():
      success = False
      break
    if 0 != ssh.wait():
      success = False
      break
  
  if not success:
    dev_sess.rollback()
    
    print '==============================================='
    print 'ERROR: VM %s failed to migrate' % machine.name
    print '==============================================='
  else:
    restore_data(machine, prod_sess)
    dev_sess.commit()
  
  return success

if __name__ == '__main__':
  while True:
    r.checkKinit(**kinit)
    p = subprocess.Popen(['curl', '-s', '-k', '--negotiate', '-u', ':', 'https://xvm.mit.edu:442/offlist'], stdout=subprocess.PIPE)
    if 0 != p.wait():
        subprocess.call(['zwrite', '-d', '-c', 'xvm', '-i', 'migration', '-s', 'XVM Migration Script', '-m',
                         'Failed to get list of remaining VMs. Will try again in 15 seconds'])
        time.sleep(15)
        continue
    
    next_line = p.stdout.read().split('\n')[0]
    if next_line == '':
      subprocess.call(['zwrite', '-d', '-c', 'xvm', '-i', 'migration', '-s', 'XVM Migration Script', '-m',
                       'XVM migration complete'])
      break
    
    next, uptime = next_line.split('\t')[:2]
    
    print '==============================================='
    print 'Migrating %s' % next
    print '==============================================='
    if not migrate_vm(next):
      subprocess.call(['zwrite', '-d', '-c', 'xvm', '-i', 'migration', '-s', 'XVM Migration Script', '-m',
                       'Error in migrating %s' % next])
      sys.exit(1)
    if uptime.strip() != '':
      r.remctl('xvm-remote-dev.mit.edu', 'control', next, 'create', **kinit)
    
    subprocess.call(['zwrite', '-d', '-c', 'xvm-auto', '-i', 'migration', '-s', 'XVM Migration Script', '-m',
                     'done'])