#!/usr/bin/ruby

############################################################
#
# Title: Parallel rsync transfer utility
# Created: Nov 30, 2012
# Author: rshott@sfu.ca
############################################################

############################################################
# CONSTANTS
############################################################

A_TO_Z = ('a'..'z').to_a

############################################################
# Global - Needed for keeping track of child processes
############################################################

$children = Array.new

############################################################
# METHODS
############################################################

def start_parallel_rsync(group)
  for i in group do
    while max_prsyncs_reached?
      sleep 30
    end
    $children << fork { system("/usr/bin/rsync" \
                               " #{ARGV[0]}" \
                               " --include='#{i}*'" \
                               " --include='#{i.upcase}*'" \
                               " --exclude='/*'" \
                               " #{ARGV[1].chomp('/')}/" \
                               " #{ARGV[2].chomp('/')}" \
                               " 1> #{ARGV[3].chomp('/')}/#{i}.log" \
                               " 2> #{ARGV[3].chomp('/')}/#{i}_error.log") 
                        check_for_errors($?, i) }
    Process.detach($children.last)
  end
end

def max_prsyncs_reached?
  return false if $children.length < 4
  limit = Array.new
  $children.each do |x|
    limit << true if pid_exists?(x)
    return true if limit.length == 4
  end
  return false
end

def wait_for_any_pids_running
  $children.each { |p|
    Process.wait(p) if pid_exists?(p) 
  }
end

def non_alpha_character_rsync
  result = system("/usr/bin/rsync" \
                  " #{ARGV[0]}" \
                  " --exclude='[a-z]*'" \
                  " --exclude='[A-Z]*'" \
                  " --exclude='.*'" \
                  " --exclude=' *'" \
                  " #{ARGV[1].chomp('/')}/" \
                  " #{ARGV[2].chomp('/')}" \
                  " 1> #{ARGV[3].chomp('/')}/non-alpha_character.log" \
                  " 2> #{ARGV[3].chomp('/')}/non-alpha_character_error.log")
  if result
    puts "The non-alpha character transfers completed successfully"
  else
    puts "The non-alpha character transfers did not exit with 0 - #{$?} - check its logs"
  end
end

def leading_whitespace_rsync
  result = system("/usr/bin/rsync" \
                  " #{ARGV[0]}" \
                  " --include=' *'" \
                  " --exclude='/*'" \
                  " #{ARGV[1].chomp('/')}/" \
                  " #{ARGV[2].chomp('/')}" \
                  " 1> #{ARGV[3].chomp('/')}/leading_whitespace.log" \
                  " 2> #{ARGV[3].chomp('/')}/leading_whitespace_error.log")
  if result
    puts "The leading whitespace transfers completed successfully"
  else
    puts "The leading whitespace transfers did not exit with 0 - #{$?} - check its logs"
  end
end

def hidden_rsync
  result = system("/usr/bin/rsync" \
                  " #{ARGV[0]}" \
                  " --include='.*'" \
                  " --exclude='/*'" \
                  " #{ARGV[1].chomp('/')}/" \
                  " #{ARGV[2].chomp('/')}" \
                  " 1> #{ARGV[3].chomp('/')}/hidden_folders.log" \
                  " 2> #{ARGV[3].chomp('/')}/hidden_folders_error.log")
  if result
    puts "The hidden folders transfer completed successfully"
  else
    puts "The hidden folders transfer did not exit with 0 - #{$?} - check its logs"
  end
end

def check_for_errors(exitcode, letter)
    unless exitcode != 0
      puts "The '#{letter}*' & '#{letter.upcase}*' transfers completed successfully"
    else
      puts "The '#{letter}*' & '#{letter.upcase}*' transfers did not exit with 0 - #{exitcode} - check its logs"
    end
end

def pid_exists?(pid)
  begin
    Process.kill(0, pid)
    return true
  rescue
    return false
  end
end

def ctrl_c
  puts " ctrl+c caught"
  $children.each { |c| if pid_exists?(c)
                         puts "Killing child process #{c}" 
                         Process.kill('KILL', c) 
                       end }
  exit 1
end

############################################################
# MAIN
############################################################

if ARGV[0].eql?('-h') or ARGV[0].eql?('--help')
  puts "Version 2.0\n" \
       "Author: Riley Shott (rshott@sfu.ca)\n" \
       "https://github.com/Ginja/Admin_Scripts/blob/master/prsync_transfer\n\n" \
       "Capabilities: This utility will rsync all the contents of one folder into another but perform some of the transfer in parallel.\n" \
       "The alpha character rsync will be done in batches of 4. New alpha character transfers will start when old transfers finish (give or take 30 seconds).\n" \
       "All other jobs are done serially. You will receive a report at the end of the command letting you know which transfers experienced errors.\n\n" \
       "Job options & order:\n" \
       "  1.  alpha characters (i.e. Aa-Zz) \n" \
       "  2.  non-alpha characters (ex: 1_folder, _file)\n" \
       "  3.  files and folders with a leading whitespace (ex: ' somefoldername')\n" \
       "  4.  hidden files and folders (ex: .itsasecret) \n\n" \
       "Usage: #{__FILE__} <-rsync_options> <source> <target> <log_location> (jobs_to_run EX: 1-3)\n" \
       "By default, this utility will run jobs 1-4. You can specify the transfers you want to do:\n" \
       "  ./prsync_transfer -avP ~/source ~/target ~/log 3-4\n" \
       "  ./prsync_transfer -avP ~/source ~/target ~/log 1-2\n" \
       "You can also specify your own includes and excludes (mind the quotes):\n" \
       "  ./prsync_transfer \"-avP --exclude=Caches*\" ~/source ~/target ~/log 1\n"
  exit 0
end

# Fail early, fail hard
begin
  message = Array.new
  message << "Usage: #{__FILE__} <-rsync_options> <source> <target> <log_location> (jobs_to_run EX: 1-3)" unless ARGV.length == 4 or ARGV.length == 5 
  message << "Error: Could not find rsync at /usr/bin/rsync" unless File.exists?('/usr/bin/rsync')
  message << "Error: Source destination does not exist" unless File.exists?(ARGV[1])
  message << "Error: Target destination does not exist" unless File.exists?(ARGV[2])
  message << "Error: Log location does not exist" unless File.exists?(ARGV[3])
  message << "Error: Source destination must be a folder" unless File.directory?(ARGV[1])
  message << "Error: Target desintation must be a folder" unless File.directory?(ARGV[2])
  message << "Error: Log location must be a folder" unless File.directory?(ARGV[3])
  raise unless message.empty?
rescue
  message.each {|m| puts m} 
  exit 1
end

unless ARGV.length == 4
  if ARGV[4].to_s.length == 1
    rsyncs_to_run = ARGV[4] 
  else
    rsyncs_to_run = ARGV[4].gsub(/[-]/,' ').split
    rsyncs_to_run = (rsyncs_to_run[0]..rsyncs_to_run[1]).to_a
  end
else
  rsyncs_to_run = ('1'..'4').to_a
end

rsyncs_to_run.each {|n|
  case n
  when '1'
    trap("SIGINT") { ctrl_c }
    start_parallel_rsync(A_TO_Z)
  when '2'
    non_alpha_character_rsync
  when '3'
    leading_whitespace_rsync
  when '4'
    hidden_rsync
  else
    puts "Did not recognize the following job option - #{n}\n" \
    "Job options:\n" \
       "  1.  alpha characters\n" \
       "  2.  non-alpha characters\n" \
       "  3.  files and folders with a leading whitespace\n" \
       "  4.  hidden files and folders\n"
  end
}
wait_for_any_pids_running
