#! /usr/bin/perl -w
##**************************************************************
##
## Copyright (C) 1990-2007, Condor Team, Computer Sciences Department,
## University of Wisconsin-Madison, WI.
## 
## Licensed under the Apache License, Version 2.0 (the "License"); you
## may not use this file except in compliance with the License.  You may
## obtain a copy of the License at
## 
##    http://www.apache.org/licenses/LICENSE-2.0
## 
## Unless required by applicable law or agreed to in writing, software
## distributed under the License is distributed on an "AS IS" BASIS,
## WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
## See the License for the specific language governing permissions and
## limitations under the License.
##
##**************************************************************


##*******************************************************************
## Cleans up and possibly archives a local directory structure
## for a Condor installation. 
##
## Author: Joe Meehean
## Date: 7/28/05
##*******************************************************************

#***
# Uses
#***
use strict;
use FindBin;
use lib ($FindBin::Bin, "$FindBin::Bin/lib", "$FindBin::Bin/../lib");
use Execute;
use Getopt::Long;
use File::Spec;
use File::Path;
use File::Copy;
use File::Basename;
use Cwd;

#***
# Constant Static Variables
#***
my $USAGE = "Usage: condor_cleanup_local <-localdir local_dir>\n".
    "<-basedir release_dir>".
    "[-archiveonly]\n".
    "[-logarchive log_archive] [-spoolarchive spool_archive]\n".
    "[-execarchive exec_archive]\n".
    "[-configfile configfile]\n".
    "[-nogurl] [-globuslocation globus_location]\n".
    "[-help]\n";

my $CONDOR_CONFIG_VAL = 'condor_config_val';
my $TAR_PROG = 'tar';
my $GZIP_PROG = 'gzip';
my $GURL_CMD = 'globus-url-copy';
my $GURL_FILE_PROTOCOL_HEADER = 'file://';

my $DEFAULT_CONFIG_NAME = 'condor_config.local_dir';

# Acceptable network protocols for fetching a file with globus-url-copy
my %GURL_PROTOCOLS = ( 'http'=>1,
		       'ftp'=>1,
		       'gsiftp'=>1,
		       );

# Environment variables
my $LD_LIBRARY_PATH = 'LD_LIBRARY_PATH';
my $GLOBAL_CONFIG_ENV = 'CONDOR_CONFIG';

# Condor Configuration Attributes
my $LOG_DIR_ATTR = 'LOG';
my $SPOOL_DIR_ATTR = 'SPOOL';
my $EXEC_DIR_ATTR = 'EXECUTE';

#***
# Non-Constant Static Variables
#***
my $local_dir = 0;
my $release_dir = 0;
my $archive_only = 0;
my $log_archive = 0;
my $spool_archive = 0;
my $exec_archive = 0;
my $local_install_config = 0;
my $gurl_flag = 1;
my $globus_location = '/opt/globus';
my $help = 0;

my $condor_config_val_cmd = 0;
my @gurl_args = ();

#***
# Main Function
#***
GetOptions('localdir=s'=>\$local_dir,
	   'basedir=s'=>\$release_dir,
	   'archiveonly'=>\$archive_only,
	   'logarchive=s'=>\$log_archive,
	   'spoolarchive=s'=>\$spool_archive,
	   'execarchive=s'=>\$exec_archive,
	   'configfile=s'=>\$local_install_config,
	   'gurl!'=>\$gurl_flag,
	   'globuslocation=s'=>\$globus_location,
	   'help'=>\$help,
	   );

# Check the usage
die $USAGE if( !$local_dir || !$release_dir || $help );

# Setup globus-url-copy, if necessary
my $archive_flag = $log_archive || $spool_archive || $exec_archive;
if( $archive_flag && $gurl_flag ){
    &SetupGurl();
}

# Default the local installation config if neccessary
if( !$local_install_config ){
     $local_install_config = File::Spec->catpath(0, $local_dir, $DEFAULT_CONFIG_NAME);
}

#Find the log, spool, and execute directories
my($log_dir, $spool_dir, $exec_dir) = &FindDirs();

