ubicloud/bin/ubi

#!/usr/bin/env ruby
# frozen_string_literal: true

begin
  require_relative "../.env"
rescue LoadError
end

require "net/http"
require "json"

unless (token = ENV["UBI_TOKEN"])
  warn "! Personal access token must be provided in UBI_TOKEN env variable for use"
  exit 1
end

url = ENV["UBI_URL"] || "https://api.ubicloud.com/cli"
ssh_progs = %w[ssh scp sftp].freeze.each(&:freeze)
pg_progs = %w[psql pg_dump pg_dumpall].freeze.each(&:freeze)
allowed_progs = (ssh_progs + pg_progs).freeze

get_prog = lambda do |prog|
  return unless allowed_progs.include?(prog)
  ENV["UBI_#{prog.upcase}"] || prog
end

uri = URI(url)
argv = ARGV
headers = {
  "authorization" => "Bearer: #{token}",
  "x-ubi-version" => File.read(File.expand_path("../../cli/version.txt", __FILE__)).chomp,
  "content-type" => "application/json",
  "accept" => "text/plain",
  "connection" => "close"
}
confirmation_prompt = false
confirmation = nil

1.times do
  if ENV["UBI_DEBUG"] == "1"
    p [:sending, *argv]
  end
  response = Net::HTTP.post(uri, {"argv" => argv}.to_json, headers)

  case response.code.to_i
  when 200...300
    if (prog_type = response["ubi-command-execute"])
      unless ARGV.include?(prog_type)
        warn "! Invalid server response, not executing program not in original argv"
        exit 1
      end

      unless (prog = get_prog[prog_type])
        warn "! Invalid server response, unsupported program requested"
        exit 1
      end

      if pg_progs.include?(prog_type) && (pgpassword = response["ubi-pgpassword"])
        ENV["PGPASSWORD"] = pgpassword
      end

      argv_set = ARGV.to_set

      args = response.body.split("\0")
      invalid_message = nil
      sep_seen = false
      custom_arg_seen = false
      pg_dumpall = false

      args.each do |arg|
        if arg == "--"
          sep_seen = true
        elsif !argv_set.include?(arg)
          if custom_arg_seen
            invalid_message = "! Invalid server response, multiple arguments not in submitted argv"
            break
          elsif sep_seen
            custom_arg_seen = true
          elsif prog_type == "pg_dumpall" && arg.start_with?("-d")
            pg_dumpall = true
            custom_arg_seen = true
          else
            invalid_message = "! Invalid server response, argument before '--' not in submitted argv"
            break
          end
        end
      end

      unless sep_seen || pg_dumpall
        invalid_message ||= "! Invalid server response, no '--' in returned argv"
      end

      if invalid_message
        if ENV["UBI_DEBUG"] == "1"
          p [:failure, prog, *args]
        end
        warn invalid_message
        exit 1
      else
        if ENV["UBI_DEBUG"] == "1"
          p [:exec, prog, *args]
        end
        Process.exec(prog, *args)
      end
    else
      if (confirmation_prompt = response["ubi-confirm"])
        if confirmation
          warn "! Invalid server response, repeated confirmation attempt"
          exit 1
        else
          $stdout.print response.body
          $stdout.print "\n#{confirmation_prompt}: "
          confirmation = $stdin.readline.chomp
          argv.unshift(confirmation)
          argv.unshift("--confirm")
          redo
        end
      else
        $stdout.print response.body
      end
      exit 0
    end
  else
    warn response.body
    exit 1
  end
end