#!/usr/bin/perl
# IBM_PROLOG_BEGIN_TAG 
# This is an automatically generated prolog. 
#  
# bos720 src/bos/usr/bin/cdat/helpers/get_file.pl 1.2 
#  
# Licensed Materials - Property of IBM 
#  
# COPYRIGHT International Business Machines Corp. 2010,2011 
# All Rights Reserved 
#  
# US Government Users Restricted Rights - Use, duplication or 
# disclosure restricted by GSA ADP Schedule Contract with IBM Corp. 
#  
# IBM_PROLOG_END_TAG 
# @(#)88    1.2  src/bos/usr/bin/cdat/helpers/get_file.pl, cdat, bos720 7/14/11 19:39:25
use warnings;
use strict;
use Net::FTP;
use POSIX qw(ceil floor);	# needed for ceil, floor
use File::Basename;		# needed for dirname, basename

use lib '/usr/lib/cdat/';
use cdat;

#
# Constants.
#
my $CHFS = '/usr/sbin/chfs';
my $DF = '/usr/bin/df';
my $DU = '/usr/bin/du';
my $AWK = '/usr/bin/awk';
my $SSH = '/usr/bin/ssh';
my $SCP = '/usr/bin/scp';
my $SSH_TIMEOUT = 2;
my $MB = 1024 * 1024;
my $EXTRA = 16;	# when extending fs, always allocate an extra 16MB

#
# Globals.
#
my ($growfs, $fsname, $fsfree);

######################################################################
# Function:	usage
# Purpose:	Display usage.
# Tasks:	Print usage and exit.
# Input:	None
# Output:	None
######################################################################
sub usage
{
    print(STDERR "Usage: get_file SourceFile TargetFile\n");
    print(STDERR "       get_file SourceFile ... TargetDirectory\n");
    exit(1);
}

######################################################################
# Function:	ftp_get_file
# Purpose:	Retrieve a file or directory using FTP.
# Tasks:	Retrieve a file or directory using FTP.
# Input:	Net::FTP handle, remote file name, local destination,
#		whether destination is a directory or not
# Output:	0 if success, 1 otherwise
######################################################################
sub ftp_get_file
{
    my ($ftp, $file, $dest, $isdir) = @_;
    my $rc;

    # determine whether remote filename is a file or directory
    if ($ftp->cwd($file)) {
        # remote filename is a directory
        if (!$isdir) {
	    print(STDERR "cannot retrieve directory, destination is a file\n");
	    return 1;
	}
	print("Creating directory $dest/$file\n");
	if (!mkdir("$dest/$file")) {
	    print(STDERR "cannot create directory $dest/$file\n");
	    return 1;
	}
	# Recursively call ourselves for each directory entry
	my $list = $ftp->ls();
	foreach my $elem (@$list) {
	    # call ourserlves recursively
	    if (ftp_get_file($ftp, $elem, "$dest/$file", 1)) {
	        return 1;
	    }
	}
	if (!$ftp->cdup()) {
	    print(STDERR "cannot CDUP from $file\n");
	    return 1;
	}

    } else {
	# remote filename is a file, retrieve it
	if (defined($growfs)) {
	    my $size = $ftp->size($file);
	    if (defined($size)) {
	        reserve_space($size / $MB);
	    }
	}
        if ($isdir) {
	    print("Copying $file to $dest/$file\n");
	    $rc = $ftp->get($file, "$dest/$file");
        } else {
	    print("Copying $file to $dest\n");
	    $rc = $ftp->get($file, $dest);
        }
        if (!$rc) {
	    print(STDERR "could not retrieve file $file\n");
	    return 1;
        }
    }
    return 0;
}

######################################################################
# Function:	ftp_get
# Purpose:	Retrieve remote files/directories using FTP.
# Tasks:	Connect via FTP to the specified remote node using the
#		specified login and retrieve files and directories.
# Input:	user, host name, remote files, local destination
# Output:	None
######################################################################
sub ftp_get
{
    my ($hosts, $user, $hostname, $sources, $dest) = @_;
    my $rc;

    my $passwd = $hosts->{$hostname}{users}{$user};
    if (!defined($passwd)) {
        $passwd = cdat::get_password("Enter password for \"$user\" user: ");
    }
    if (!defined($passwd)) {
	print(STDERR "No password to connect to $user\@$hostname\n");
        exit(1);
    }

    my $isdir = -d $dest;

    my $ftp = Net::FTP->new($hostname);
    if (!defined($ftp)) {
	print(STDERR "Cannot open FTP connection to $hostname\n");
        exit(1);
    }
    $rc = $ftp->login($user, $passwd);
    if (!$rc) {
	print(STDERR "FTP to $hostname failed: ".$ftp->message);
        exit(1);
    }

    # Transfer files in binary mode
    $ftp->binary();

    foreach my $file (@$sources) {
        my $dirname = dirname($file);
        if (!$ftp->cwd($dirname)) {
	    print(STDERR "Cannot CWD to $dirname\n");
	    next;
	}
	last if ftp_get_file($ftp, basename($file), $dest, $isdir);
    }

    # Close FTP session
    $ftp->quit();
}

