Thor task for sftp syncing

sftp.thor.rb

# module: sftp

# A generic Thor module for sftp syncing.
#
#  1. Call `thor sftp:setup` to create config file.
#  2. Edit the config file
#  3. Call `tor sftp:sync` start the sync
#
# Ze end.

require 'rubygems'
gem 'net-sftp', '<2.0.0'
require 'net/sftp'
require 'yaml'
require 'find'

class Sftp < Thor

  desc "upload ENV", "Synchronous a local folder with a remote server as defined in an sftp.yaml file. (ENV defaults to 'default')."
  method_options :verbose => :boolean, :silent => :boolean, :dry => :boolean
  def upload(*args) #env='', opts={}
    env = (args.first.is_a?(Hash)) ? 'default' : args.shift
    opts = args.last
    if File.exists?('sftp.yaml')
      @dry = opts["dry"] || false
      if @dry
        @verbose = true
        @silent = false
      else
        @verbose = opts['verbose'] || false
        @silent = opts['silent'] || false
      end
      begin
        @config = YAML.load( File.open( 'sftp.yaml' ) )
      rescue
        puts "Error: Unable to parse 'sftp.yaml' file."
        exit(1)
      end
      if @config.has_key?(env)
        @settings = @config[env]
        require_keys = %w(host user path) # bare minimums
        if (require_keys - @settings.keys).length == 0
          sync_it(@settings)
        else
          puts "Error: The required key(s) '#{(require_keys - @settings.keys).join("', '")}' are missing."
        end
      else
        puts "Error: There is no definition for the '#{env}' environment."
      end
    else
      puts "Error: 'sftp.yaml' file not found. You may run 'thor sftp:setup' to create one."
    end
  end


  desc "setup", "Setup a folder for remote synchronization."
  def setup
    if File.exists?("sftp.yaml")
      puts "sftp.yaml already exists!"
      exit(1)
    end
    settings =<<-EOYAML
