ton_upload

#!/usr/bin/env ruby -KU -rubygems

# The TON API allows OAuth 1.0A or with whitelisting OAuth 2 (application-only auth)
# Default is OAuth 1.0A for all applications accessing the TON API.
#
# Documentation for doing application-only auth using the Twitter API:
#   https://dev.twitter.com/docs/auth/3-legged-authorization
#   https://dev.twitter.com/docs/auth/application-only-auth and
#   https://dev.twitter.com/docs/api/1.1/post/oauth2/token
#
# This script uses the .twurlrc file managed by https://github.com/twitter/twurl

require 'bundler/setup'

require 'digest/md5'
require 'json'
require 'logger'
require 'mime/types'
require 'rest-client'
require 'twurl'
require 'time'
require 'optparse'
require 'oauth'

RestClient.log = Logger.new(STDERR)

API_DOMAIN = "ton.twitter.com"

MISSING_TWURL_CONFIG = "[ERROR] Your ~/.twurlrc could not be found. " +
  "Please install and setup twurl then try again: \n\n" +
  "  gem install twurl\n" +
  "  twurl authorize --consumer-key key --consumer-secret secret\n\n"

TWURL_PROFILE = begin
  default = Twurl::RCFile.load["configuration"]["default_profile"]
  abort MISSING_TWURL_CONFIG unless default
  profile = Twurl::OAuthClient.load_client_for_username_and_consumer_key(default[0], default[1])
end

CONSUMER_KEY    = TWURL_PROFILE.consumer_key
CONSUMER_SECRET = TWURL_PROFILE.consumer_secret
USER_TOKEN      = TWURL_PROFILE.token
USER_SECRET     = TWURL_PROFILE.secret

MISSING_SECRET_ERROR = "[ERROR] Your ~/.twurlrc does not contain a consumer secret " +
  "for the test application (#{CONSUMER_KEY}). Follow the twurl setup " +
  "instructions and try again."

MISSING_TOKEN_ERROR = "[ERROR] Your ~/.twurlrc does not contain a user token " +
  "and secret for the application (#{CONSUMER_KEY}). Follow the twurl setup " +
  "instructions and try again."

# This can go upto 1024*1024*64
SINGLE_CHUNK_UPLOAD_LIMIT_IN_BYTES = 1024*1024*8

DOWNLOAD_FILE = 'sample_download_file'

EXPIRE_TIME = (Time.now + 10*24*60*60).httpdate

OAUTH_METHODS = {
  :post => Net::HTTP::Post,
  :get => Net::HTTP::Get,
  :put => Net::HTTP::Put,
  :delete => Net::HTTP::Delete,
  :options => Net::HTTP::Options,
  :head => Net::HTTP::Head,
  :copy => Net::HTTP::Copy
}

OPERATIONS = [
  :upload,
  :download,
  :verify_upload
]

module AuthenticationMode
  USER = "1"
  APP = "2"
end

def get_basic_authorization_header
  # Used for app-auth only
  credentials = "#{CONSUMER_KEY}:#{CONSUMER_SECRET}"
  "Basic #{Base64.encode64(credentials).gsub("\n", '')}"
end

# Exchange the consumer key and secret for a bearer token (used for
# application-only authentication)
# TODO: when twurl supports app-auth, use twurl rather than RestClient
def get_bearer_token
  # Used for app-auth only
  rsp = RestClient::Request.execute(
    :method => :post,
    :url => "https://api.twitter.com/oauth2/token",
    :headers => {"Authorization" => get_basic_authorization_header},
    :payload => {:grant_type => 'client_credentials'}
  )

  JSON.parse(rsp.body)["access_token"]
end

def single_chunk_upload(file_path, auth_profile, api_domain, file_size, content_type, bucket_name, trace)

  headers = {
    'Content-Type' => content_type,
    'Content-Length' => file_size,
    'X-TON-Expires' => EXPIRE_TIME
  }

  url = "https://#{api_domain}/1.1/ton/bucket/#{bucket_name}"
  body = File.read(file_path)

  if auth_profile[:authmode] == AuthenticationMode::USER
    # User-auth
    rsp_headers = user_auth_request(headers, :post, auth_profile, trace, api_domain, url, body)
    rsp_headers[:location]
  else
    # App-auth
    rsp_headers = app_auth_request(headers, :post, trace, auth_profile, url, body)
    rsp_headers[:location]
  end
end

def initiate_multi_part_upload(auth_profile, api_domain, file_size, content_type, bucket_name, trace)

  headers = {
    'X-TON-Content-Type' => content_type,
    'X-TON-Content-Length' => file_size,
    'X-TON-Expires' => EXPIRE_TIME,
    'Content-Length' => 0,
    'Content-Type' => content_type
  }

  url = "https://#{api_domain}/1.1/ton/bucket/#{bucket_name}?resumable=true"
  body = nil

  if auth_profile[:authmode] == AuthenticationMode::USER
    # User-auth
    rsp_headers = user_auth_request(headers, :post, auth_profile, trace, api_domain, url, body)
  else
    # App-auth
    rsp_headers = app_auth_request(headers, :post, auth_profile, trace, url, body)
  end
