#!/usr/bin/perl -w

# AUTHOR
#   Daniel Berteaud <daniel@firewall-services.com>
#
# COPYRIGHT
#   Copyright (C) 2009  Daniel Berteaud
#
#   This program is free software; you can redistribute it and/or modify
#   it under the terms of the GNU General Public License as published by
#   the Free Software Foundation; either version 2 of the License, or
#   (at your option) any later version.
#
#   This program is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU General Public License for more details.
#
#   You should have received a copy of the GNU General Public License
#   along with this program; if not, write to the Free Software
#   Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA



# This script allows you to backup Virtual Manchines managed by libvirt.
# It has only be tested with KVM based VM, but should work with others (like Xen)
# This script will dump:
# * each block devices
# * optionnally the memory (if --state flag is given)
# * the XML description of the VM

# These files are writen in a temporary backup dir. Everything is done in order to minimize
# donwtime of the guest. For example, it takes a snapshot of the block devices (if backed with LVM)
# so the guest is just paused for a couple of seconds. Once this is done, the guest is resumed,
# and the script starts to dump the snapshot.

# Once a backup is finished, you'll have several files in the backup directory. Let's take an
# example with a VM called my_vm which has two virtual disks: hda and hdb. You have passed the --state flag:
# * my_vm.lock: this file is just there to prevent another backup to run at the same time
# * my_vm.xml: this file is the XML description of the VM (for libvirt configuraiton)
# * my_vm_hda.img: this file is an image of the hda drive of the guest
# * my_vm_hdb.img: this file is an image of the hdb drive of the guest
# * my_vm.state: this is a dump of the memory (equivalent to virsh save my_vm /backupdir/my_vm.state)

# This script was made to be ran with BackupPC pre/post commands. In the pre-backup phase, you dump everything
# then, backuppc backups, compress, pools etc... the dumped file. Eventually, when the backup is finished
# The script is called with the --post flag, which cleanups everything.

# Some examples:
#
# Backup the VM named mail01 and devsrv. Also dump the memory.
# Exclude any virtual disk attached as vdb or hdb and on the fly compress the dumped disks (uses gzip by default)
# virt-backup.pl --pre --vm=mail01,devsrv --state --exclude=vdb,hdb --compress

# Remove all the files related to mail01 VM in the backup directory
# virt-backup.pl --post --vm=mail01

# Backup devsrv, use 10G for LVM snapshots (if available), do not dump the memory
# (the guest will just be paused while we take a snapshot)
# Keep the lock file present after the dump
# virt-backup.pl --pre --vm=devsrv --snapsize=10G --keep-lock

# Backup devsrv, and disable LVM snapshots
# virt-backup.pl --pre --vm=devsrv --no-snapshot

# Backup mail01, and enable debug (verbose output)
# virt-backup.pl --pre --vm=mail01 --debug

# Backup mail01. virt-backup is ran on a NAS device (where the disks of the VM are stored.
# Libvirt is running on another host called virthost and is available with SSH
# Specify pbzip2 as compression utility
# Use /tmp/backup as backup directory
# virt-backup.pl --connect=qemu+ssh://user@virthost/system --pre --vm=mail01 --compress=pbzip2 --backupdir=/tmp/backup


use XML::Simple;
use Sys::Virt;
use Getopt::Long;

# Some constant

our %opts = ();
our @vms = ();
our @excludes = ();

# Sets some defaults values
$opts{backupdir} = '/var/lib/libvirt/backup';
$opts{snapsize} = '5G';
$opts{state} = 0;
$opts{debug} = 0;
$opts{keeplock} = 0;
$opts{snapshot} = 1;
$opts{connect} = "qemu:///system";
$opts{compress} = 'none';
$opts{lvcreate} = '/usr/sbin/lvcreate';
$opts{lvremove} = '/usr/sbin/lvremove';
$opts{nice} = 'nice -n 19';
$opts{ionice} = 'ionice -c 2 -n 7';
$opts{livebackup} = 1;
$opts{wasrunning} = 1;

