#!/usr/bin/ruby
# Copyright 2006, Daniel Erat <dan-ruby@erat.org>
# All rights reserved.

require 'S3'

NAME = 's3'
MAX_ARGS = 2048
AUTH_FILE = "#{ENV['HOME']}/.s3_auth"

Thread.abort_on_exception = true

class S3Command
  def initialize num_args, short_doc, long_doc, block
    @num_args = num_args
    @short_doc = short_doc
    @long_doc = long_doc
    @block = block
  end
  attr_reader :short_doc, :long_doc

  def S3Command.check_response r
    if not r.http_response.is_a? Net::HTTPSuccess
      raise RuntimeError, "Response was #{r.http_response.code}: " \
        "#{r.http_response.message}"
    end
  end

  def run s3, args
    if (@num_args.is_a? Range and not @num_args.member? args.size) or \
      (not @num_args.is_a? Range and args.size != @num_args)
      $stderr.write "#@num_args argument(s) required " \
        "(got #{args.size})\n"
      return 1
    end
    begin
      @block.call(s3, args)
    rescue RuntimeError
      $stderr.write "#{$!}\n"
      return 1
    end
    0
  end
end


cmds = {}

cmds['mkbucket'] = S3Command.new((1..MAX_ARGS), 'Create bucket', """\
Usage: #{NAME} mkbucket [BUCKET] ...
Create a new bucket.""",
  proc do |s3, a|
    a.each do |arg|
      r = s3.create_bucket arg
      S3Command.check_response r
    end
  end
)

cmds['ls'] = S3Command.new((0..MAX_ARGS), 'List buckets or objects', """\
Usage: #{NAME} ls
       #{NAME} ls [BUCKET][:PREFIX] ...
List all buckets with creation times (if called with no arguments)
or all objects associeted keys matching a pattern in a bucket (if
called with a bucket name or colon-joined bucket and key prefix).""",
  proc do |s3, a|
    if a.size == 0
      r = s3.list_all_my_buckets
      S3Command.check_response r
      r.entries.each do |e|
        time_str =
          Time.parse(e.creation_date).localtime.strftime '%Y-%m-%d %H:%M:%S'
        puts "#{time_str}  #{e.name}"
      end
    else
      first = true
      a.each do |arg|
        bucket, prefix = arg.split(':', 2)
        r = s3.list_bucket bucket, { :prefix => prefix }
        S3Command.check_response r

        # Make a pass through the list to create format strings to line
        # things up nicely.
        owner_max, size_max = 0, 0
        r.entries.each do |e|
          s = e.size.to_s.size
          size_max = s if s > size_max
          s = e.owner.display_name.to_s.size
          owner_max = s if s > owner_max
        end
        owner_fmt = "%#{owner_max}s"
        size_fmt = "%#{size_max}d"

        r.entries.each do |e|
          time_str =
            Time.parse(e.last_modified).localtime.strftime '%Y-%m-%d %H:%M:%S'
          puts "#{owner_fmt % e.owner.display_name} #{size_fmt % e.size} " \
            "#{time_str}  #{bucket}:#{e.key}"
        end
      end
    end
  end
)

cmds['rm'] = S3Command.new((1..MAX_ARGS), 'Remove key or empty bucket', """\
Usage: #{NAME} rm [EMPTY_BUCKET] ...
       #{NAME} rm [BUCKET:KEY] ...
Remove an empty bucket, or an object from a bucket.""",
  proc do |s3, a|
    a.each do |arg|
      bucket, key = arg.split(':', 2)
      if not key
        r = s3.delete_bucket bucket
      else
        r = s3.delete bucket, key
      end
      S3Command.check_response r
    end
  end
)

##
# Does a path look like it could be an S3 bucket:key?
#
# @param path  path to check
#
def maybe_remote? path
  path =~ /^.+:/ ? true : false
end

##
# Get the number of columns displayable by the terminal.  Non-portable.
# From Usenet <20020514160432.A2167@atdesk.com> by Paul Brannan.
#
def num_columns
  tiocgwinsz = 0x5413
  str = [ 0, 0, 0, 0 ].pack("SSSS")
  if $stdin.ioctl(tiocgwinsz, str) >= 0 then
    rows, cols, xpixels, ypixels = str.unpack("SSSS")
    cols
  else
    -1
  end
end

##
# Return a nicely-formatted string for a file size.
#
# @param bytes  file size to format
#
def format_size bytes
  if bytes < 1024
    "#{bytes.to_i}B"
  elsif bytes < 1024 * 1024
    "#{'%.1f' % (bytes.to_f/1024)}KB"
  else
    "#{'%.1f' % (bytes.to_f/(1024*1024))}MB"
  end
end

##
# Return a nicely-formatted string for an interval of time.
#
# @param secs  time interval to format, in seconds
#
def format_time secs
  if secs >= 3600
    s = "#{(secs/3600).to_i}:#{'%02d' % (secs%3600)/60}"
  else
    s = (secs/60).to_i.to_s
  end
  s += ":#{'%02d' % (secs%60)}"
end

