#!/usr/bin/ruby

require 'fileutils'
require 'open3'
require 'puavo/conf'
require 'resolv'
require 'socket'
require 'syslog'
require 'uri'

class Aria
  attr_reader :last_msgs, :progress, :status, :stdout, :url

  DOWNLOAD_STATE_FIELDMAP = {
    'CN'    => 'connections',
    'DL'    => 'download_speed',
    'ETA'   => 'download_eta',
    'SD'    => 'seeders',
    'SEED'  => 'seed_ratio',
    'UL'    => 'upload_speed',
  }

  def initialize(url, use_rate_limit, seed_filename=nil)
    @url = URI(url)
    @control_file_path = nil
    @connections_state = {}
    @in_download_stage = false
    @last_msgs = []
    @network_torrent_friendliness_logged = false
    @outmsg = ''
    @previous_state_log_time = Process.clock_gettime(Process::CLOCK_MONOTONIC)
    @progress = nil
    @status = nil

    start_download(use_rate_limit, seed_filename)
  end

  def ratelimit(puavo_conf_key)
    puavoconf = Puavo::Conf.new
    rate_limit = puavoconf.get(puavo_conf_key)
    puavoconf.close

    rate_limit = nil if rate_limit == ''

    if rate_limit then
      if !rate_limit.match(/^[0-9]+[km]$/) then
        raise "the rate limit was not understood, got '#{ rate_limit }'"
      end
      rate_limit.upcase!        # aria2c(1) says uppercase letters must be used
    end

    rate_limit
  end

  def remove_corrupted_control_file
    return unless @control_file_path
    # aria2c refuses to do anything with a corrupt control file (which may be
    # empty, truncated or something else), so remove the control file if
    # there is any.
    Syslog.log(Syslog::LOG_NOTICE,
               'removing a corrupted control file: %s', @control_file_path)
    FileUtils.rm_f(@control_file_path)
  end

  def start_download(use_rate_limit, seed_filename)
    filename = File.basename(@url.to_s)
    output_filename = (seed_filename || "#{ filename }.tmp")
    @control_file_path = File.join($target_dir, "#{ output_filename }.aria2")

    torrent_url = @url.clone
    torrent_url.path = "/torrents/#{ filename }.torrent"

    download_complete_path \
      = '/usr/lib/puavo-ltsp-install/puavo-torrent-download-complete'

    aria_args = [ '/usr/bin/env',
                  'LANG=C',
                  "PUAVO_TORRENT_UPDATE_PID=#{ Process.pid }",
                  'aria2c',
                  '--bt-enable-lpd=true',
                  '--bt-stop-timeout=300',
                  '--ca-certificate=/etc/puavo-conf/rootca.pem',
                  '--certificate=/etc/puavo/certs/hostorgcabundle.pem',
                  '--continue',
                  "--dir=#{ $target_dir }",
                  '--enable-color=false',
                  '--file-allocation=falloc',
                  '--follow-torrent=mem',
                  '--human-readable=false',
                  "--index-out=1=#{ output_filename }",
                  '--no-conf=true',
                  "--on-bt-download-complete=#{ download_complete_path }",
                  '--private-key=/etc/puavo/certs/host.key',
                  '--seed-ratio=0.0',
                  '--summary-interval=1' ]

    if use_rate_limit then
      download_rate_limit = ratelimit('puavo.image.torrent.download.ratelimit')
      aria_args << "--max-overall-download-limit=#{ download_rate_limit }" \
        if download_rate_limit
      upload_rate_limit = ratelimit('puavo.image.torrent.upload.ratelimit')
      aria_args << "--max-overall-upload-limit=#{ upload_rate_limit }" \
        if upload_rate_limit
    end

    # If case we are seeding, do not verify the contents as they should be
    # okay if we are entering the seeding mode, otherwise check file integrity.
    aria_args << (seed_filename                   \
                    ? '--bt-seed-unverified=true' \
                    : '--check-integrity=true')

    aria_args << torrent_url.to_s

    Syslog.log(Syslog::LOG_NOTICE,
               'starting aria2c with args %s', aria_args.join(' '))

    stdin, @stdout, stderr, @wait_thr \
      = Open3.popen3(*aria_args, :err => '/dev/null', :in  => '/dev/null')
  end

  def read_stdout
    begin
      new_outmsg = ''
      @outmsg += @stdout.read_nonblock(1024)
      @outmsg.each_line do |line|
        if line.match(/\n$/) then
          update_connections_state(line.chomp)
          @last_msgs = (@last_msgs + [ line ]).last(20)
        else
          new_outmsg = line
        end
      end
      @outmsg = new_outmsg
    rescue EOFError, IOError
      @stdout.close
      @stdout = nil
    end
  end

  def update_connections_state(line)
    # progress may be reported on "FILE: [MEMORY]" downloads (torrent files),
    # so do not update connection state until we are downloading the actual
    # payload
    if line.match(/\AFILE: \//) then
      @in_download_stage = true
    end

    return unless @in_download_stage

    m = line.match(/\A\[(.*?)\]/)
    return unless m

    gid, aria2_progress, *the_rest = * m[1].split

    # assignments to match, not comparisons
    if match = aria2_progress.match(/\((\d+)%\)$/) then
      @progress = match[1]
    elsif match = aria2_progress.match(/^SEED\((.*)\)$/) then
      @connections_state['DL']   = 0
      @connections_state['ETA']  = 0
      @connections_state['SEED'] = match[1]
    end

    the_rest.each do |aria2_info|
      field, value = *aria2_info.split(':')
      if DOWNLOAD_STATE_FIELDMAP.has_key?(field) then
        @connections_state[ field ] = value
      end
    end

    # if we have not logged the connections status in the last five minutes,
    # log it now
    time_now = Process.clock_gettime(Process::CLOCK_MONOTONIC)
    if (time_now - @previous_state_log_time) >= 300 then
      @previous_state_log_time = time_now
      log_connections_state
    end
  end

  def log_connections_state
    logfields = []
    @connections_state.sort.each do |field, value|
      fieldname = DOWNLOAD_STATE_FIELDMAP[field].to_s
      logfields << "#{ fieldname }=#{ value }"
    end
    return if logfields.empty?

    Syslog.log(Syslog::LOG_NOTICE, 'connections state for %s: %s',
                                   @url, logfields.join(' '))

    log_network_torrent_friendliness
  end

  def get_network_identifier
    begin
      output, cmdstatus = Open3.capture2('ip', 'route', 'show', 'default')
      raise 'ip route returned failure' unless cmdstatus.success?
      interface = output.split[4]
      output, cmdstatus = Open3.capture2('iw', 'dev', interface, 'link')
      return interface unless cmdstatus.success?
      return "#{ interface }:#{ output[/^\s*SSID:\s*(.*)$/, 1] }"
    rescue StandardError => e
      Syslog.log(Syslog::LOG_WARNING,
                 'could not determine network identifier: %s', e.message)
      return nil
    end
  end

  def log_network_torrent_friendliness
    return if @network_torrent_friendliness_logged

    # if we see some seeders, or we have done some uploading,
    # then it looks like this network is torrent-friendly and we report it
    if @connections_state['SD'].to_i > 0 || @connections_state['UL'] then
      network_id = get_network_identifier()
      if network_id then
        Syslog.log(Syslog::LOG_NOTICE,
                   "network #{ network_id } is torrent-friendly")
      end
      @network_torrent_friendliness_logged = true
    end
  end

  def stop
    log_connections_state

    if @wait_thr.alive? then
      Syslog.log(Syslog::LOG_INFO, 'killing the current aria2c process')
      Process.kill('TERM', @wait_thr.pid)
      wait
    end
  end

  def wait
    @status = @wait_thr.value
  end
end

class Server
  def initialize(server_socket, aria)
    @server_socket = server_socket
    @aria = aria

    @client_message = nil
    @client_socket = nil
    @read_filename = false

    @current_progress = nil
    @write_progress   = false

    @selfpipe_r, @selfpipe_w = IO.pipe
  end

  def loop
    readable = [ @server_socket, @selfpipe_r ]
    writeable = []

    if @client_socket then
      readable  << @client_socket if @read_filename
      writeable << @client_socket if @write_progress
    end

    if @aria then
      readable << @aria.stdout if @aria.stdout
    end

    ios = IO.select(readable, writeable)

    if ios[0].include?(@selfpipe_r) then
      # we got "bt-download-complete" signal from aria2c, all is good and
      # we can close the client connection
      @selfpipe_r.read_nonblock(1)
      if @client_socket then
        begin
          @client_socket.write_nonblock("OK\n")
          write_torrent_url
        rescue Errno::EPIPE
          Syslog.log(Syslog::LOG_INFO,
                     'download finished, yet client has closed the connection')
        end
        @client_socket.close
        @client_socket = nil
      end
      return
    end

    if ios[0].include?(@server_socket) then
      Syslog.log(Syslog::LOG_INFO, 'new client connection')
      if @client_socket then
        Syslog.log(Syslog::LOG_INFO, 'breaking old client connection')
        @client_socket.close
        @client_socket = nil
      end

      @client_socket, client_addr = @server_socket.accept_nonblock
      @client_message = ''
      @read_filename = true

      # @client_socket has been switched, must return
      return
    end

    if ios[1].include?(@client_socket) then
      # we should write progress information to client
      begin
        @client_socket.write_nonblock("#{ @current_progress }\n")
      rescue Errno::EPIPE
        Syslog.log(Syslog::LOG_INFO,
                   'client has unexpectedly closed the connection')
        @client_socket.close
        @client_socket = nil
        if @aria then
          @aria.stop
          @aria = nil
        end
        return
      end
      @write_progress = false
    end

    if ios[0].include?(@client_socket) then
      # client is sending us a url to download
      begin
        @client_message += @client_socket.read_nonblock(1024)
      rescue EOFError
        @client_message.chomp!

        @client_socket.shutdown(:RD)
        @read_filename = false
        @current_progress = nil

        # stop previous download, if any
        @aria.stop if @aria

        # provide a mechanism for client to make aria stop seeding previous
        # downloads
        if @client_message == 'stop' then
          if @aria then
            Syslog.log(Syslog::LOG_NOTICE,
                       'stopping aria because of client request')
          else
            Syslog.log(Syslog::LOG_DEBUG,
              'aria is not run (at least not by us),' \
                + ' even though client requests it to stop')
          end

          @aria = nil
          begin
            @client_socket.write_nonblock("STOPPED\n")
          rescue Errno::EPIPE
          end
          @client_socket.close
          @client_socket = nil
          return
        end

        rate_limit_or_seed, requested_url, seedpatch_dir \
          = * @client_message.split("\n")
        if rate_limit_or_seed == 'seed' then
          Syslog.log(Syslog::LOG_NOTICE, 'client requested seeding')
          @aria = handle_aria_seeding_request(requested_url, seedpatch_dir)
        else
          rate_limit = rate_limit_or_seed
          use_rate_limit = (rate_limit != '')
          @url_to_download = requested_url

          Syslog.log(Syslog::LOG_NOTICE,
                     "client asks to download %s",
                     @url_to_download)

          begin
            # We put the url into this path only after the initial download
            # has been successful.  How that we are starting to download
            # something (to .tmp file), we should not seed anything old.
            File.unlink($torrent_url_path)
          rescue Errno::ENOENT
          end

          # start a new download
          @aria = Aria.new(@url_to_download, use_rate_limit)
        end
      end
    end

    if @aria then
      # aria is reporting us something, possibly new progress information
      @aria.read_stdout if @aria.stdout && ios[0].include?(@aria.stdout)

      new_progress = @aria.progress
      if new_progress != @current_progress then
        @current_progress = new_progress
        @write_progress = true
        Syslog.log(Syslog::LOG_INFO,
                   "download progress for %s is now %s%%",
                   @aria.url,
                   @current_progress)
      end

      if !@aria.stdout then
        @aria.wait
        if @aria.status.success? then
          errmsg = 'aria has exited with success, even though it should be' \
                     + " seeding '%s'"
          Syslog.log(Syslog::LOG_ERR, errmsg, @aria.url)
        else
          exitcode = @aria.status.exitstatus    # may be nil if signaled
          messages = @aria.last_msgs.join
          errmsg = 'aria2c has exited with failure or terminated' \
                     + ' by signal with messages (exitcode=%s): %s'
          Syslog.log(Syslog::LOG_ERR, errmsg, exitcode, messages)
          if messages.match(/Failed to read from the segment file/) \
            || messages.match(/info hash mismatch/) then
              # We do not restart but the next time it should be better
              # (because a corrupted control file makes aria2c non-functional).
              @aria.remove_corrupted_control_file
          end
        end
        @aria = nil

        if @client_socket then
          begin
            @client_socket.write_nonblock("FAIL\n")
          rescue Errno::EPIPE
          end
          @client_socket.close
          @client_socket = nil
        end
      end
    end
  end

  def download_complete
    Syslog.log(Syslog::LOG_INFO, 'bittorrent download complete')
    @selfpipe_w.write_nonblock('o')
  end

  def handle_aria_seeding_request(requested_url, seedpatch_dir)
    # We should use the same server for seeding where we have obtained
    # the patch for the seed file, as the seed torrent will more likely
    # not exist in the server where we have obtained the patch from.
    # (This is for the image updates, for cases where we have obtained
    # the image from an rdiff).
    seedpatch_uri = IO.read("#{ seedpatch_dir }/.aria2_torrent_url").chomp \
                      rescue nil
    if seedpatch_uri then
      uri = URI(requested_url)
      uri.host = URI(seedpatch_uri).host
      @url_to_download = uri.to_s
    else
      @url_to_download = requested_url
    end

    write_torrent_url
    @aria = start_aria_seeding(false)
  end

  def write_torrent_url
    # This is a sign for subsequent startups of this tool that it should
    # start seeding the url in this file.
    tmpfile = "#{ $torrent_url_path }.tmp"
    File.open(tmpfile, 'w') { |f| f.puts(@url_to_download) }
    FileUtils.mv(tmpfile, $torrent_url_path)
  end
end

def wait_for_dns(host)
  Syslog.log(Syslog::LOG_NOTICE,
             "checking that dns resolves for server #{ host }")

  result = nil
  Resolv::DNS.open do |dns|
    dns.timeouts = 5
    15.times do |i|
      begin
        result = dns.getaddress(host)
        break if result
      rescue StandardError => e
      end
      sleep(i)
    end
  end

  unless result then
    raise "Could not resolve #{ host }"
  end
end

def start_aria_seeding(wait_until_dns_is_available)
  begin
    torrent_url = IO.read($torrent_url_path).chomp
  rescue Errno::ENOENT
    # we have no torrent url, we are not going to seed
    return nil
  end

  seed_filename = File.basename(torrent_url)
  return nil if seed_filename == ''

  seed_path = "#{ $target_dir }/#{ seed_filename }"
  return nil unless File.exist?(seed_path)

  host = URI(torrent_url).host

  if wait_until_dns_is_available then
    begin
      # We wait for DNS when starting to seed right after boot, because
      # seeding should right away and we may not have a working DNS/VPN
      # at that time.  In other situations it is better to skip this check
      # or return an error right away.  Note that server is listening for
      # connections at this time, we will deal with those later.
      wait_for_dns(host)
    rescue StandardError => e
      Syslog.log(Syslog::LOG_ERR, "could not resolve #{ host }, not seeding")
      return nil
    end
  end

  # we have torrent url, we have file on the right path,
  # we have a working DNS, we can seed:

  Syslog.log(Syslog::LOG_NOTICE,
             "starting to seed #{ torrent_url } from #{ seed_path }")

  Aria.new(torrent_url, true, seed_filename)
end

Syslog.open( File.basename($0) )

$target_dir = ARGV[0]
if !$target_dir then
  progname = File.basename($0)
  Syslog.log(Syslog::LOG_ERR,
             "#{ progname } was not given a target directory as an argument")
  exit 1
end

$torrent_url_path = "#{ $target_dir }/.aria2_torrent_url"

server_socket = Socket.for_fd(3)

aria = nil
begin
  # this will return nil if there is nothing to seed
  aria = start_aria_seeding(true)
rescue StandardError => e
  Syslog.log(Syslog::LOG_ERR,
             'could not start aria2 for seeding: %s', e.message)
end

server = Server.new(server_socket, aria)

Signal.trap('SIGUSR1', proc { server.download_complete })

while true do
  begin
    server.loop
  rescue StandardError => e
    Syslog.log(Syslog::LOG_ERR, 'we got some unexpected error: %s', e.message)
    exit(1)
  end
end

Syslog.close
