#!/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

  def initialize(url, use_rate_limit, seed_filename=nil)
    @url = URI(url)
    @last_msgs = []
    @outmsg = ''
    @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 start_download(use_rate_limit, seed_filename)
    filename = File.basename(@url.to_s)
    output_filename = (seed_filename || "#{ filename }.tmp")

    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',
                  '--check-integrity=true',
                  "--dir=#{ $target_dir }",
                  '--enable-color=false',
                  '--file-allocation=none',
                  '--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

    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|
        update_progress(line)
        if line.match(/\n$/) then
          @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_progress(line)
    fields = line.split
    return unless fields[1] && fields[5] && fields[5].match(/^ETA/)

    match = fields[1].match(/\((\d+)%\)$/)
    if match then
      @progress = match[1]
    end
  end

  def stop
    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_NOTICE,
              '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, @url_to_download = * @client_message.split("\n")
        use_rate_limit = (rate_limit != '')

        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

    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
          errmsg = 'aria2c has exited with failure or terminated' \
                     + ' by signal with messages: %s'
          Syslog.log(Syslog::LOG_ERR, errmsg, @aria.last_msgs.join)
        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 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(torrent_url_path)
  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

  begin
    # We wait for DNS when starting to seed, because seeding should start
    # right after boot and we may not have a working DNS/VPN at that time.
    # In other situations it is better to 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

  # we have torrent url, we have rdiff 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($torrent_url_path)
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