##
# Get a string containing the status of an ongoing copy.
#
# @param filename    name of file being copied
# @param size        total size of file being copied, in bytes
# @param copied      number of bytes that have been copied so far
# @param start_time  time at which the copy started
# @param last_msg    last status message from this copy.  if non-nil,
#                    the returned string will be prefixed with
#                    last_msg.length backspaces.
#
def get_copy_status filename, size, copied, start_time, last_msg=nil
  clear = last_msg ? "\b" * last_msg.length : ''

  elapsed = Time.now - start_time
  rate = copied / elapsed
  estimated = size >= 0 ? (size - copied) / rate : 1.0/0
  status = "#{'%.1f' % (100.0*copied/size)}%  "
  status += "#{format_size(copied)}  "
  status += "#{format_size(rate)}/s  "
  if not estimated.infinite?
    if copied != size
      status += "#{format_time estimated} ETA"
    else
      status += "#{format_time elapsed}"
    end
  end

  num_spaces = num_columns - filename.length - status.length - 1
  num_spaces = 2 if num_spaces < 0
  clear + filename + ' ' * num_spaces + status
end

cmds['cp'] = S3Command.new((2..MAX_ARGS),
  'Copy local file to S3 or vice versa', """\
Usage: #{NAME} cp [LOCAL_FILE] ... [BUCKET:KEY]
       #{NAME} cp [BUCKET:KEY] ... [LOCAL_FILE]
Copy a local file to a remote object or vice versa.
If the destination is a directory or bucket then multiple sources
can be specified, and their names will be preserved.""",
  proc do |s3, a|
    dest = a.pop

    if File.exists? a[0] and maybe_remote? dest
      # Local src, remote dest... probably.
      a.each do |arg|
        bucket, key = dest.split(':', 2)
        if key == '.' or key == ''
          arg =~ /([^\/]+$)/
          key = $1
        elsif a.size > 1
          raise RuntimeError, "Must provide bucket as destination when " \
            "giving multiple sources"
        end

        copy_done = false
        File.open(arg) do |f|
          arg =~ /([^\/]+)$/ and filename = $1

          # Start status thread.
          t = Thread.new do
            start = Time.now
            size = File.size arg
            msg = nil
            loop do
              msg = get_copy_status filename, size, f.tell, start, msg
              $stderr.write msg
              break if copy_done
              sleep 1
            end
            $stderr.write "\n"
          end

          r = s3.put bucket, key, f
          copy_done = true
          t.join
          S3Command.check_response r
        end
      end
    else
      # Remote src, local dest

      if false
        new_a = []
        a.each do |arg|
          bucket, key = arg.split ':'
          r = s3.list_bucket bucket, { :prefix => key }
          S3Command.check_response r
          r.entries.each {|e| new_a << "#{bucket}:#{e.key}" }
        end
        a = new_a
      end

      a.each do |arg|
        bucket, key = arg.split ':'

        this_dest = dest
        if File.directory? this_dest
          key =~ /([^\/]+$)/
          this_dest = "#{this_dest}/#{$1}"
        end

        File.open this_dest, File::RDWR|File::CREAT|File::EXCL, 0600 do |f|
          # Get the file size so we can show progress.
          r = s3.list_bucket bucket, { :prefix => key }
          S3Command.check_response r
          e = r.entries.find {|i| i.key == key }
          size = e ? e.size : -1
          copied = 0
          copy_done = false

          # Start status thread.
          t = Thread.new do
            start = Time.now
            msg = nil
            loop do
              msg = get_copy_status key, size, copied, start, msg
              $stderr.write msg
              break if copy_done
              sleep 1
            end
            $stderr.write "\n"
          end

          r = s3.get bucket, key do |h|
            h.read_body do |seg|
              f.write seg
              copied += seg.length
            end
          end
          copy_done = true
          t.join
          S3Command.check_response r
        end
      end
    end
  end
)

cmds['cat'] = S3Command.new((1..MAX_ARGS), 'Print contents of object', """\
Usage: #{NAME} cat [BUCKET:KEY] ...
Print the contents of one or more objects to stdout.""",
  proc do |s3, a|
    a.each do |arg|
      bucket, key = arg.split ':'
      r = s3.get bucket, key do |h|
        h.read_body do |seg|
          puts seg
        end
      end
      S3Command.check_response r
    end
  end
)

cmds['help'] = S3Command.new((0..1), 'Display help', """\
Usage #{NAME} help
      #{NAME} help [COMMAND]
Display a short summary of all commands or detailed help about a
specific command.""",
  proc do |s3, a|
    if a.size == 0 or not cmds[a[0]]
      $stderr.write """\
Usage: #{NAME} [COMMAND] [ARG1] [ARG2] ...
Perform operations using Amazon's S3 online storage service.

Available commands (run \"#{NAME} help [COMMAND]\" for detailed help):\n"""
      cmds.sort.each do |name, cmd|
        $stderr.write "  #{'%-8s' % name}  #{cmd.short_doc}\n"
      end
    else
      $stderr.write "#{cmds[a.shift].long_doc}\n"
    end
  end
)


if $0 == __FILE__
  cmd = ARGV.shift

  if cmd == 'help' or not cmds[cmd]
    $stderr.write "Unknown command \"#{cmd}\"\n" if not cmds[cmd]
    exit cmds['help'].run(nil, ARGV)
  end

  if not File.exists? AUTH_FILE
    $stderr.write """\
  Please create a file named #{AUTH_FILE} containing your
  S3 key ID and secret, separated by whitespace.\n"""
    exit 1
  end

  s3 = S3::AWSAuthConnection.new *(File.open(AUTH_FILE) {|f| f.readline.split })
  exit cmds[cmd].run(s3, ARGV)
end