# get command line arguments
GetOptions(
	"debug" 	=> \$opts{debug},
	"keep-lock"	=> \$opts{keeplock},
	"state" 	=> \$opts{state},
	"snapsize=s" 	=> \$opts{snapsize},
	"backupdir=s"	=> \$opts{backupdir},
	"vm=s"		=> \@vms,
	"post"		=> \$opts{post},
	"pre"		=> \$opts{pre},
	"connect=s"	=> \$opts{connect},
	"snapshot!"	=> \$opts{snapshot},
	"compress:s"	=> \$opts{compress},
	"exclude=s"	=> \@excludes
);


# Set compression
if ($opts{compress} eq 'lzop'){
	$opts{compext} = ".lzo";
	$opts{compcmd} = "lzop -c";
}
elsif ($opts{compress} eq 'bzip2'){
	$opts{compext} = ".bz2";
	$opts{compcmd} = "bzip2 -c";
}
elsif ($opts{compress} eq 'pbzip2'){
	$opts{compext} = ".bz2";
	$opts{compcmd} = "pbzip2 -c";
}
elsif ($opts{compress} eq 'xz'){
	$opts{compext} = ".xz";
	$opts{compcmd} = "xz -c";
}
# Default is gzip
elsif (($opts{compress} eq 'gzip') || ($opts{compress} eq '')) {
	$opts{compext} = ".gz";
	$opts{compcmd} = "gzip -c";
}
if($opts{compress} eq 'none'){
	$opts{compext} = "";
	$opts{compcmd} = "cat";
}

# Allow comma separated multi-argument
@vms = split(/,/,join(',',@vms));
@excludes = split(/,/,join(',',@excludes));



# Stop here if we either have pre and post, no pre and no post, or no vm
if (
	($opts{pre} && $opts{post}) ||
	(!$opts{pre} && !$opts{post}) ||
	(!@vms)
	){
	usage();
	exit 1;
}

if (! -d $opts{backupdir} ){
	print "$opts{backupdir} is not a valid directory\n";
	exit 1;
}

### TODO:
# - Make it more robust (script crash sometime while trying to restore
#		Probably a bug somewhere between libvirt and Sys::Virt)
# - Test with images as backend. Should just work, but without the snapshot function
# - Add snapshot support for image based disk ? (should we detect the mount moint, and block device
#	of the storage	or let the user specify it with a --logical ?)
# - Additionnal check that the vm is available after a restore (via $dom->get_info->{status}, ping ?)
# - Support multiple VM in one run (with multiple --vm=<vmname>): DONE
# - Support other compression (lzop, bzip2, pbzip2, xz): DONE
# - Check if compression utilies are available
# - Support virtual drives excludes: DONE (--exclude=hda,vdb. LIMIT: excldues will apply on all VMS for the run)
# - Support per vm excludes in one run
# - add a --keep-lock, which make the lock file persistent until --post is run (if not specified, remove it at the end of --pre): DONE


# Connect to libvirt
print "\n\nConnecting to libvirt daemon using $opts{connect} as URI\n" if ($opts{debug});
our $libvirt = Sys::Virt->new( uri => $opts{connect} ) || 
	die "Error connecting to libvirt on URI: $opts{connect}";



print "\n" if ($opts{debug});


foreach our $vm (@vms){
	# Create a new object representing the VM
	print "Retrieving $vm status from libvirt\n" if ($opts{debug});
	our $dom = $libvirt->get_domain_by_name($vm) ||
		die "Error opening $vm object";
	if ($opts{pre}){
		print "Running pre-backup routine for $vm, as requested by the --pre flag\n" if ($opts{debug});
		pre();
	}
	elsif ($opts{post}){
		print "Running post-backup routine for $vm, as requested by the --post flag\n" if ($opts{debug});
		post();
	}
	else {
		usage();
		exit 1;
	}
}