######################################################################
# Function:	scp_get
# Purpose:	Retrieve remote files/directories using SCP.
# Tasks:	Connect via SSH to the specified remote node using the
#		specified login and retrieve files and directories.
# Input:	user, host name, remote files, local destination
# Output:	None
######################################################################
sub scp_get
{
    my ($hosts, $user, $hostname, $sources, $dest) = @_;

    if (defined($growfs)) {
        my @sizes = `$SSH -i $ENV{HOME}/.ssh/id_rsa -o UserKnownHostsFile=$ENV{HOME}/.ssh/known_hosts -o StrictHostKeyChecking=yes $user\@$hostname 'LANG=C $DU -sm @$sources 2>/dev/null' | $AWK '{print \$1}'`;
	if ($? == 0) {
	    my $size = 0;
	    $size += $_ foreach (@sizes);
	    reserve_space($size);
	}
    }

    system("$SCP -i $ENV{HOME}/.ssh/id_rsa -o UserKnownHostsFile=$ENV{HOME}/.ssh/known_hosts -o StrictHostKeyChecking=yes -q -r $user\@$hostname:'@$sources' '$dest'");
    if ($? != 0) {
	print(STDERR "scp command failed.\n");
        exit(1);
    }
}

######################################################################
# Function:	extendfs
# Purpose:	Extend the size of the destination filesystem.
# Tasks:	Change the size of the destination filesystem using
#		the chfs command.
# Input:	size to add in megabytes
# Output:	1 if successful, 0 otherwise
######################################################################
sub extendfs
{
    my ($size) = @_;

    print("Extending $fsname file system size by $size"."M\n");
    system("$CHFS -a size=+$size"."M '$fsname' >/dev/null 2>&1");
    return ($? == 0);
}

######################################################################
# Function:	get_free_space
# Purpose:	Get the free space on the destination filesystem.
# Tasks:	Get the free space on the destination filesystem in
#		megabytes using the df command.
# Input:	None
# Output:	free space in megabytes or undef if it fails
######################################################################
sub get_free_space
{
    my $size = `LANG=C $DF -m '$fsname' | $AWK '(NR==2){print \$3}'`;
    return undef if ($? != 0);
    chomp($size);
    return $size;
}

######################################################################
# Function:	reserve_space
# Purpose:	Make sure we have space for the specified file size.
# Tasks:	Extend the destination file system if needed.
# Input:	File size in megabytes.
# Output:	None
######################################################################
sub reserve_space
{
    my ($size) = @_;

    for (my $ratio = 1.0; $ratio <= 2.5; $ratio += .5) {
        my $free = get_free_space;
        last if (!defined($free) || ($free >= $size));
	print("File system $fsname has $free".
	    "M free, need at least $size"."M\n");
        last if (!extendfs(ceil($EXTRA + $ratio * ($size - $free))));
    }
}

######################################################################
# Function:	main
# Purpose:	Entry point of the get_file command.
# Tasks:	Retrieve the specified files/directories from the
#		specified remote node using either SCP or FTP.
# Input:	user, host name, node type, remote files,
#		local destination
# Output:	None
######################################################################
sub main
{
    usage() if (@ARGV < 2);

    my $user = $ENV{CDAT_USER};
    my $hostname = $ENV{CDAT_HOST};
    my $type = $ENV{CDAT_TYPE};
    my $dest = pop @ARGV;
    # @ARGV now only contains the remote files

   if (!defined($user) || !defined($hostname) || !defined($type)) {
       printf(STDERR "Bad environment\n");
       exit(1);
   }

    # If there is more than 1 remote file, check that destination
    # exists and is a directory.
    if (@ARGV > 1 && ! -d $dest) {
	print(STDERR "Destination directory \"$dest\" not found.\n");
	exit(1);
    }

    my %hosts;
    cdat::read_host_db(\%hosts);

    $growfs = $ENV{CDAT_GROWFS};
    if (defined($growfs)) {
        # Destination may not exist yet.
	my $file = $dest;
	$file = dirname($file) if (! -e $file);
	# Get name of destination filesystem.
	$fsname = `LANG=C $DF '$file' | $AWK '(NR==2){print \$7}'`;
	if ($? == 0) {
	    chomp($fsname);
	} else {
	    # we won't be able to extend fs, but continue anyway...
	    undef($growfs);
	}
    }

    my $protocol = $hosts{$hostname}{protocol};
    if (!defined($protocol)) {
        if ( -x $SCP && cdat::check_remote_protocol($hostname, 'ssh')) {
            $protocol = "ssh";
        } elsif (cdat::check_remote_protocol($hostname, 'exec')) {
            $protocol = "exec";
        } else {
            $protocol = "telnet";
        }
    }

    print("Retrieving ".join(', ', @ARGV)." from $type $hostname using ");
    if ($protocol eq "ssh") {
	print("scp\n");
        scp_get(\%hosts, $user, $hostname, \@ARGV, $dest);
    } else {
	print("ftp\n");
	ftp_get(\%hosts, $user, $hostname, \@ARGV, $dest);
    }
}

main;