#!/usr/bin/env python

import os, sys
sys.path.append('/usr/lib/python')

from xen.xend import XendClient, sxp
from xen.util.xmlrpclib2 import ServerProxy
import xmlrpclib

class LVM:
    def __init__(self, name):
        self._name = name

    def run(self, cmd):
        fd = os.popen(cmd)
        for line in fd:
            print line.rstrip()

        return not fd.close()

    def path(self):
        return self._name

    def exists(self):
        return self.run("lvdisplay %s > /dev/null" % (self._name))

    def snapshot(self, target):
        return self.run("lvcreate -s -L4M -n %s %s" % (target, self._name))

    def remove(self):
        return self.run("lvremove -f %s" % self._name)

def createSnapshot(lvm, dest):
    snapshot = os.path.basename(dest)

    if not lvm.snapshot(snapshot):
        slvm = LVM(os.path.join(lvm.path(), snapshot))
        if slvm.remove():
            print "Removed old snapshot %s" % snapshot
            if lvm.snapshot(snapshot):
                return True
        print "Snapshot failed: %s" % slvm.path()
        return False

    return True

def getBlockDevice(dominfo):
    vbd = None
    for child in sxp.children(dominfo, 'device'):
        for child in sxp.children(child, 'vbd'):
            vbd = sxp.child_value(child, 'uname')

    if not vbd:
        return None

    if vbd.startswith('phy:'):
        vbd = '/dev/' + vbd[4:]

    return vbd

server = ServerProxy(XendClient.uri)

host, domname, step = sys.argv[1:]

# Only take snapshots for local saves
if step != "2" or not host.startswith('/'):
    sys.exit(0)

try:
    dominfo = server.xend.domain(domname, True)
except xmlrpclib.Fault:
    try:
        dominfo = server.xend.domain('migrating-' + domname, True)
    except xmlrpclib.Fault, inst:
        print inst.faultString
        sys.exit(1)

vbd = getBlockDevice(dominfo)
lvm = LVM(vbd)
rc = createSnapshot(lvm, host)
if rc:
    sys.exit(0)
else:
    sys.exit(1)