############################################################################
##############		VARIOUS SUBS			####################
############################################################################


sub pre{
	# Create a new XML object
	my $xml = new XML::Simple ();
	my $data = $xml->XMLin( $dom->get_xml_description() );

	die "Another backup is running\n" if ( -e "$opts{backupdir}/$vm.lock" );

	# Lock VM: Create a lock file so only one backup process can run
	lock_vm();

	# Save the XML description
	save_xml();

	# Save the VM state if it's running and --state is present (else, just suspend the VM)
	$opts{wasrunning} = 0 unless (is_running());

	if ($opts{wasrunning}){
		if ($opts{state}){
			save_vm_state();
		}
		else{
			suspend_vm();
		}
	}

	my @disks;

	# Create a list of disks for the VM
	foreach $disk (@{$data->{devices}->{disk}}){

		my $source = $disk->{source}->{dev} || '';
		my $target = $disk->{target}->{dev};

		if (!!grep { $_ eq "$target" } @excludes){
			print "\n\nSkiping $source for vm $vm as it's matching one of the excludes: ".join(",",@excludes)."\n" if ($opts{debug});
			next;
		}

		# If the device is a disk (and not a cdrom) and the source dev exists
		if (($disk->{device} eq 'disk') && (-e $source) && (!$disk->{readonly})){

			print "\n\nAnalysing disk $source connected on $vm as $target\n" if ($opts{debug});

			# If it's a block device
			if ($disk->{type} eq 'block'){
				
				my $time = "_".time();
				# Try to snapshot the source is snapshot is enabled
				if ( ($opts{snapshot}) && (create_snapshot($source,$time)) ){
					print "$source seems to be a valid logical volume (LVM), a snapshot has been taken as ". $source . $time ."\n" if ($opts{debug});
					$source = $source.$time;
					push (@disks, {source => $source, target => $target, type => 'snapshot'});
				}
				# Snapshot failed, or disabled: disabling live backups
				else{
					if ($opts{snapshot}){
						print "Snapshoting $source has failed (not managed by LVM, or already a snapshot ?), live backup will be disabled\n" if ($opts{debug}) ;
					}
					else{
						print "Not using LVM snapshots, live backups will be disabled\n" if ($opts{debug});
					}
					$opts{livebackup} = 0;
					push (@disks, {source => $source, target => $target, type => 'block'});
				}
			}
			elsif ($disk->{type} eq 'file'){
				$opts{livebackup} = 0;
				push (@disks, {source => $source, target => $target, type => 'file'});
			}
			print "Adding $source to the list of disks to be backed up\n" if ($opts{debug});
		}
	}

	if ($opts{debug}){
		print "\n\nThe following disks will be dumped:\n";
		foreach $disk (@disks){
			print "Source: $disk->{source}\tDest: $opts{backupdir}/$vm" . '_' . $disk->{target} . ".img$opts{compext}\n";		
		}
		print "\n";
	}

	# If livebackup is possible (every block devices can be snapshoted)
	# We can restore the VM now, in order to minimize the downtime
	if ($opts{livebackup}){
		print "\nWe can run a live backup\n" if ($opts{debug});
		if ($opts{wasrunning}){
			if ($opts{state}){
				restore_vm();
			}
			else{
				resume_vm();
			}
		}
	}

	foreach $disk (@disks){

		my $source = $disk->{source};
		my $dest = "$opts{backupdir}/$vm" . '_' . $disk->{target} . ".img$opts{compext}";

		print "\nStarting dump of $source to $dest\n" if ($opts{debug});
		my $ddcmd = "$opts{ionice} dd bs=16k if=$source ";
		if ($opts{compress} ne 'none'){
			$ddcmd .=  "2>/dev/null | $opts{nice} $opts{compcmd} > $dest 2>/dev/null";
		}
		else{
			$ddcmd .= "of=$dest 2>/dev/null";
		}
		unless( system("$ddcmd") == 0 ){
			die "Couldn't dump the block device/file $source to $dest with dd\n";
		}
		destroy_snapshot($source) if ($disk->{type} eq 'snapshot');
	}
	if ($opts{wasrunning}){
		if ($opts{state}){
			restore_vm();
		}
		else{
			resume_vm();
		}
	}
	unlock_vm() unless ($opts{keeplock});
}