# Archive the directories if required
&ArchiveDir($log_dir, $log_archive) if( $log_archive );
&ArchiveDir($spool_dir, $spool_archive) if( $spool_archive );
&ArchiveDir($exec_dir, $exec_archive) if( $exec_archive );

# Cleanup the directories if required
if( !$archive_only ){
    rmtree($log_dir, 0, 0) if( -d $log_dir );
    rmtree($spool_dir, 0, 0) if( -d $spool_dir );
    rmtree($exec_dir, 0, 0) if( -d $exec_dir );

    # Remove the configuration file as well
    unlink $local_install_config
	or die "FAILED: Could not unlink $local_install_config: $!";
}

#*********************************************************************
# Finds the log, spool, and execute directories using
# the local installation configuration
#*********************************************************************
sub FindDirs(){
    
    # Find condor_config_val
    $condor_config_val_cmd = &FindBinary($CONDOR_CONFIG_VAL);

    # Ensure the configuration file exist
    if( !-e $local_install_config ){
	die "Could not find the local installation configuration: $local_install_config"; 
    }
    
    # Set up the environment to use the local install config
    $ENV{$GLOBAL_CONFIG_ENV} = $local_install_config;

    # Query for the spool, log, and execute directories
    my $log_dir = &QueryCondorConfig($LOG_DIR_ATTR);
    my $spool_dir = &QueryCondorConfig($SPOOL_DIR_ATTR);
    my $exec_dir = &QueryCondorConfig($EXEC_DIR_ATTR);
    
    return($log_dir, $spool_dir, $exec_dir);
}

#*********************************************************************
# Searches for the given binary name under the release directory.
# Currently it only searches in the release_dir or in release_dir/bin,
# but eventually it should perform a more stringent search.
#*********************************************************************
sub FindBinary(){
    my $binaryName = shift @_;  #Name the parameter
    my $binaryPath = 0;

    # set up the 2 possible paths
    my $flat_path = $release_dir.'/'.$binaryName;
    my $bin_path = $release_dir.'/bin/'.$binaryName;
    my $sbin_path = $release_dir.'/sbin/'.$binaryName;
    
    # determine which to use
    if( -e $flat_path ){
	$binaryPath = $flat_path;
    } elsif( -e $bin_path ){
	$binaryPath = $bin_path;
    } elsif( -e $sbin_path ){
	$binaryPath = $sbin_path;
    } else{
	die "Failed to locate $binaryName in either $flat_path OR $bin_path";
    }

    return($binaryPath);
}

#*********************************************************************
# Queries the condor configuration with the given parameters.
#*********************************************************************
sub QueryCondorConfig(){

    my $returnHash = ExecuteAndCapture($condor_config_val_cmd, @_);
    
    if( $returnHash->{EXIT_VALUE} ){
	#error return undef
	return;
    }
    
    # Check to see if the result is more than one element
    if( @{$returnHash->{OUTPUT}} > 1 ){
	warn "WARNING: compressing results of condor config_val";
    }

    # compress the results into a single element
    my $result = '';
    for( @{$returnHash->{OUTPUT}} ){
	$result = $result.$_."\n";
    }

    #remove the last endline
    chomp($result);

    #return the result
    return $result;
}

