daemon

#!/usr/bin/env ruby
PROTOCOL_VERSION = "1.0"
%w(fileutils tempfile pathname optparse uri net/http net/https
net/http/post/multipart openssl zlib json logger digest/sha2).each{|g| require g}

Signal.trap("INT") { exit }

BASEDIR = File.absolute_path(File.dirname(__FILE__))
API_KEY_FILE = "#{BASEDIR}/api.key"
LOGGER = Logger.new(STDOUT)
USER_AGENT = "Electric Sheep HD Client/#{PROTOCOL_VERSION}"
LOGGER.formatter = proc do |severity, datetime, progname, msg|
  "[#{datetime.strftime("%H:%M:%S")}] #{severity}: #{msg}\n"
end

@argv = ARGV.dup
@options={}
OptionParser.new do |opts|
  opts.banner = "Usage: #{File.basename(__FILE__)} [OPTIONS]"
  opts.on("--server SERVER", "Control server to connect to") {|v| @options["server"] = v.to_s}
  opts.on("--apikey APIKEY", "Your api key - Default: read from ./api.key file") {|v| @options["apikey"] = v}
  opts.on("--nd", "--no-download", "Do not download movies") {|v| @options["no-download"] = true}
  opts.on("--np", "--no-progress", "Hide progress and ETA - This will speed up rendering a bit") {|v| @options["no-progress"] = true}
  opts.on("-k", "--keep-frames", "Do not delete rendered frames after upload") {|v| @options["keep_frames"] = true}
  opts.on("--nice NICENESS", "Niceness (Higher values result in lower process priority (default: 19, max: 19))") {|v| @options["nice"] = v}
  opts.on("--gpu", "Use GPU renderer (Fractorium - http://fracotrium.com)") {|v| @options["gpu"] = true}
  opts.on("--gpu-devices [Device-Ids]", "Use device(s) with given ids e.g.: '2,3'. Use --gpu-list-devices to get a list of your available devices.") {|v| @options["gpu-devices"] = v}
  opts.on("--gpu-list-devices", "Returns a list of your installed OpenCL devices") {|v| @options["gpu-list-devices"] = true}
  if Gem.win_platform?
    opts.on("--gpu-priority", "Set GPU render priority (-2: lowest, 2: highest)") {|v| @options["priority"] = v}
  else
    opts.on("--gpu-priority", "Set GPU render priority (1: lowest, 99: highest)") {|v| @options["priority"] = v}
  end
  opts.on("--insecure", "Ignore ssl certificate errors") {|v| @options["insecure"] = true}
  opts.on("--debug", "Debug-mode") {|v| @options["debug"] = true}
end.parse!
@options["apikey"] ? APIKEY_SWITCH_USED = true : APIKEY_SWITCH_USED = false

`which flam3-animate` ; raise "Missing flame3-animate. You have to install flam3 first" if $? != 0

unless Gem.win_platform?
  `which git` ; LOGGER.error "Missing git. Autoupdate wont work. Please install git" if $? != 0
  `which emberrender`; raise "emberanimate not found" if $? != 0 && @options["gpu"]
else
  raise "emberanimate not found in #{ENV["APPDATA"]}\\Fractorium" if @options["gpu"] &&
                                                                    !File.exist?("#{ENV["APPDATA"]}\\Fractorium\\emberrender.exe")
end

if @options["gpu-list-devices"]
  system "emberanimate --openclinfo"
  exit 0
end

@options["server"] ||= "https://sheeps.triple6.org"

# Autoupdater
if !ENV["electricsheep-hd-client-updated"]
  ENV["electricsheep-hd-client-updated"] = "true"
  Dir.chdir(BASEDIR)
  `git pull origin master` unless Gem.win_platform?
  if !ENV["OCRA"]
    exec("ruby #{__FILE__} #{@argv.join(' ')}")
  else
    exec("ocra #{__FILE__} #{@argv.join(' ')}")
  end
  exit
end