end

def read_file_in_chunks(file_path, chunk_size_in_bytes, file_size)
  File.open(file_path) do|file|
    total_length_read = 0
    i = 0
    until file.eof?
      start_index = total_length_read
      bytes_to_read = [chunk_size_in_bytes, file_size - total_length_read].min
      bytes = file.read(bytes_to_read)
      total_length_read += bytes_to_read
      i += 1
      yield i, bytes, start_index, bytes_to_read, total_length_read
    end
  end
end

def multi_chunk_upload(file_path, auth_profile, api_domain, file_size, content_type, bucket_name, trace)
  puts "Uploading file in multiple chunks..."
  rsp_headers = initiate_multi_part_upload(auth_profile, api_domain, file_size, content_type, bucket_name, trace)

  chunk_size = rsp_headers[:x_ton_min_chunk_size].to_i

  url = "https://#{api_domain}#{rsp_headers[:location]}"

  read_file_in_chunks(file_path, chunk_size, file_size) do |i, buffer, start_index, bytes_to_read, total_length_read|
    puts "Uploading chunk #{i} (#{start_index}-#{total_length_read-1}/#{file_size})..."

    headers = {
      "Content-Type" => content_type,
      "Content-Length" => bytes_to_read,
      "Content-Range" => "bytes #{start_index}-#{total_length_read-1}/#{file_size}"
    }

    if auth_profile[:authmode] == AuthenticationMode::USER
      # User-auth
      chunk_rsp_headers = user_auth_request(headers, :put, auth_profile, trace, api_domain, url, buffer)
    else
      # App-auth

      # TODO move to using app_auth_request
      RestClient::Request.execute(
        :method => :put,
        :url => "https://#{api_domain}#{rsp_headers[:location]}",
        :headers => {
          "Authorization" => "Bearer #{auth_profile[:bearer_token]}",
          "Content-Type" => content_type,
          "Content-Length" => bytes_to_read,
          "Content-Range" => "bytes #{start_index}-#{total_length_read-1}/#{file_size}"
        },
        :payload => buffer
      ){|response, request, result, &block|
        case response.code
        when 308
        else response.return!(request, result, &block)
        end
      }
    end
  end
  rsp_headers[:location]
end

# Use the provided bearer token to upload a file on disk to TON.
# Files that are less than 64MB can be uploaded all at once, though you might break them into smaller chunks as well;
# Files that are larger than 64MB must be uploaded in some smaller chunk sizes.
def upload_to_ton(file_path, auth_profile, api_domain, bucket_name, trace)
  file_size = File.size(file_path)
  content_type = MIME::Types.type_for(file_path).first.to_s
  content_type = 'text/plain' if content_type.empty?

  puts "Uploading #{file_path} (#{pretty_size(file_size)}) as #{content_type} via TON API"
  @rsp_location =
    if file_size < SINGLE_CHUNK_UPLOAD_LIMIT_IN_BYTES
      single_chunk_upload(file_path, auth_profile, api_domain, file_size, content_type, bucket_name, trace)
    else
      multi_chunk_upload(file_path, auth_profile, api_domain, file_size, content_type, bucket_name, trace)
    end
  puts "File stored at #{@rsp_location}"
  @rsp_location
end

# Download a file from TON
#
# Can be used to test that a file was uploaded correctly
def get_from_ton(file_key, auth_profile, api_domain, bucket_name, trace)
  # TODO update to support user-auth
  # TODO trace not implemented without user-auth
  puts "Downloading the file"
  rsp = RestClient::Request.execute(
    :method => :get,
    :url => "https://#{api_domain}/1.1/ton/bucket/#{bucket_name}/#{file_key}",
    :headers => {"Authorization" => "Bearer #{auth_profile[:bearer_token]}"}
  )
  puts "Successfully downloaded"
  File.open(DOWNLOAD_FILE, 'w') { |file| file.write(rsp) }
end

def pretty_size(bytes)
  kb = bytes / 1024
  mb = bytes / 1024 ** 2
  mb > 0 ? "#{mb} MB" : "#{kb} KB"
end

def app_auth_request(headers, method, auth_profile, url, body)
  # App-auth
  # For docs on RestClient::Request.execute, see
  # https://github.com/rest-client/rest-client/blob/master/lib/restclient/request.rb

  headers['Authorization'] = "Bearer #{auth_profile[:bearer_token]}"

  rsp = RestClient::Request.execute(
    :method => method,
    :url => url,
    :headers => headers,
    :payload => body
  )
  rsp.headers
end

