Ruby file segment downloader

download.rb

#!/usr/bin/env ruby

require 'httparty'
require 'parallel'
require 'optparse'
require 'curses'

PART_CHAR = "\u2588"

def parse_options
  options = {parts: nil, threads: 20}
  OptionParser.new do |opts|
    opts.banner = "Usage: dl.rb [options] <URL to download>"
    opts.on("-p", "--parts PARTS", OptionParser::DecimalInteger, "How many parts to split the file into", "Default is enough parts to make each part 10MB") do |p|
      options[:parts] = p.to_i
    end
    opts.on("-t", "--threads THREADS", OptionParser::DecimalInteger, "How many threads to run in", "Default is 20") do |t|
      options[:threads] = t.to_i
    end
    opts.on_tail("-h", "--help", "Show this message") do
      puts opts
      exit
    end
  end.parse!

  url = ARGV.pop
  unless url
    puts "Please provide a URL!".red
    exit 1
  end

  [options, url]
end

def get_content_length(httpartyOptions, url)
  head = HTTParty.head(url, httpartyOptions)
  unless head.ok?
    puts 'Could not get header for file, aborting!'
    exit 1
  end
  head['content-length'].to_i
end

def chunk_content(content_length, num_parts)
  seg_size = (content_length / num_parts.to_f).ceil
  segs = (1..num_parts).map{|i| [seg_size * (i - 1), seg_size * i - 1]}
  segs[-1][1] = content_length - 1
  segs
end

def file_ok?(filename, from_byte, to_byte)
  File.exists?(filename) && File.file?(filename) && (to_byte - from_byte + 1 == File.size(filename))
end

def download_part(filename, seg_number, url, httpartyOptions, from_byte, to_byte)
  return true if file_ok?(filename, from_byte, to_byte)

  opts = httpartyOptions.merge(headers: {"Range" => "bytes=#{from_byte}-#{to_byte}"})
  retries = 0
  begin
    File.open(filename, "wb") do |f|
      f.binmode
      HTTParty.get(url, opts) do |chunk|
        f.write chunk
      end
      f.close
    end

    unless file_ok?(filename, from_byte, to_byte)
      raise 'Invalid file size'
    end
  rescue
    if retries < 5
      retries += 1
      retry
    else
      return false
    end
  end

  true
end

def initUI(filename)
  Curses.init_screen
  Curses.start_color
  Curses.init_pair(1, Curses::COLOR_RED, Curses::COLOR_BLACK)
  Curses.init_pair(2, Curses::COLOR_BLUE, Curses::COLOR_BLACK)
  Curses.init_pair(3, Curses::COLOR_GREEN, Curses::COLOR_BLACK)
  Curses.setpos(0, 0)
  Curses.addstr("Downloading: #{filename}")
end

def closeUI
  Curses.close_screen
end

@updating_ui = false
def updateUI(parts)
  return if @updating_ui
  @updating_ui = true
  Curses.setpos(1, 0)
  parts.each_with_index do |part, index|
    Curses.attron(Curses.color_pair(part+1)|Curses::A_NORMAL){
      Curses.addstr(PART_CHAR)
    }
  end

  rows = parts.length / `tput cols`.to_i
  Curses.setpos(rows + 2, 0)
  Curses.addstr("Downloaded: #{parts.find_all{|i| i == 2}.length.to_s.rjust(parts.length.to_s.length, ' ')} / #{parts.length}")
  Curses.refresh
  @updating_ui = false
end

def part_filename(filename, part, num_parts)
  "#{filename}.#{part.to_s.rjust((num_parts - 1).to_s.length, '0')}"
end

def join_parts(filename, num_parts)
  File.open(filename, "wb") do |mf|
    mf.binmode
    num_parts.times do |i|
      file = part_filename filename, i, num_parts
      mf.write(File.open(file, 'rb').read)
      File.delete(file)
    end
  end
end


def run
  options, url = parse_options

  uri = URI.parse(url)

  httpartyOptions = {}
  httpartyOptions.merge!(verify: false) if uri.scheme == 'https'
  if uri.user
    # Use digest auth as default because I am lazy
    httpartyOptions.merge!(digest_auth: {username: uri.user, password: uri.password})
    uri.user = uri.password = nil
  end
  url = uri.to_s

  content_length = get_content_length(httpartyOptions, url)
  if options[:parts].nil?
    options[:parts] = content_length / (10 * 1024 * 1024)
  end
  segs = chunk_content(content_length, options[:parts])

  filename = URI.unescape(url.split('/').last)
  digLength = (options[:parts] - 1).to_s.length

  httpartyOptions.merge!(stream_body: true)
  parallelOptions = {in_threads: options[:threads]}

  parts = [0] * options[:parts]

  begin
    initUI(filename)
    updateUI(parts)
    Parallel.each(0...options[:parts], parallelOptions) do |i|
      parts[i] = 1
      updateUI(parts)
      res = download_part part_filename(filename, i, options[:parts]), i, url, httpartyOptions, segs[i][0], segs[i][1]
      if res
        parts[i] = 2
        updateUI(parts)
      else
        closeUI
        exit 1
      end
    end
  ensure
    closeUI
  end

  join_parts(filename, options[:parts])
end

run