# Process lock
unless Gem.win_platform?
  LOCKFILE = "daemon.lock"
  if File.exist?(LOCKFILE)
    pid = File.read(LOCKFILE).strip.to_i
    if pid == 0
      File.delete(LOCKFILE)
      LOGGER.error "Broken pid file - start again"
      exit 1
    end

    # Check if pid is still running
    begin
      Process.kill(0, pid)
      # Locked
      raise "daemon is already running as pid: #{pid}"
    rescue Errno::ESRCH
      # Not locked
      File.open(LOCKFILE, "w") { |f| f.write Process.pid }
    end
  else
    File.open(LOCKFILE, "w") { |f| f.write Process.pid }
  end
end


class HTTPHalt < StandardError; end
class API
  attr_reader :http
  attr_reader :fingerprint
  attr_reader :key

  def initialize(options={})
    @server = options[:server]
    @key = options[:key]
    @uri = URI.parse(@server)
    @http = Net::HTTP.new(@uri.host, @uri.port)
    @http.read_timeout = 600
    if options[:server].start_with?("https://")
      @http.use_ssl = true
      @http.ca_file = "#{File.dirname(__FILE__)}/sheeps.crt"
      @http.verify_mode = OpenSSL::SSL::VERIFY_PEER
      @http.verify_mode = OpenSSL::SSL::VERIFY_NONE if options["insecure"]
      @http.verify_depth = 5
    end
  end

  def get(url)
    uri = @uri + url
    request = Net::HTTP::Get.new(uri.request_uri)
    request["user-agent"] = USER_AGENT
    response = @http.request(request)
    LOGGER.warn "Server response #{response.code}: #{response.body}" if response.code != '200'
    response
  end

  def pow
    start = Time.now
    LOGGER.info "Solving PoW..."
    pow = JSON.parse(get("api/fingerprint?apikey=#{@key}").body)
    sha = ""
    solution = 0
    until sha.start_with?(pow["pattern"])
      solution += 1
      sha = Digest::SHA2.hexdigest(pow["fingerprint"] + '^' + solution.to_s)
    end
    LOGGER.info "Solved after #{(Time.now - start).round} seconds"
    response = get("api/pow?apikey=#{@key}&fingerprint=#{pow["fingerprint"]}&pattern=#{pow["pattern"]}&solution=#{solution}")
    return nil if response.code != '200'
    @fingerprint = pow["fingerprint"]
    @fingerprint
  end

  def renew_key
    return if APIKEY_SWITCH_USED
    loop do
      print "Enter your apikey: "
      @key = gets.strip
      next if @key.nil? || @key.empty?
      if !verify_key
        File.delete(API_KEY_FILE) if File.exist?(API_KEY_FILE)
        next
      end
      File.open(API_KEY_FILE, "w" ) { |f| f.write @key }
      break
    end
    return self
  end

  def verify_key
    api = API.new({server: @server, insecure: @insecure, key: @key })
    response = api.get("api/verify?apikey=#{@key}")
    response.code == '200'
  end
end

# Default options:
@options["apikey"] ||= File.read(API_KEY_FILE).strip if File.exist?(API_KEY_FILE)
@options["debug"] ||= false
@options["no-download"] ||= false
@options["nice"] ||= "19"
$DEBUG = @options["debug"]
@options["gpu"] ||= false
@options["insecure"] ||= false
@options["keep_frames"] ||= false
@options["gpu-devices"] ||= "0"
Gem.win_platform? ? @options["priority"] ||= -2 : @options["priority"] ||= 1

# Initialize API
@api = API.new({server: @options["server"], key: @options["apikey"], insecure: @options["insecure"] })
if APIKEY_SWITCH_USED
  exit unless @api.verify_key
else
  @api = @api.renew_key unless @api.verify_key
end