def user_auth_request(headers, method, auth_profile, trace, api_domain, url, body)
  # User-auth
  request_class = OAUTH_METHODS.fetch(method)
  request = request_class.new(url)

  headers.each {|k, v| request[k] = v }

  request.body = body

  consumer = OAuth::Consumer.new(
    auth_profile[:consumer_key],
    auth_profile[:consumer_secret],
    :site => "https://#{api_domain}")
  consumer.http.use_ssl = true
  consumer.http.set_debug_output(STDERR) if trace
  consumer.http.verify_mode = OpenSSL::SSL::VERIFY_NONE

  access_token = OAuth::AccessToken.new(consumer, auth_profile[:token], auth_profile[:secret])

  request.oauth!(consumer.http, consumer, access_token)
  rsp = consumer.http.request(request)

  rsp_headers = {}

  # Normalize response headers for compatibility with requestclient
  rsp.each {|k, v| rsp_headers[k.gsub(/\-/, "_").to_sym] = v }
  rsp_headers
end

def parse_args(args)
  options = {:file_path => nil}

  # Options
  # --mode | -m upload|download|verify_upload
  # --file | -f <file_path> path of file to be uploaded
  # --bucket | -b <bucket_name>
  # --app-auth
  parser = OptionParser.new do |opts|
    opts.banner = "\nUsage: #{$0} [options]"

    opts.on('-h', '--help', "Display usage")

    # Mode
    opts.on('-m', '--mode <mode>', "One of #{OPERATIONS.join(', ')}") do |mode|
      normalized_mode = mode.downcase.to_sym
      abort "ERROR: Unknown mode '#{mode}'" unless OPERATIONS.include?(normalized_mode)
      options[:mode] = normalized_mode
    end

    # Bucket
    opts.on('-b', '--bucket <bucket_name>', 'Bucket name for upload (required)') do |bucket_name|
      abort "ERROR: invalid bucket_name" unless (/^[a-z_0-9]+$/.match(bucket_name))
      options[:bucket_name] = bucket_name
    end

    # File path for upload
    opts.on('-f', '--file <file_path>', 'Path of the file to be uploaded') do |file_path|
      if File.file?(file_path)
        options[:file_path] = file_path
      else
        abort "ERROR: File not found"
      end
    end

    # App-auth (by default use user auth)
    opts.on('--app-auth', 'Use app-auth rather than user-auth') do |app_auth|
      options[:app_auth] = app_auth
    end

    opts.on('-t', '--trace', 'Trace request/response traffic') do |trace|
      options[:trace] = trace
    end
  end

  parser.parse!(args)

  warn "Missing required argument --mode" if options[:mode].nil?
  warn "Missing required argument --bucket" if options[:bucket_name].nil?

  if options[:bucket_name].nil? || options[:mode].nil?
    abort parser.to_s
  end

  options
end

def execute(mode, trace, auth_profile, file_path, bucket_name)
  api_domain = API_DOMAIN

  if mode == :upload
    upload_to_ton(file_path, auth_profile, api_domain, bucket_name, trace)
  elsif mode == :verify_upload
    file_key = upload_to_ton(file_path, auth_profile, api_domain, bucket_name, trace)
    get_from_ton(file_key, auth_profile, api_domain, bucket_name, trace)

    # Verifying if downloaded and uploaded files are same
    actual_md5 = Digest::MD5.hexdigest(File.read("#{DOWNLOAD_FILE}"))
    expected_md5 = Digest::MD5.hexdigest(File.read("#{file_path}"))
    if(actual_md5 == expected_md5)
      puts("SUCCESS!!! Uploaded and downloaded files match")
    else
      puts("SOME ERROR WAS ENCOUNTERED")
    end
  elsif mode == :download
    get_from_ton(file_path, auth_profile, api_domain, bucket_name, trace)
  end
end

def get_oauth_profile(app_auth)
  auth_profile = {}
  if app_auth.nil?
    # User-auth (OAuth 1.0A)
    auth_profile[:consumer_key] = CONSUMER_KEY
    auth_profile[:consumer_secret] = CONSUMER_SECRET
    auth_profile[:token] = USER_TOKEN
    auth_profile[:secret] = USER_SECRET
    auth_profile[:authmode] = AuthenticationMode::USER
  else
    # App-auth (OAuth 2)
    auth_profile[:bearer_token] = get_bearer_token
    auth_profile[:authmode] = AuthenticationMode::APP
  end
  return auth_profile
end

def main(args)
  abort MISSING_SECRET_ERROR unless CONSUMER_SECRET

  parsed_arguments = parse_args(args)
  if parsed_arguments[:mode] != :download && parsed_arguments[:file_path].nil?
    abort "file path is required for --mode=#{parsed_arguments[:mode]}"
  end

  abort MISSING_TOKEN_ERROR unless parsed_arguments[:app_auth] || (USER_TOKEN && USER_SECRET)

  auth_profile = get_oauth_profile(parsed_arguments[:app_auth])

  execute(parsed_arguments[:mode], parsed_arguments[:trace], auth_profile, parsed_arguments[:file_path], parsed_arguments[:bucket_name])
end

main(ARGV)