#!/usr/bin/perl
#
#
# Copyright (c) 2006 Zmanda Inc.  All Rights Reserved.
#
# This program is free software; you can redistribute it and/or modify it
# under the terms of the GNU General Public License version 2 as published
# by the Free Software Foundation.
#
# 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
#
# Contact information: Zmanda Inc, 505 N Mathlida Ave, Suite 120
# Sunnyvale, CA 94085, USA, or: http://www.zmanda.com
#
#

use strict;
use warnings;
use File::Copy;
use File::Find;
use File::Path;
use File::Spec::Functions;
use File::Basename;
use File::Temp qw/ :POSIX /;
use lib '/usr/local/lib/mysql-zrm';
use ZRM::Common;
use ZRM::MySQL;

my $MYSQLBINLOG="mysqlbinlog";
my $nextBinLog;

#usage strings 
my $USAGE_RESTORE_STRING=  
		"\t\t[--user <user>] [--password <password>]\n".  
		"\t\t[--host <hostname>] [--port <portnumber>]\n".  
		"\t\t[--socket <name>] [--ssl-options <\"MySQL\ ssl\ options\">]\n".
		"\t\t[--mysql-binpath\ <mysql\ binaries\ directory>]\n".
		"\t\t[--mysql-shutdown|--no-mysql-shutdown]\n".
		"\t\t[--retry-count\ <count>]\n".
		"\t\t[--retry-delay\ <delay in seconds>]\n".
		"\t\t[--source-directory <directory\ name>]\n".  
		"\t\t[--bin-logs <\"/fullpath/name1\ /fullpath/name2\ ...\">]\n".
		"\t\t[--all-databases] [--databases <\"name1\ name2\ ...\">]\n".  
		"\t\t[--replication\ |\ --noreplication]\n".
		"\t\t[--start-position <#>] [--stop-position <#>] [--offset <#>]\n".
        	"\t\t[--start-datetime <name>] [--stop-datetime <name>]\n".
		"\t\t[--copy-plugin <plugin>]\n".
		"\t\t[--passfile <path>] [--ssh-user <user>]\n".
		"\t\t[--socket-remote-port <port>]\n";

my @RESTOREOPT =  qw/
			user=s
              		password=s
              		host=s
              		port=i
              		socket=s
              		ssl-options=s
			mysql-binpath=s
			mysql-shutdown!
			retry-count=i
			retry-delay=i
              		source-directory=s
              		bin-logs=s
			to-file:s
			from-file=s
              		all-databases
              		databases=s
			replication!
              		start-position=i
              		stop-position=i
              		offset=i
              		start-datetime=s
              		stop-datetime=s
              		passfile=s
              		copy-plugin=s
              		ssh-user=s
              		socket-remote-port=i
			/;

my $retryCount = 3;
my $sleepTime = 5;
my $currentRetrys = 0;
my $snapshotHost = "";
my $snapshotPlugin = "";
my %mountDetails;