# Main functions
def request_work(keep_frames: false)
  response = @api.get("api/request_work?apikey=#{@api.key}&fingerprint=#{@api.pow}&gpu=#{@options["gpu"]}")
  work = JSON.parse(response.body)
  if @options["gpu"]
    genome_file = "#{@animated_genome_dir}/#{work['sequence']}.flame"
    unless File.exist?(genome_file)
      LOGGER.info "Downloading #{work['sequence']} ..."
      download_animated_genome(work['sequence'])
    end
    LOGGER.info "Rendering... #{work['sequence']} Frame #{work['begin']}-#{work['end']}"
    frames = ember_render_frame(genome_file, work['begin'], work['end'])
    frames.sort.each_with_index do |frame, i|
      frame_time = i + work['begin'].to_i
      w = { "sequence": work['sequence'],
            "frame": frame_time }
      LOGGER.info "Uploading... #{work['sequence']} Frame #{frame_time}"
      upload_frame(frame, w)

      File.delete(frame) unless keep_frames
    end
  else
    work.each do |w|
      genome_file = "#{@animated_genome_dir}/#{w['sequence']}.flame"
      unless File.exist?(genome_file)
        LOGGER.info "Downloading #{w['sequence']} ..."
        download_animated_genome(w['sequence'])
      end
      LOGGER.info "Rendering... #{w['sequence']} Frame #{w['frame']}"
      frame = flam3_render_frame(genome_file, w["frame"])
      LOGGER.info "Uploading... #{w['sequence']} Frame #{w['frame']}"
      upload_frame(frame, w)
      File.delete(frame) unless keep_frames
    end

  end
end

def download_animated_genome(gid)
  output_file = "#{@animated_genome_dir}/#{gid}.flame"
  response = @api.get("#{@branch}/animated_genomes/#{gid}.flame.gz")
  raise HTTPHalt, "Error in download: #{response.body}" if response.code != "200"
  File.open("#{@animated_genome_dir}/#{gid}.flame.gz", "wb") { |f| f.write(response.body) }
  Zlib::GzipReader.open("#{@animated_genome_dir}/#{gid}.flame.gz") do |gz|
    File.open(output_file, "wb") do |g|
      IO.copy_stream(gz, g)
    end
  end
  File.delete("#{@animated_genome_dir}/#{gid}.flame.gz")
  output_file
end

def ember_render_frame(genome_file, begin_frame, end_frame)
  genome = {}
  genome["file"] = Pathname.new(genome_file.to_s)
  genome["id"] = genome["file"].basename.sub(/\.flame$/,"")

  concat_name = "#{genome['id']}"
  Dir.mkdir("#{@frame_dir}/#{concat_name}") unless File.exist?("#{@frame_dir}/#{concat_name}")
  if Gem.win_platform?
    emberanimate = "#{ENV["APPDATA"]}/Fractorium/emberanimate"
  else
    emberanimate = "emberanimate"
  end
  system(emberanimate,
          "--in=#{File.expand_path(genome_file)}",
          "--format=jpg",
          "--opencl", # TODO Control with flag
          "--priority=#{@options["priority"]}",
          "--sp",
          "#{@options['no-progress'] ? '' : '--progress'}",
          "--prefix=../frames/#{concat_name}/",
          "--device=#{@options["gpu-devices"]}",
          "--begin=#{(begin_frame-1)}",
          "--end=#{(end_frame)}",
          "--isaac_seed=fractorium",
          "--enable_comments")
  result = $? # TODO Does it always return 0 ?
  fail "Emberrender exited: #{result.exitstatus}" unless result.success?
  # TODO is it always 04d ? George added 03d here
  frames = ((begin_frame.to_i)..(end_frame.to_i)).map{ |i| "#{@frame_dir}/#{concat_name}/#{"%04d" % i.to_s}.jpg" }
  frames
end