sub post{
	
	if (! is_running()){
		restore_vm();
	}
	elsif ($dom->get_info->{state} == Sys::Virt::Domain::STATE_PAUSED){
		resume_vm();
	}
	cleanup();
}

sub usage{
	print "usage:\n$0 --pre|--post --vm=name[,vm2,vm3] [--debug] [--exclude=hda,hdb] [--compress] [--state] [--no-snapshot] [--snapsize=<size>] [--backupdir=/path/to/dir] [--connect=<URI>] [--keep-lock]\n" .
	"\n\nMandatory options:\n" .
	"\t--pre: Run the pre backup routine (dump disk image to temp dir, pausing the VM if needed)\n\n" .
	"\t--post: Run the post-backup routine, cleaning up the temp dir\n\n" .
	"\t\tNote: --pre and --post are mutually exclusive\n\n" .
	"\t--vm=name: The VM you want to work on (as known by libvirt). You can backup several VMs in one shot if you separate them with comma, or with multiple --vm argument\n\n" .
	"\n\nOther options:\n\n" .
	"\t--state: Cleaner way to take backups. If this flag is present, the script will save the current state of the VM (if running) instead of just suspending it. With this you should be able to restore the VM at the exact state it was when the backup started. The reason this flag is optional is that with the currents libvirt versions, saving the state shuts down the VM, and sometimes, the restoration fails, which result in a stoped (in a unclean way) VM\n\n" .
	"\t--no-snapshot: Do not attempt to use LVM snapshots. If not present, the script will try to take a snapshot of each disk of type 'block'. If all disk can be snapshoted, the VM is resumed, or restored (depending on the --state flag) immediatly after the snapshots have been taken, resulting in no downtime (or very few). If at least one disk cannot be snapshoted, the VM is suspended (or stoped) for the time the disks are dumped in the temp dir. That's why you should use a fast support for the temp dir (fast disks, RAID0 or RAID10)\n\n" .
	"\t--snapsize=<snapsize>: The amount of space to use for snapshots. Use the same format as -L option of lvcreate. eg: --snapsize=15G. Default is 5G\n\n" .
	"\t--compress[=[gzip|bzip2|pbzip2|lzop|xz]]: On the fly compress the disks images during the dump. If you don't specify a compression algo, gzip will be used.\n\n" .
	"\t--exclude=hda,hdb: Prevent the disks listed from being dumped. The names are from the VM perspective, as configured in livirt as the target element. It can be usefull for example if you want to dump the system disk of a VM, but not the data one which can be backed up separatly, at the files level.\n\n" .
	"\t--backupdir=/path/to/backup: Use an alternate backup dir. The directory must exists and be writable. The default is /var/lib/libvirt/backup\n\n" .
	"\t--connect=<URI>: URI to connect to libvirt daemon (to suspend, resume, save, restore VM etc...). The default is qemu:///system.\n\n" .
	"\t--keep-lock: Let the lock file present until --post is called for the VM(s). This prevent another --pre to run while an third party backup software (BackupPC for example) saves the dumped files.\n\n";
}

# Save a runnign VM
sub save_vm_state{
	if (is_running()){
		print "$vm is running, saving state....\n" if ($opts{debug});
#		$dom->suspend();
#		sleep(2);
		$dom->save("$opts{backupdir}/$vm.state");
		print "$vm state saved as $opts{backupdir}/$vm.state\n" if ($opts{debug});
	}
	else{
		print "$vm is not running, nothing to do\n" if ($opts{debug});
	}
}