#Handles the restore action
sub doRestore()
{
	my $r;
	# For some reason at times, even after 
	# mysqladmin --shutdown returns, it takes a few seconds 
	# for the mysql server to actually go down. 
	# So adding a retry just for restore of the first set of files.
	if( defined $inputs{"retry-count"} ){
		$retryCount = $inputs{"retry-count"};
		$retryCount += 1;
	}
	if( $retryCount <= 0 ){
		$retryCount = 1;
	}
	if( defined $inputs{"retry-delay"} ){
		$sleepTime = $inputs{"retry-delay"};	
	}
	if( $sleepTime < 0 ){
		$sleepTime = 0;
	}

	if( $inputs{"source-directory"} ) {
		if( $inputs{"bin-logs"} ){
			&usage( "Both --source-directory and --bin-logs cannot be specified at the same time. Please supply only one of these.\n" );
		}
		if( !-d $inputs{"source-directory"} ) {
			&printAndDie( "Cannot find source directory ".$inputs{"source-directory"}."\n" );
		}
		my $filename = catfile( $inputs{"source-directory"}, $INDEX_FILENAME );
		$r = &parseIndexFile( $filename );
		if( $r == 0 ){
			&printAndDie("cannot open index file $filename $!\n");
		}
		if( defined $indexdata{"mysql-server-os"} ){
			if( $indexdata{"mysql-server-os"} ne $mysql_server_os ){
				&printAndDie( "backup data from ".$indexdata{"mysql-server-os"}." cannot be restored to $mysql_server_os\n" );
			
			}
		}
		if( !defined $indexdata{"backup-status"} || 
		   $indexdata{"backup-status"} ne "Backup succeeded" ){
			&printWarning( "Restoring from a backup which was not successful.\n" ); 
			if( defined $indexdata{"backup-status"} ){
				&printWarning( "Backup status reported in the index file is '".$indexdata{"backup-status"}."'\n" );
			}else{
				&printWarning( "Backup status is not available in the index file. It looks like this is a completly failed backup\n" );
			}
			&printWarning( "Restoring from this may not be successful\n" );
		}
		if( defined $indexdata{"compress"} || defined $indexdata{"encrypt"} ){
			$r = 1;
			$r = &uncompressBackup( $inputs{"source-directory"} );
			if( $r == 0 ){
				&printAndDie( "Unable to uncompress backup\n" );
			}
		}
	} else{
		if( !$inputs{"bin-logs"} ){
			&usage( "For restore action specify either --source-directory or --bin-logs\n" );
		}
	}
	my $isReplicating = 0;
	$r = 0;
        $r = isSlave();
        if( $r ) {
                $r = stopSlave();
                if( $r == 1 ){
                        $isReplicating = 1;
                }
        }
	$r = 0;
	if( $indexdata{"backup-level"} ){
		$r = $indexdata{"backup-level"};
	}
	if( $inputs{"source-directory"} && $r == 0 ){
		doFullRestore();
	} else {
		if( $inputs{"databases"} ) {
			my @db = split( " ", $inputs{"databases"} );
			foreach( @db ) {
				doIncRestore( $_ );
			}
		} else {
			doIncRestore();
		}
	}
	if( $zrm_error == 0 && defined $inputs{"mysql-shutdown"} && $inputs{"mysql-shutdown"} == 1 && $mysql_shutdown == 0 ){
		&shutdownMySQL();
	}
	if( $zrm_error == 0 && $indexdata{"replication"} ) {
		restoreReplicationData();
	}	
	if( $mysql_shutdown == 0 && $isReplicating eq "1" ) {
                &startSlave();
        }

	if( defined $indexdata{"compress"} ){
		&removeUncompressedBackup( $inputs{"source-directory"}, \%indexdata );
	}
	my $time_taken = time() - $start_time;
	my $msg;
	if( $zrm_error > 0 ){
		&printError( "Restore failed\n" );
	}else{
		&printLog("Restore done in ".$time_taken." seconds.\n");
	}
	if( $mysql_shutdown ){
		print "MySQL server has been shutdown. Please restart after verification.\n";
	}
}

#restores the replication slaves data
sub restoreReplicationData()
{
	if( defined $inputs{"replication"} && $inputs{"replication"} == 0 ){
		return;
	}
	my $dest = $datadir;
	my $r = 1;
	if( $indexdata{"replication"} ) {
		$r = &copyFiles( $dest, $inputs{"source-directory"}, $indexdata{"replication"} );
		if( $r == 1 ){
			&printLog( "Restored replication related files '".$indexdata{"replication"}."'\n");
		}
	}
	$r = 1;
	if( $indexdata{"slave-load-files"} ) {
		my $sqlLoadFile = $indexdata{"slave-load-files"};
		$r = &copyFiles( $slave_load_tmpdir, $inputs{"source-directory"}, $sqlLoadFile );
		if( $r == 1 ){
			&printLog( "Restored slave-load-files '".$indexdata{"slave-load-files"}."' to ".$slave_load_tmpdir."\n");
		}
	}
}