default: # <- Environment can be whatever you'd like. However 'default' is used if you don't specify a target.
  host: YOUR_DOMAIN.com
  user: YOUR_USERNAME
  #pass: YOUR_PASSWORD # OPTIONAL (will prompt if not specified)
  path: www # path on the server to the web root
  checksum_file: .checksums
  folder: . # Local folder to sync from
  ignore: [.svn, .git, .DS_Store, sftp.yaml, .checksums]
    EOYAML
    File.open('sftp.yaml', 'w') do |f|
      f.write( settings )
    end
    `mate sftp.yaml`
    puts "Done."
  end
  
  
  # ===============
  # = Sync Method =
  # ===============
  def sync_it(settings)
    puts_if_verbose "Generating local checksums..."
    checksum_file = settings.fetch('checksum_file', ".checksums")
    ignore_these = settings.fetch('ignore', [])
    local_checksums = checksum_generate_from(settings.fetch('folder', '.'), ignore_these)

    # We use this to upload the latest checksums
    File.open( local_path(checksum_file), 'w' ) do |f|
      f.write local_checksums.to_yaml
    end

    # Fetch password....
    valid_pass = settings.has_key?('pass')
    while !valid_pass
      puts
      print "Specify password for #{ settings['host'] }: "
      password = STDIN.gets.chomp
      valid_pass = (password.length > 0)
      if valid_pass
        settings['pass'] = password
      else
        puts "Password cannot be blank."
      end
    end
    
    puts_if_verbose "Connecting to #{settings['host']}..."
    begin
    
      Net::SFTP.start(settings['host'], settings['user'], settings['pass'], :timeout=>settings.fetch('timeout', 30)) do |sftp|
      
        puts_if_verbose "Fetching checksums...", '.', true
        remote_checksums = {}
        begin
          checksums_src = ''
          sftp.open_handle( remote_path(checksum_file)) do |handle|
            checksums_src = sftp.read(handle)
          end
          remote_checksums = YAML::load(checksums_src) unless checksums_src.nil? or checksums_src==''
        rescue Net::SFTP::Operations::StatusException=>se
          # server's checksum.yml is missing -- don't worry about it, we'll upload it when we're done.
        end
      
        puts_if_verbose "Comparing checksum data..."
        to_upload, to_remove = checksum_diff(local_checksums, remote_checksums)
      
        if to_upload.length > 0 or to_remove.length > 0
          puts_if_verbose "Differences found:"
          to_upload.sort.each {|f| puts_if_verbose " - (Upload) #{f}" }
          to_remove.sort.each {|f| puts_if_verbose " - (Delete) #{f}" }
          if @dry
            puts "Those files would be effected, if this weren't a dry run."
            puts "Done."
            exit(0)
          end
          puts_if_verbose "Beginning sync..."
        
          to_upload.each do |filename|
            begin
              puts_if_verbose " - #{remote_path(filename)}", '+', true

              dir_name = File.dirname(filename)
              dir_segs = dir_name.split('/')
              prog_path = []
              dir_segs.each do |partial_dir|
                begin
                  prog_path << partial_dir
                  sftp.mkdir( remote_path( prog_path ), :permissions=>0755 )
                  puts_if_verbose " + #{remote_path( prog_path )}"
                rescue Net::SFTP::Operations::StatusException=>se
                  # don't worry about it
                end
              end

              sftp.put_file local_path(filename), remote_path(filename)
              sftp.open_handle( remote_path(filename) ) do |handle|
                sftp.fsetstat( handle, :permissions=>0644 )
              end
            rescue  Net::SFTP::Operations::StatusException=>se
              puts_if_verbose " ! Error uploading '#{filename}': #{se.description}"
              puts_if_verbose; puts_if_verbose "Halted execution of upload."
              exit(1)
            end
          end
        
          to_remove.each do |filename|
            begin
              sftp.remove remote_path(filename)
              puts_if_verbose " x #{remote_path(filename)}", 'x', true
            rescue
              puts_if_verbose " ! Error removing '#{filename}': #{$!}"
            end
          end

          begin
            sftp.put_file local_path(checksum_file), remote_path(checksum_file)
            sftp.open_handle( remote_path(checksum_file) ) do |handle|
              sftp.fsetstat( handle, :permissions=>0644 )
            end
          rescue
            puts_if_verbose " ! Error uploading checksum file: #{$!}", '!', true
          end

        
          summary = "\n#{to_upload.length} file(s) uploaded"
          summary += ", #{to_remove.length} files(s) deleted" if to_remove.length > 0
        
          puts summary
        else
          puts "\nNo changes made. The server is up to date!"
        end
      end
    rescue Net::SSH::AuthenticationFailed=>ae
      puts "Authentication failed!"
    end
  end
  
  def remote_path(path)
    [@settings['path'], path].flatten.join('/')
  end
  
  def local_path(path)
    File.join(@settings.fetch('folder', '.'), path)
  end
  
  
  # returns [files_to_upload, files_to_delete]
  def checksum_diff(source={}, target={})
    # look for differences...
    src_files = source.fetch('files', {})
    tgt_files = target.fetch('files', {})
    to_update = []; to_delete = []; to_upload = []

    tgt_files.each do |filename, checksum|
      if src_files.has_key? filename
        to_update << filename unless src_files[filename] == checksum
      else
        to_delete << filename
      end
    end

    to_upload = src_files.keys - tgt_files.keys

    # returns [files_to_upload, files_to_delete]
    [[to_upload, to_update].flatten, to_delete]
  end

  # Create the checksums...
  def checksum_generate_from(cache_dir, exclusions=[])
    checksums = { 'generated_on'=>Time.now, 'files'=>{} }
    Find.find( cache_dir ) do |f|
      next if File.directory?( f )
      relative_path = f.gsub("#{cache_dir}/", '')
      tester = file_fails_test(relative_path, exclusions)
      next if tester
      checksums['files'][relative_path] = Digest::MD5.hexdigest( File.read(f) )
    end
    checksums
  end
  
  def file_fails_test(path, exclusions)
    return true if exclusions.include?(path)
    matches = exclusions.select {|ignore| path.match(ignore) }
    return true if matches.nil?
    return (matches.length > 0) ? true : false
  end
  

  def puts_if_verbose(msg='', alt_char='.', server_progress=false)
    if @verbose
      puts(msg) 
    elsif !@silent
      if server_progress
        STDOUT.print(alt_char)
        STDOUT.flush
      end
    end
  end
  
end