sub restore_vm{
	if (! is_running()){
		if (-e "$opts{backupdir}/$vm.state"){
			print "\nTrying to restore $vm from $opts{backupdir}/$vm.state\n" if ($opts{debug});
			$libvirt->restore_domain("$opts{backupdir}/$vm.state");
			print "Waiting for restoration to complete\n" if ($opts{debug});
			my $i = 0;
			while ((!is_running()) && ($i < 120)){
				sleep(5);
				$i = $i+5;
			}
#			sleep(2);
#			$dom->resume();
			print "Timeout while trying to restore $vm, aborting\n" if (($i > 120) && ($opts{debug}));
		}
		else{
			print "\nRestoration impossible, $opts{backupdir}/$vm.state is missing\n" if ($opts{debug});
		}
	}
	else{
		print "\nCannot start domain restoration, $vm is running (already restored after a live backup ?)\n" if ($opts{debug});
	}
}

sub suspend_vm(){
	if (is_running()){
		print "$vm is running, suspending\n" if ($opts{debug});
		$dom->suspend();
		print "$vm now suspended\n" if ($opts{debug});
	}
	else{
		print "$vm is not running, nothing to do\n" if ($opts{debug});
	}
}

sub resume_vm(){
	if ($dom->get_info->{state} == Sys::Virt::Domain::STATE_PAUSED){
		print "$vm is suspended, resuming\n" if ($opts{debug});
		$dom->resume();
		print "$vm now resumed\n" if ($opts{debug});
	}
	else{
		print "$vm is not suspended, nothing to do\n" if ($opts{debug});
	}
}

sub save_xml{
	print "\nSaving XML description for $vm to $opts{backupdir}/$vm.xml\n" if ($opts{debug});
	open(XML, ">$opts{backupdir}/$vm" . ".xml") || die $!;
	print XML $dom->get_xml_description();
	close XML;
}

sub cleanup{
	print "\nRemoving backup files\n" if ($opts{debug});
	my $cnt= unlink <$opts{backupdir}/$vm*>;
	print "$cnt file(s) removed\n" if $opts{debug};
}

sub create_snapshot{
	my ($blk,$suffix) = @_;
	my $ret = 0;
	print "Running: $opts{lvcreate} -p r -s -n " . $blk . $suffix ." -L $opts{snapsize} $blk > /dev/null 2>&1\n" if $opts{debug};
	if ( system("$opts{lvcreate} -s -n " . $blk . $suffix ." -L $opts{snapsize} $blk > /dev/null 2>&1") == 0 ) {
		$ret = 1;
	}
	return $ret;
}

sub destroy_snapshot{
	my $ret = 0;
	my ($snap) = @_;
	print "Removing snapshot $snap\n" if $opts{debug};
	if (system ("$opts{lvremove} -f $snap > /dev/null 2>&1") == 0 ){
		$ret = 1;
	}
	return $ret;
}

sub lock_vm{
	print "Locking $vm\n" if $opts{debug};
	open ( LOCK, ">$opts{backupdir}/$vm.lock" ) || die $!;
	print LOCK "";
	close LOCK;
}

sub unlock_vm{
	print "Removing lock file for $vm\n" if $opts{debug};
	unlink <$opts{backupdir}/$vm.lock>;
}

# Check if a VM is running
sub is_running{
	my $status = $dom->get_info->{state};
	my $ret = 0;
	
	if (($status == Sys::Virt::Domain::STATE_RUNNING) ||
	   ($status == Sys::Virt::Domain::STATE_BLOCKED) ||
	   ($status == Sys::Virt::Domain::STATE_SHUTDOWN) ||
	   ($status == Sys::Virt::Domain::STATE_PAUSED)) {
		$ret = 1
	}
	return $ret;
}