def flam3_render_frame(genome_file, frame)
  genome = {}
  genome["file"] = Pathname.new(genome_file.to_s)
  genome["id"] = genome["file"].basename.sub(/\.flame$/,"")

  concat_name = "#{genome['id']}"
  Dir.mkdir("#{@frame_dir}/#{concat_name}") unless File.exist?("#{@frame_dir}/#{concat_name}")
  ENV["in"] = File.expand_path(genome_file)
  ENV["prefix"]= "#{@frame_dir}/#{concat_name}/"
  ENV["format"] = "jpg"
  ENV["jpeg"] = "95"
  ENV["begin"] = frame.to_s
  ENV["end"] = frame.to_s
  @options["no-progress"] ? ENV["verbose"] = "0" : ENV["verbose"] = "1"
  if Gem.win_platform?
    `start /low /b flam3-animate`
  else
    `nice -#{@options["nice"]} flam3-animate`
  end
  "#{@frame_dir}/#{concat_name}/#{"%05d" % frame.to_s}.jpg"
end

def upload_frame(frame_file, work_set)
  url = URI.parse(@options["server"] + "/api/upload")
  File.open(frame_file) do |jpg|
    req = Net::HTTP::Post::Multipart.new(url.path,
                                         "file" => UploadIO.new(jpg, "image/jpeg", "image.jpg"),
                                         "apikey" => @api.key,
                                         "work_set" => work_set.to_json,
                                         "branch" => @branch,
                                         "gpu" => @options["gpu"])
    req["user-agent"] = USER_AGENT
    @api.http.verify_mode = OpenSSL::SSL::VERIFY_NONE if @options["insecure"]
    response = @api.http.request(req)
    raise HTTPHalt, "Error in upload: #{response.body}" if response.code != "200"
  end
end

def download_sequence
  response = @api.get("api/request_sequence_list?apikey=#{@api.key}&gpu=#{@options["gpu"]}")
  if response.code != "200"
    return response.body
  end
  sequences = JSON.parse(response.body)
  wanted = sequences.select{|s| !File.exist?("#{@movie_dir}/#{s}")}.first
  return "Nothing to download" if wanted.nil?
  LOGGER.info "Downloading #{wanted}..."
  response = @api.get("api/request_sequence?apikey=#{@api.key}&fingerprint=#{@api.pow}&seq=#{wanted}&gpu=#{@options["gpu"]}")
  File.open("#{@movie_dir}/#{wanted}", "wb") {|f| f.write(response.body)}
  return "OK"
end

# Loop
loop do
  sleep 1
  begin
    response = @api.get("api/active_season?apikey=#{@api.key}&gpu=#{@options["gpu"]}")
    raise HTTPHalt, "Error in season: #{response.body}" if response.code != "200"
    @season = JSON.parse(response.body)
  rescue JSON::ParserError, Errno::ECONNREFUSED, Errno::ECONNRESET,
         SocketError, Errno::ETIMEDOUT, Errno::ENETUNREACH,
         HTTPHalt, Net::ReadTimeout,
         Net::OpenTimeout => e
    LOGGER.error "#{e}"
    sleep 60
    next
  end

  @branch = @season["name"]
  LOGGER.info "Active season: #{@branch}"
  File.open("#{BASEDIR}/.active-season", "w") { |f| f.write @season["name"]}
  @branch_dir = "#{BASEDIR}/branches/#{@branch}"
  @genome_dir = "#{@branch_dir}/genomes"
  @animated_genome_dir = "#{@branch_dir}/animated_genomes"
  @frame_dir = "#{@branch_dir}/frames"
  @movie_dir = "#{@branch_dir}/movies"
  FileUtils.mkdir_p(@genome_dir)
  FileUtils.mkdir_p(@animated_genome_dir)
  FileUtils.mkdir_p(@frame_dir)
  FileUtils.mkdir_p(@movie_dir)

  begin
    unless @options["no-download"] || rand(1..10) != 1
      LOGGER.info "Requesting new sequence ..."
      LOGGER.info download_sequence
    end
    LOGGER.info "Requesting new work..."
    request_work(keep_frames: @options["keep_frames"])
    break if ENV["OCRA"]
  rescue JSON::ParserError, Errno::ECONNREFUSED, Errno::ECONNRESET,
         Errno::ENETUNREACH, SocketError, Errno::ETIMEDOUT, HTTPHalt,
         Net::ReadTimeout, Net::OpenTimeout => e
    LOGGER.error "#{e}"
    sleep 60
    next
  end
end