#!/usr/bin/env ruby
# Sets up vagrant for your developer machine. This will:
# 1. Modify .ssh/config file so you can log in to Vagrant using `ssh barkeep_vagrant`
#    instead of `vagrant ssh` (which is required to deploy to Vagrant).
# 2. Add the vagrant user's public ssh key to root's .ssh/authorized_keys file, so you can login as root.

require "bundler/setup"
require "timeout"
require "open3"

def hostname() "barkeep_vagrant" end

def setup_vagrant
  virtual_box_installed = `which VBoxManage`.size > 0
  unless virtual_box_installed
    puts "Vagrant requires VirtualBox to be installed. Download it here:"
    puts "https://www.virtualbox.org/wiki/Downloads"
    exit 1
  end
  stream_output("vagrant up")
  setup_ssh_config
  # Ensure no old packages are lingering around. This will avoid possible 404's when installing packages.
  puts "* Updating Vagrant's apt packages. This will take a minute."
  # apt-get update can hang while downloading packages from unresponsive mirrors, and never times itself
  # out. Often the problem can be solved by retrying.
  unless try_n_times(2, "ssh root@#{hostname} apt-get update -y -qq", 90)
    puts "Unable to complete `apt-get update` inside your Vagrant VM. " +
        "This can happen when Ubuntu's apt mirrors are being unresponsive."
    exit 1
  end
end

def setup_ssh_config
  ssh_config_path = File.expand_path("~/.ssh/config")
  FileUtils.touch(ssh_config_path) unless File.exists?(ssh_config_path)

  unless File.read(ssh_config_path).include?(hostname)
    # Create an entry in our ~/.ssh/config which allows us to ssh into this vagrant box by hostname.
    original_ssh_config = File.read(ssh_config_path)
    vagrant_ssh_config = `vagrant ssh-config --host "#{hostname}"`
    # The ssh config block generated by Vagrant looks like this if you're curious:
    # Host barkeep_vagrant
    #   HostName 127.0.0.1
    #   User vagrant
    #   Port 2222
    #   StrictHostKeyChecking no
    #   PasswordAuthentication no
    #   IdentityFile /Users/philc/.vagrant.d/insecure_private_key
    #   IdentitiesOnly yes

    # Don't check the "known hosts file" when sshing into Vagrant.
    vagrant_ssh_config = vagrant_ssh_config.split("\n").
        reject { |line| line.match(/UserKnownHostsFile/) }.join("\n") + "\n"
    File.open(ssh_config_path, "w") { |file| file.write(vagrant_ssh_config + original_ssh_config) }
  end

  # The vagrant user has the default "vagrant public key" in authorized_keys. Make it so for root as well.
  remote_commands = "sudo mkdir /root/.ssh 2> /dev/null;" +
      "sudo cp .ssh/authorized_keys /root/.ssh/authorized_keys"
  stream_output "ssh vagrant@#{hostname} '#{remote_commands}'"
end

# Runs the command and raises an exception if its status code is nonzero.
def stream_output(command)
  puts command
  exit_status = nil
  Open3.popen3(command) do |stdin, stdout, stderr, wait_thread|
    stdout.each { |line| puts line }
    stderr.each { |line| puts line }
    exit_status = wait_thread.value.to_i
  end
  raise %Q(The command "#{command}" failed.) unless exit_status == 0
end

def try_n_times(n, command, timeout)
  attempt = 0
  while (attempt < n)
    error = nil
    attempt += 1
    begin
      Timeout::timeout(timeout) { stream_output(command) }
    rescue Timeout::Error => error
      puts "The command '#{command}' didn't finish within #{timeout} seconds."
    rescue StandardError => error
      puts error.message
    end
    return true unless error
    puts "Trying again." if attempt < n
  end
  false
end

setup_vagrant
