#!/bin/ksh
#
# CDDL HEADER START
#
# The contents of this file are subject to the terms of the
# Common Development and Distribution License (the "License").
# You may not use this file except in compliance with the License.
#
# You can obtain a copy of the license at usr/src/OPENSOLARIS.LICENSE
# or http://www.opensolaris.org/os/licensing.
# See the License for the specific language governing permissions
# and limitations under the License.
#
# When distributing Covered Code, include this CDDL HEADER in each
# file and include the License file at usr/src/OPENSOLARIS.LICENSE.
# If applicable, add the following below this CDDL HEADER, with the
# fields enclosed by brackets "[]" replaced with your own identifying
# information: Portions Copyright [yyyy] [name of copyright owner]
#
# CDDL HEADER END
#
# Copyright 2007 Jan Pechanec.  All rights reserved.
# Use is subject to license terms.
#

#
# This is basically just a wrapper script around SSH. The fact that
# SSH connection is single threaded means that all those virtual
# CPU's on Niagara machines won't help us -- only 1 CPU is actually
# working. So, the idea is to fire up several SSH connections, with
# every such a session transfering a piece of the file. No multiple
# files or recursive copy is supported for now; patches are welcome.
# This is more or less a proof of concept, I personally think that
# this might be a nice performance enhancement to rsync(1).
#
# Anyway, don't expect miracles -- on a 1Gb network or above you
# quickly hit disk I/O bottleneck. I was able to trasfer up to 20MB
# of data over localhost on our T1000 with 10 parallel connections.
# That's 4x speed up against single threaded scp(1). If you want to
# test the transfer speed only, use -t option that uses fake data so
# no disk I/O is performed.
#
# The script was developed on Nevada (Solaris 11) and only partially
# tested. You use it at your own risk. I shall not be liable for any
# damages resulting from you using this software.
#
# If there is any new version you should be able to find it on
# http://blogs.sun.com/janp
#


# default number of parallel connections. Use -c option to change.
nconns=6
# block size for dd(1); should be enough but you can experiment with
# -b option
bs=65536
# DisableBanner=yes instead of LogLevel would be better but that's a new option
# in OpenSolaris
ssh_opts="-T -x -o PreferredAuthentications=publickey -o LogLevel=error"
# min file size in bytes to transfer in parallel, let's say it's 20MB. If less
# we transfer single threaded.
minsize=20971520
# digest command. OpenSSL was chosen for portability.
dgst="openssl sha1"

# was experimenting with various settings
FILE=file
# we can't use any special options if we want to be reasonably able to run on
# many systems. So, this is the best guess what can be a regular file when using
# file(1) without any options. It would be great if S10 had stat(1) command :-/
# Changed by Eli
#filepat=":.*file|:.*text|:.*data"
filepat=":.*"
# let's believe this covers many systems, including Solaris, FreeBSD and Linux
# distros
dirpat=":.*(sticky )?directory$"
# this is really nasty but that's the closest to portability I can get :-/
nofilepat=": No such file or directory$|: cannot open "

# we might have many processes running on the background. If the use ^C this
# script we should try kill all of them.
handler ()
{
	[ -n "$called" ] && return

	called=1
	printf \
	    "\nSIGINT or SIGTERM caught, trying to kill all my background jobs"
	printf "\ngive me some time to do it...\n\n"
	# it looks like ksh(1) puts all background processes, when running as
	# non-interactive shell, into one process group -- the group of the
	# script itself. So, we just kill $$ and that should kill all local
	# processes (and thus remote, too).
	pkill -g $$ 2>/dev/null
	[ $? != 0 ] && echo "pkill failed" >&2
	sleep 3
	echo "hopefully done"
}

# if $1 is "INFO" we take it as the default info log level
dbg ()
{
	[ -n "$quiet" ] || [ -z "$debug" -a "$1" != "INFO" ] && return
	# shift to remove "INFO"
	[ "$1" != "INFO" ] && label="DBG: " || shift
	echo "$label$@"
}

# check if the $3 string is on line $2 of data in $1
checkline ()
{
	echo "$1" | head -$2 | tail -1 | egrep -e "$3" >/dev/null
}

usage ()
{
	cat <<END
pscp - parallel data transfer based on SSH
Version 1.0
Usage: pscp user@host:file1 user@host:[dir | file2]
Options:
	-b	block size in bytes (default: $bs)
	-c	number of parallel connections (default: $nconns)
	-d	debug mode
	-h	this help
	-q	quiet mode; print only errors
	-s	SHA1 verification of the transfered file (use for debug only)
	-t 	test mode; no reads, no writes, transfer only

Notes:
	- doesn't support local-to-local and remote-to-remote transfers
	- you might consider to increase sshd's MaxStartups on server side
END
	exit 1
}