sub doRestoreFromSQLFile()
{
	my $f = $_[0];
	my $y = &addMySQLParams($MYSQL);
	my $sql_cmd = "set character_set_client=utf8;";
	$sql_cmd = $sql_cmd."set character_set_connection=utf8;";
	$sql_cmd = $sql_cmd."set character_set_database=utf8;";
	$sql_cmd = $sql_cmd."set character_set_results=utf8;";
	$sql_cmd = $sql_cmd."set character_set_server=utf8;";
	#$sql_cmd = $sql_cmd."set character_set_system=utf8;";
	$sql_cmd = $sql_cmd."source ".$f.";";
	my $x = " -e \"$sql_cmd\"";
	if( $_[1] ){
		$x .= " $_[1]";
	}
	if( $verbose ) {
		&printLog( "restoring using command ".$y.$x."\n" );
	}
	$x = system( $y.$x );
	return $x;
}

# $_ list of databases to restore.
sub extractToTmpFile()
{
        my $sqlmaster = catfile( $inputs{"source-directory"}, $BACKUPSQLFILE );
	my $dbs = $_[0];
        $dbs =~s/\s/|/g;
        open( MSQL, $sqlmaster ) or &printAndDie( "Could not open file $sqlmaster $!\n" );
	my $tf = tmpnam();
        open(OF,">$tf" ) or &printAndDie( "Coupld not open tmp file $tf $!\n" );
        my $found = 0;
        while(<MSQL>){
                if( /^-- Current Database: / ){
                        # If the line starts with "-- Current Database"
                        # then checkif the db name in that line
                        # is one that we are interested in.
                        $found = /\`($dbs)\`/ ? 1:0;
                }
                if( $found ){
                        print OF $_;
                }elsif( /^\/\*\!\d*\s*SET / ){
                        # Also print if statement is "/*! SET "
                        print OF $_;
                }
        }
        close( OF );
	close( MSQL );
	return $tf;
} 
#executes the logical restore
#$_[0] list of databases logically backed
sub doLogicalRestore()
{
	my $f;
	my $x;
	if( $indexdata{"logical-tables"} ){
		$f = catfile( $inputs{"source-directory"}, $BACKUPSQLFILE );
		$x = &doRestoreFromSQLFile( $f, $_[0] );
	}else{
		$f = &extractToTmpFile( $_[0] );
		$x = &doRestoreFromSQLFile( $f );
		unlink $f;
	}
	if( $x > 0 ) {
		&printError( "Restore from logical backup failed\n" );
	} else {
		&printLog( "Restored database(s) from logical backup: ".$_[0]."\n");
	}
}

#$_[0] specifies the database name
sub copyDataUsingCopyPlugin()
{
        my $cmd = $inputs{"copy-plugin"};
        my @params;
        push @params, "--source-host";
	if( $snapshotHost eq "" ){
        	push @params, "localhost";
	}else{
		push @params, $snapshotHost;
	}
        push @params, "--source-file";
	if( $snapshotHost eq "" ){
        	push @params, catfile( $inputs{"source-directory"}, $_[0] );
	}else{
        	my $d = catfile( $inputs{"source-directory"}, $LINK_POINT );
        	push @params, catfile( $d, $_[0] );
	}
        push @params, "--destination-host";
        if( $inputs{"host"} ){
                push @params, $inputs{"host"};
        }else{
                push @params,  "localhost";
        }
        push @params, "--destination-directory";
        push @params, catfile( $datadir );
        if( $verbose ){
                &printLog( "Command being executed is $cmd @params\n" );
        }
        my $r;
	while( 1 ){
		$r = system( $cmd, @params );
		$currentRetrys += 1;
		if( $r == 0 ){
			$currentRetrys = $retryCount;
		}
		if( $currentRetrys >= $retryCount ){
			last;
		}
		&printLog( "Retrying... (retry count: $currentRetrys)\n" );
		sleep( $sleepTime );
	}
        if( $r > 0 ){
                &printError( "Could not copy database $_[0] $!\n" );
                &printError( "copy-plugin exited with error $r\n" );
                return 0;
        }else{
		&printLog("Restored database from raw backup: ".$_[0]."\n");
	}
        return 1;
}

#Creates the database directory and copies all of the table info into it
#$_[0] specifies the database name
sub createDirAndCopyData()
{
	my $x;
	if( $snapshotHost eq "" ){
		$x = catfile( $inputs{"source-directory"}, $_[0] );
	}else{
		$x = catfile( $inputs{"source-directory"}, $LINK_POINT );
		$x = catfile( $x, $_[0] );
	}
	my $f = catfile( $datadir, $_[0] );
	if( ! -e $f ) {
		my ($sdev,$sino,$smode,$snlink,$suid,$sgid,$srdev,
		$ssize, $satime,$smtime,$sctime,$sblksize,$sblocks)
		= stat $x;
		mkpath( $f, 0, $smode & 07777 );
		chown( $suid, $sgid, $f );
	}
	my $r = 1;
	my $y = $ARGV[0];
	if( $x=~/\s/ ){
        	$x = "\"$x\"";
	}
	$x = catfile( $x, "*" );
	my @ff = glob $x;
	foreach( @ff ) {
		$r = &copyOneFileLocal( $_, $f );
	}
	if( $r == 1 ){
		&printLog("Restored database from raw backup: ".$_[0]."\n");
	}else{
		&printError("Failed to restore database from raw backup: ".$_[0]."\n");
	}
}

#Create list of databases to be restored
#$_[0] specifies the list from the index file
sub createDBList()
{
	my $list = "";
	if( $inputs{"databases"} ){
		my @input = split( " ", $inputs{ "databases" } );
		my @index = split( " ", $_[0] );
		foreach( @input ){
			my $inp = $_;
			foreach( @index ) {
				if( $inp eq $_ ){
					$list = $list." ".$_;
				}
			}
		}
	} else {
		$list = $_[0];
	}
	return $list;
}

# Restores the specified innodb files
# $_[0] should be either "innodb_data" or "innodb_logs"
sub doInnoDBRestore()
{
	if( ! $indexdata{$_[0]} ){
		return;
	}	
	my @y;
	my $msg;
	if( $_[0] eq "innodb-logs" ){
		my $dir = dirname($indexdata{$_[0]});
		if( $snapshotHost eq "" || $snapshotHost eq "localhost" ){
			my $x;
			if( $snapshotHost eq "" ){
				$x = catfile( $inputs{"source-directory"}, $indexdata{$_[0]} );
			}else{
				$x = catfile( $inputs{"source-directory"}, $LINK_POINT );
				$x = catfile( $x, $indexdata{$_[0]} );
			}
			my @t = glob $x;
			my @suf;
			foreach( @t ){
				my $name = basename( $_ );	
				$name = catfile( $dir, $name );
				push @y, $name;	
			}
		}else{
			push @y, $indexdata{$_[0]};
		}
		$msg = "log";
	}else{
		@y = split( /;/, $indexdata{$_[0]} );
		$msg = "data file";
	}
	foreach( @y ){	
		my $file = $_;
		my $dir = dirname( $file );
		my $source;
		if( $snapshotHost ne "" ){
			$source = catfile( $inputs{"source-directory"}, $LINK_POINT );
			$source = catfile( $source, $file );
		}else{
			$source = catfile( $inputs{"source-directory"}, $file );
		}
		my $r;
		while( 1 ){
			$r = &copyTo( $source, $dir, $snapshotHost, "no error" );
			$currentRetrys += 1;
			if( $r == 1 ){
				$currentRetrys = $retryCount;
			}
			if( $currentRetrys >= $retryCount ){
				last;
			}
			&printLog( "Retrying... (retry count: $currentRetrys)\n" );
			sleep( $sleepTime);
		}
		if( $r == 1 ){
			&printLog( "Restored innodb $msg '$_'\n");
		}else{
			&printError( "Failed to restore innodb $msg '$_'\n" );
		}
	}
}

sub mountAll()
{
	foreach( keys%mountDetails ){
		my $tmpDir = $_;
		my $snp = $mountDetails{$_};
		my @a = split /;/, $snp;
		$snp = $a[0];
		my $fs = $a[1];
                my $command = $snapshotPlugin;
                $command .= " --action mount";
                $command .= " --dev $snp";
                $command .= " --directory $tmpDir";
                $command .= " --fstype $fs";
                if( $snapshotHost ne "localhost" ){
                        $command .= " --host ".$snapshotHost;
                }
		$command .= " > $LOGGER 2>&1";
                if( $verbose ) {
                        &printLog( "Mounting snapshot\n" );
                        &printLog( $command."\n" );
                }
		$ENV{'ZRM_NO_ERROR'} = 1;
                my $r = system( $command );
		delete $ENV{'ZRM_NO_ERROR'};
                if( $r ) {
			if( $verbose ){
                        	&printCommandOutputToLog( "INFO", "mount", $LOGGER );
			}
                }
	}

}

#executes the full restore
sub doFullRestore()
{
	my $db_list = "";
	if( $indexdata{"logical-databases"} ) {
		$db_list = &createDBList( $indexdata{"logical-databases"} );
	}
	if( $indexdata{"logical-database"} ) {
		$db_list = &createDBList( $indexdata{"logical-database"} );
	}
	if( $indexdata{"logical_database"} ) {
		$db_list = &createDBList( $indexdata{"logical_database"} );
	}

	my $x = "";
	if( $indexdata{"raw-databases"} ) {
		$x = $indexdata{"raw-databases"};
	}	
	if( $indexdata{"raw-databases-snapshot"} ) {
		$x = $x." ".$indexdata{"raw-databases-snapshot"};
	}	

	# create the list of raw dbs to be restored
	my $db_list_raw = "";
	if( $x ) {
		$db_list_raw = &createDBList( $x );
	}
	if( !$db_list && !$db_list_raw ){
		if( $inputs{"databases"} ){
			&printError( "databases ".$inputs{"databases"}." not found in backup\n" );
		}
		&printAndDie( "Nothing to Restore\n" );
	}

	if( $db_list ){
		&doLogicalRestore( $db_list );
	}

	# if there was an error in restoring te logical stuff return;	
	# no point going forward
	if( $zrm_error != 0 ){
		return;
	}
		
	# If there is no raw data to restore again return;
	if( !$indexdata{"innodb-logs"} && !$indexdata{"innodb-data"} && 
		$db_list_raw eq "" ){
		return;
	}

	my $doShutdown = 1;
	if( defined $inputs{"mysql-shutdown"} && $inputs{"mysql-shutdown"} == 0 ){
		&printWarning( "This backup has raw backup data and the --no-mysql-shutdown flag has been specified.\n" );
		&printWarning( "Hence restore will be done without shutting down the MySQL server.\n" );
		&printWarning( "Restoring from a raw backup without shutting down the MySQL server can result in unexpected problems.\n" );
		&printWarning( "After restore is completed, please verify the restored data.\n" );
		$doShutdown = 0;
	}

	if( $doShutdown == 1 ){
		&shutdownMySQL();
	}

	my $f = catfile( $inputs{"source-directory"}, $ZRM_MOUNT_DETAILS );
	if( -f $f ){
		&parseGeneralOptionFile( $f, \%mountDetails );	
		$snapshotPlugin = $mountDetails{"snapshot-plugin"};
		delete $mountDetails{"snapshot-plugin"};
		$snapshotHost = $mountDetails{"host"};
		delete $mountDetails{"host"};
		delete $mountDetails{"backup-set"};
		&mountAll();
	}

	# Restore innodb shared data and logs
	&doInnoDBRestore( "innodb-logs" );
	&doInnoDBRestore( "innodb-data" );

	# Do actual restore 
	if( $db_list_raw ) {
		my @dbs = split( " ", $db_list_raw );
		foreach( @dbs ) {
			if( $inputs{"copy-plugin"} ){
				&copyDataUsingCopyPlugin( $_ );
			}else{
				&createDirAndCopyData( $_ );
			}
		}
	}
}

#executes the incremental restore 
sub doIncRestore()
{
	my $x = &addMySQLParams($MYSQLBINLOG);
	if( $inputs{"start-position"} ) {
		$x = $x." --start-position=".$inputs{"start-position"};
	}	
	if( $inputs{"stop-position"} ) {
		$x = $x." --stop-position=".$inputs{"stop-position"};
	}	
	if( $inputs{"offset"} ) {
		$x = $x." --offset=".$inputs{"offset"};
	}	
	if( $inputs{"start-datetime"} ) {
		$x = $x." --start-datetime=".$inputs{"start-datetime"};
	}
	if( $inputs{"stop-datetime"} ) {
		$x = $x." --stop-datetime=".$inputs{"stop-datetime"};
	}	
	if( $_[0] ) {
		$x = $x." --database=".$_[0];
	}
	my $inp = $_[0];
	
	my $file;
	my $delFile = 0;
	my $msg = "Incremental restore done";
	my $printToScreen = 0;
	if( defined $inputs{"to-file"} && defined $inputs{"from-file"} ){
		&printAndDie( "Both from-file and to-file cannot be specified at the same time\n" );
	}elsif( defined $inputs{"to-file"} ){
		if( $inputs{"to-file"} eq "" ){
			$file = tmpnam();
			$printToScreen = 1;
		}else{
			$file = $inputs{"to-file"};
		}
		if( -e $file && ! -f $file ){
			&printAndDie( "$file does not seem to be a regular file\n" );
		}
		$msg .= " to file=$file";
		# This is so that we do not unnecessarly shutdown mysql server 
		# since we are not touching it in this case.
		if( defined $inputs{"mysql-shutdown"} && $inputs{"mysql-shutdown"} == 1 ){
			$inputs{"mysql-shutdown"} = 0;
		}
	}elsif( defined $inputs{"from-file"} ){
		$file = $inputs{"from-file"};
		if( ! -f $file ){
			&printAndDie( "Unable to find file $file\n" );
		}
		$msg .= " from file=$file";
	}else{
		$file = tmpnam();
		$delFile = 1;
	}
	my $err_msg = "Incremental restore failed\n";
	my $r;	
	if( ! defined $inputs{"from-file"} ){
		my $res="";
		if( $inputs{"bin-logs"} ){
			$res = $inputs{"bin-logs"};
		}else{
			my $t = "\"".$inputs{"source-directory"}."\"";
			$res = catfile( $t, $indexdata{"incremental"} );
		}
		$x = $x." ".$res." >> ".$file;
		if( $verbose ) {
			&printLog( "Restoring incremental to file\n" );
			&printLog( $x."\n" );
		}
		$r = system($x);
		if( $r != 0  ) {
			print "ERROR: $err_msg";
			&printError( $err_msg );
			if( $delFile ){
				unlink( $file );
			}
			return;
		}
	}

	if( ! defined $inputs{"to-file"} ){
        	$x = &addMySQLParams($MYSQL);
        	$x = $x." -e \"source $file;\"";
        	if( $verbose ) {
                	&printLog( "restoring using command ".$x."\n" );
        	}
        	$r = &execCmdAndGetOutput($x);
		if( $delFile ){
			unlink( $file );
		}
		if( !defined $r ) {
			print "ERROR: $err_msg";
			&printError( $err_msg );
			return;
		}
	}
	if( $printToScreen == 1 ){
		print "$msg\n";
	}
	if( $inp ) {
		 $msg = $msg." for database ".$_[0];
	}
	&printLog( $msg."\n" );
}


#Sets up defaults for backup
sub setupDefaults()
{
        $action = "restore";
        $USAGE_STRING = $USAGE_RESTORE_STRING.$USAGE_COMMON_OPTIONS_STRING;
}


sub main()
{
	&setupDefaults();
        &initMySQL(@RESTOREOPT);
        &doRestore();
	&my_exit();
}
&main();