#*********************************************************************
# Archives the given directory (tar and gzip) and moves it to
# the given archive name.
#*********************************************************************
sub ArchiveDir(){
    my ($dir_path, $archive_path) = @_;  #Name the parameters

    #0. Ensure the directory actually exists
    die "FAILED: $dir_path doesn't exist" if( !-e $dir_path );
    die "FAILED: $dir_path is not a directory" if( !-d $dir_path );

    #1. Tar it locally
    
    # Store the initial working directory
    my $iwd = cwd;

    # Move to the directory just above the directory to archive
    chdir $dir_path
	or die "FAILED: Cannot chdir to $dir_path:$!";
    my $working_dir = File::Spec->updir();
    chdir $working_dir
	or die "FAILED: Cannot chdir to $working_dir:$!";
    my $dir_name = File::Spec->abs2rel( $dir_path );
    
    # Create a name for the tar
    my $tar_name = $dir_name.'.tar';

    # Create the tar arguments
    my @tar_args = ('cf', $tar_name, $dir_name);

    # tar cf temp.tar dir_path
    !system $TAR_PROG, @tar_args
	or die "FAILED: Failed to tar $dir_name";

    #2.  Gzip it locally
    
    # Create a name for the gzip
    my $gzip_name = $tar_name.'.gz';

    # Create the gzip arguments
    my @gzip_args = ('-n', $tar_name);
    
    # gzip -n temp.tar
    !system $GZIP_PROG, @gzip_args
	or die "FAILED: Failed to gzip $tar_name";

    # Compute the absolute path for the gzip
    my $gzip_path = File::Spec->rel2abs($gzip_name);

    #3. Move back to the initial working directory
    chdir $iwd
	or die "FAILED: Cannot chdir to $iwd:$!";

    #4. Transfer the archive
    &TransferArchive($gzip_path, $archive_path);
}

#*********************************************************************
# Moves the archive to the storage destination
#*********************************************************************
sub TransferArchive(){
    my ($archive_src, $archive_dest) = @_;  #Name the parameters
    
    if( $gurl_flag ){
	# transfer with globus-url-copy
	&TransferArchiveWithGurl($archive_src, $archive_dest);
    } else{

	# Transfer with Perl Move
	move($archive_src, $archive_dest)
	    or die "FAILED: Failed to move $archive_src to $archive_dest: $!";
    }
}

#*********************************************************************
# Moves the archive to the storage destination using globus-url-copy
#*********************************************************************
sub TransferArchiveWithGurl(){
    my ($archive_src, $archive_dest) = @_;  #Name the parameters

    #1. Create the source url
    my $src_url = $GURL_FILE_PROTOCOL_HEADER.$archive_src;

    #2. Create the destination url
    my $dest_url = $archive_dest;

    # See if the destination url
    # a. Has a protocol header?
    # b. That the protocol header is valid?
    # -
    # match a protocol name followed by a colon and two forward slashes
    # Ex: http://files.we.may/want
    # $1 = protocol name
    if( $dest_url =~ m{^(\w+)://} ){
	if( !exists $GURL_PROTOCOLS{$1} ){
	    &DieWithError( "This program does not currently support". 
			   " the given protocol: $1");
	}
    }
    
    # c. Assume file system if there is no protocol header
    else{
	# Ensure the file location is an absolute path (from iwd)
	if( !File::Spec->file_name_is_absolute($dest_url) ){
	    $dest_url = File::Spec->rel2abs($dest_url);
	}
	# Append the file protocol header to the location
	$dest_url = $GURL_FILE_PROTOCOL_HEADER.$dest_url;
    }
    
    #3. Perform the transfer
    #globus-url-copy [options] sourceURL destURL
    push @gurl_args, ($src_url, $dest_url);
    !system $GURL_CMD, @gurl_args
	or die "FAILED: globus-url-copy failed to transfer".
	" $src_url to $dest_url\n";
}

#*********************************************************************
# Adds the globus location to the list of paths to check for 
# dynamically linked libraries.
#*********************************************************************
sub SetupGurl(){
    # Look up the current ld path list
    my $ld_paths = $ENV{$LD_LIBRARY_PATH};

    # Ensure the globus location is an absolute path
    if( !File::Spec->file_name_is_absolute($globus_location) ){
	$globus_location = File::Spec->rel2abs($globus_location);
    }
    
    # Create globus libraries path
    my $globus_lib_path = File::Spec->catpath(0, $globus_location, 'lib');

    # Add globus location to the path list
    if( $ld_paths ){
	# colon-separated
	$ld_paths = $ld_paths.':'.$globus_lib_path;
    } else{
	$ld_paths = $globus_lib_path;
    }

    # Reset the environment variable
    $ENV{$LD_LIBRARY_PATH} = $ld_paths;
}