while getopts 'b:c:dhpqst' c; do
        case ${c} in
        b)
                bs="$OPTARG"
		[ "$nconns" -lt 1 -o "$nconns" -gt 4194304 ] && \
		    echo "value for -b doesn't seem sane" >&2 && exit 1
                ;;
        c)
                nconns=$OPTARG
		[ "$nconns" -lt 1 -o "$nconns" -gt 128 ] && \
		    echo "value for -c doesn't seem sane" >&2 && exit 1
                ;;
        d)
                debug="yes"
                ;;
        q)
                quiet="yes"
                ;;
        s)
                sha1="yes"
                ;;
        t)
                test_only="yes"
                ;;
        h|*)
                usage
                ;;
        esac
done

shift `expr $OPTIND - 1`
[ $# -ne 2 ] && usage

[ -n "$test_only" -a -n "$sha1" ] && \
    echo "can't use both -t and -s" >&2 && exit 1

trap handler INT TERM

ifile=$1
ofile=$2

dbg "process ID is $$"
dbg "number of file specifications: $#"

# for both file specifications, find out if they reference a remote or local
# file/directory
expr $1 : ".*:.*" >/dev/null && download=1
expr $2 : ".*:.*" >/dev/null && upload=1

# we are willing to do local-to-remote or remote-to-local transfers only
[ -n "$download" -a -n "$upload" ] && \
    echo "can't do remote-to-remote transfers" >&2 && exit 1
[ -z "$download" -a -z "$upload" ] && \
    echo "can't do local-to-local transfers" >&2 && exit 1

dbg "download/upload [$download/$upload]"
	
base_ifile=`basename $ifile`

if [ -n "$upload" ]; then
	[ -r "$1" ] ||
	    { echo "$1 doesn't exist or unreadable" >&2; exit 1; }
	$FILE $1 2>/dev/null | egrep "$filepat" >/dev/null || \
	    { echo "$1 is not a regular file" >&2; exit 1; }
	ofile=`echo $ofile | cut -f2 -d:`
	# get local file size
	size=`ls -ld "$1" 2>/dev/null`
	[ $? -ne 0 ] && echo "can't stat $1" >&2 && exit 1
	size=`echo $size | awk '{ print \$5 }'`
	host=`echo $2 | cut -f1 -d:`

	# if in upload mode, we have to check whether the remote file spec is
	# a directory or not. We also explicitly check by simple echo that we
	# can log in remotely.
	rem_stat="$FILE $ofile 2>&1; $FILE `dirname $ofile` 2>&1"
	is_dir=`ssh $ssh_opts $host "echo hello; $rem_stat" 2>/dev/null`
	ret=$?
	dbg "rem_stat ssh returned '$is_dir'"
	echo "$is_dir" | grep ^hello >/dev/null

	# was remote echo successful at all?
	if [ $? -eq 0 ]; then
		# does remote parent file/directory exist?
		if ! checkline "$is_dir" 3 "$nofilepat"; then
			dbg "$ofile's parent does exist"
			# is the parent really a directory?
			if checkline "$is_dir" 3 "$dirpat"; then
				# is target a dir or a file?
			    	if ! checkline "$is_dir" 2 "$nofilepat"; then
					if checkline "$is_dir" 2 "$dirpat"; then
						[ -z "$test_only" ] && \
						    ofile=$ofile/$base_ifile
						dbg "uploading to directory"
					elif checkline "$is_dir" 2 "$filepat"; then
						dbg "will overwrite existing file"
					else
						echo \
						    "$ofile not a regular file" >&2
						exit 1
					fi
				fi
			else
				# remote parent dir is not a dir
				echo "invalid path $ofile" >&2 && exit 1
			fi
		else
			# remote parent dir doesn't exist
			echo "invalid path $ofile" >&2 && exit 1
		fi
	else
		echo "can't login without password" >&2 && exit 1
	fi

	dbg "will upload to $ofile"

else
	# download
	host=`echo $1 | cut -f1 -d:`
	ifile=`echo $1 | cut -f2 -d:`
	dbg "file to download is $ifile"

	# does local parent file/dir exist and is it a directory?
	stat="$FILE `dirname $2`"
	is_dir=`$stat 2>/dev/null`
	if ! checkline "$is_dir" 1 "$nofilepat"; then
		if checkline "$is_dir" 1 "$dirpat"; then
			: everything is ok
		else
			echo "invalid path $2" >&2 && exit 1
		fi
	else
		echo "invalid path $2" >&2 && exit 1
	fi

	stat="$FILE $2"
	is_dir=`$stat 2>/dev/null`
	if ! checkline "$is_dir" 1 "$nofilepat"; then
		if checkline "$is_dir" 1 "$dirpat"; then
			[ -z "$test_only" ] && ofile=$ofile/$base_ifile
			dbg "downloading to directory"
			dbg "...to the file $ofile"
		fi
	fi


	# we need to know the remote file size and type
	rem_stat="ls -ld $ifile 2>&1 | awk '{ print \$5 }'; $FILE $ifile 2>&1"
	dbg "rem_stat is '$rem_stat'"
	dbg "will run: ssh $ssh_opts $host \"echo hello; $rem_stat\""
	size=`ssh $ssh_opts $host "echo hello; $rem_stat" 2>/dev/null`
	ret=$?
	dbg "rem_stat ssh returned '$size'"
	echo "$size" | grep ^hello >/dev/null
	if [ $? -eq 0 ]; then
		# does the remote file/dir exist?
		if ! checkline "$size" 3 "$nofilepat"; then
			checkline "$size" 3 "$filepat"
			[ $? -eq 1 ] && \
			    echo "remote $ifile is not a regular file" >&2 && \
			    exit 1
			# we have size
			size=`echo "$size" | head -2 | tail -1`
		else
			echo "can't read the file for download" >&2 && exit 1
		fi
	else
		echo "can't login without password" >&2 && exit 1
	fi
fi

dbg "host is $host"
dbg "stat output is $size bytes"

# if in test mode only, we just transfer that many zeroes to /dev/null as the
# filesize. No disk read/write I/O is performed. This tests actuall transfer
# only, avoiding possible disk I/O bottlenecks.
if [ -n "$test_only" ]; then
	dbg "test mode run only, fake data is copied"
	ofile=/dev/null
	ifile=/dev/zero
fi

# let's find out how many blocks we transfer over every connection
nblocks=`expr $size / $bs + 1`
# how many blocks per connection
chunk=`expr $nblocks / $nconns + 1`
dbg "chunk size is $chunk blocks, block-size $bs"

if [ $size -lt $minsize ]; then
	echo "file too small ($size bytes), parallel copy is not worth it"
	echo "setting parallelism to 1"
	nconns=1
	chunk=$nblocks
fi

dbg INFO "initialization done"
dbg INFO "firing up parallel transfer with $nconns connection(s) NOW"

# first python code, didn't work in 2.3 since there is no subprocess module
#"dd if='$ifile' count=$chunk skip=$skip bs=$bs 2>/dev/null" | \
#python -c "import os; import subprocess; \
#fd = os.open(\"$ofile\", os.O_CREAT | os.O_WRONLY, 0666); \
#os.lseek(fd, $seek, 0); \
#subprocess.call([\"cat\"], stdout = fd)" &

i=0
for j in `yes | head -$nconns`; do

	skip=`expr $i '*' $chunk`
	seek=`expr $skip '*' $bs`

	# why we use Python code? dd(1) have no option to always read up to the
	# whole block. And since we are behind the socket, we can't be sure that
	# you can always have a full block. In fact, you couldn't rely on such
	# transfer even if you block size was below MTU since you never know
	# what can happen.
	#
	# in my opinion, there is no command in the base system that could do
	# what we need here. So, a simple Python code that lseek(2) on the file
	# and then fork cat(1) is a simple and efficient workaround.
	if [ -n "$upload" ]; then
		dbg "ssh upload chunk#$i offset=$skip count=$chunk bs=$bs"
		dd if="$ifile" count=$chunk skip=$skip bs=$bs 2>/dev/null | \
		    ssh $ssh_opts $host \
		    "python -c 'import os; \
		     fd = os.open(\"$ofile\", os.O_CREAT | os.O_WRONLY, 0666); \
		     os.lseek(fd, $seek, 0); os.dup2(fd, 1); os.system(\"cat\")'" &
	else
		dbg "ssh download chunk#$i offset=$skip count=$chunk bs=$bs"
		ssh $ssh_opts $host \
		    "dd if='$ifile' count=$chunk skip=$skip bs=$bs 2>/dev/null" | \
		     python -c "import os; \
		     fd = os.open(\"$ofile\", os.O_CREAT | os.O_WRONLY, 0666); \
		     os.lseek(fd, $seek, 0); os.dup2(fd, 1); os.system(\"cat\")" &
	fi

	i=`expr $i + 1`
done

if [ -n "$sha1" ]; then
	dbg "computing checksum of the original file"
	if [ -n "$upload" ]; then
		sum_local=`$dgst "$ifile" | sed -e 's/.*= //g'`
	else
		sum_remote=`ssh $ssh_opts $host $dgst "$ifile" | \
		    sed -e 's/.*= //g'`
	fi
	dbg "checksum computation done"
fi

dbg INFO "waiting for all ssh clients to finish"
wait

# now do the checksum on the transfered file
if [ -n "$sha1" ]; then
	dbg INFO "computing SHA-1 checksum over the transfered file"
	if [ -n "$upload" ]; then
		sum_remote=`ssh $ssh_opts $host $dgst "$ofile" | \
		    sed -e 's/.*= //g'`
	else
		sum_local=`$dgst "$ofile" | sed -e 's/.*= //g'`
	fi

	dbg "local sum: $sum_local"
	dbg "remote sum: $sum_remote"
	if [ "$sum_local" != "$sum_remote" ]; then
		echo "checksums mismatch, transfer FAILED" >&2
		exit 1
	fi
	dbg INFO "OK, checksum match"
fi

dbg INFO "all finished"

