From a39877ac27a3d2b428a40f4c02c8e9be6f083d86 Mon Sep 17 00:00:00 2001
From: Shawn Dahlen <shawn@dahlen.me>
Date: Tue, 18 Jun 2013 17:26:29 -0400
Subject: [PATCH] Add support for managing host /etc/hosts file.

This commit also include preservation original /etc/hosts file on
guest machines.
---
 Gemfile                                       |   5 +-
 .../action/delete_guests.rb                   |  33 ---
 ...{update_guests.rb => update_hosts_file.rb} |  16 +-
 lib/vagrant-hostmanager/command.rb            |  11 +-
 lib/vagrant-hostmanager/config.rb             |  11 +-
 lib/vagrant-hostmanager/hosts_file.rb         | 216 +++++-------------
 lib/vagrant-hostmanager/plugin.rb             |  15 +-
 lib/vagrant-hostmanager/provisioner.rb        |   6 +-
 lib/vagrant-hostmanager/version.rb            |   2 +-
 test/Vagrantfile                              |   7 +-
 vagrant-hostmanager.gemspec                   |   4 +
 11 files changed, 105 insertions(+), 221 deletions(-)
 delete mode 100644 lib/vagrant-hostmanager/action/delete_guests.rb
 rename lib/vagrant-hostmanager/action/{update_guests.rb => update_hosts_file.rb} (62%)

diff --git a/Gemfile b/Gemfile
index e402da5..052a8bc 100644
--- a/Gemfile
+++ b/Gemfile
@@ -1,4 +1,7 @@
 source 'https://rubygems.org'
 
-gem 'vagrant', :git => 'git://github.com/mitchellh/vagrant.git', :tag => 'v1.2.1'
+group :development do
+  gem 'vagrant', :git => 'git://github.com/mitchellh/vagrant.git', :tag => 'v1.2.2'
+end
+
 gemspec
diff --git a/lib/vagrant-hostmanager/action/delete_guests.rb b/lib/vagrant-hostmanager/action/delete_guests.rb
deleted file mode 100644
index 244b052..0000000
--- a/lib/vagrant-hostmanager/action/delete_guests.rb
+++ /dev/null
@@ -1,33 +0,0 @@
-require 'vagrant-hostmanager/hosts_file'
-
-module VagrantPlugins
-  module HostManager
-    module Action
-      class DeleteGuests
-        include HostsFile
-
-        def initialize(app, env)
-          @app = app
-          @machine = env[:machine]
-          @logger = Log4r::Logger.new('vagrant::hostmanager::delete_hosts_file')
-        end
-
-        def call(env)
-          # check if machine is already active
-          #return @app.call(env) if @machine.id
-
-          # check config to see if the hosts file should be update automatically
-          return @app.call(env) unless @machine.config.hostmanager.enabled?
-          @logger.info 'Updating /etc/hosts file automatically'
-
-          # update /etc/hosts file on each active machine
-          delete_guests(@machine,@machine.provider_name)
-
-          # continue the action stack so the machine will be created
-          @app.call(env)
-
-        end
-      end
-    end
-  end
-end
diff --git a/lib/vagrant-hostmanager/action/update_guests.rb b/lib/vagrant-hostmanager/action/update_hosts_file.rb
similarity index 62%
rename from lib/vagrant-hostmanager/action/update_guests.rb
rename to lib/vagrant-hostmanager/action/update_hosts_file.rb
index c495bf0..a281511 100644
--- a/lib/vagrant-hostmanager/action/update_guests.rb
+++ b/lib/vagrant-hostmanager/action/update_hosts_file.rb
@@ -3,28 +3,32 @@ require 'vagrant-hostmanager/hosts_file'
 module VagrantPlugins
   module HostManager
     module Action
-      class UpdateGuests
+      class UpdateHostsFile
         include HostsFile
 
         def initialize(app, env)
           @app = app
           @machine = env[:machine]
-          @logger = Log4r::Logger.new('vagrant::hostmanager::update_guests')
+          @logger = Log4r::Logger.new('vagrant::hostmanager::update_hosts_file')
         end
 
         def call(env)
           # check if machine is already active
-          #return @app.call(env) if @machine.id
+          return @app.call(env) if @machine.id
 
           # check config to see if the hosts file should be update automatically
           return @app.call(env) unless @machine.config.hostmanager.enabled?
           @logger.info 'Updating /etc/hosts file automatically'
 
-          # continue the action stack so the machine will be created
           @app.call(env)
 
-          # update /etc/hosts file on each active machine
-          update_guests(@machine,@machine.provider_name)
+          # update /etc/hosts file on active machines
+          update_guests(@machine.env, @machine.provider_name)
+
+          # update /etc/hosts files on host if enabled
+          if @machine.config.hostmanager.manage_host?
+            update_host(@machine.env, @machine.provider_name)
+          end
         end
       end
     end
diff --git a/lib/vagrant-hostmanager/command.rb b/lib/vagrant-hostmanager/command.rb
index 52a0078..415af03 100644
--- a/lib/vagrant-hostmanager/command.rb
+++ b/lib/vagrant-hostmanager/command.rb
@@ -6,7 +6,7 @@ module VagrantPlugins
       def execute
         options = {}
         opts = OptionParser.new do |o|
-          o.banner = 'Usage: vagrant hostmanager [vm-name]'
+          o.banner = 'Usage: vagrant hostmanager'
           o.separator ''
           o.version = VagrantPlugins::HostManager::VERSION
           o.program_name = 'vagrant hostmanager'
@@ -17,14 +17,13 @@ module VagrantPlugins
           end
         end
 
-        argv = parse_options(opts)
-        return if !argv
+        parse_options(opts)
 
         options[:provider] ||= @env.default_provider
 
-        with_target_vms(argv, options) do |machine|
-          update_guests(machine, machine.provider_name)
-          update_local(machine)
+        update_guests(@env, options[:provider])
+        if (@env.config_global.hostmanager.manage_host?)
+          update_host(@env, options[:provider])
         end
       end
     end
diff --git a/lib/vagrant-hostmanager/config.rb b/lib/vagrant-hostmanager/config.rb
index e537d06..ec83717 100644
--- a/lib/vagrant-hostmanager/config.rb
+++ b/lib/vagrant-hostmanager/config.rb
@@ -2,17 +2,18 @@ module VagrantPlugins
   module HostManager
     class Config < Vagrant.plugin('2', :config)
       attr_accessor :enabled
-      attr_accessor :manage_local
+      attr_accessor :manage_host
       attr_accessor :ignore_private_ip
       attr_accessor :aliases
       attr_accessor :include_offline
 
       alias_method :enabled?, :enabled
       alias_method :include_offline?, :include_offline
+      alias_method :manage_host?, :manage_host
 
       def initialize
         @enabled = false
-        @manage_local = true
+        @manage_local = false
         @ignore_private_ip = UNSET_VALUE
         @aliases = Array.new
         @include_offline = false
@@ -33,9 +34,9 @@ module VagrantPlugins
         errors << validate_bool('hostmanager.include_offline', include_offline)
 
         # check if manage_local option is either true or false
-        if ![TrueClass, FalseClass].include?(manage_local.class)
-          errors << "A value for hostmanager.manage_local can be true or false."
-        end
+        # if ![TrueClass, FalseClass].include?(manage_local.class)
+        #   errors << "A value for hostmanager.manage_local can be true or false."
+        # end
 
         # check if ignore_private_ip option is either true or false (or UNSET_VALUE)
         if @ignore_private_ip != UNSET_VALUE
diff --git a/lib/vagrant-hostmanager/hosts_file.rb b/lib/vagrant-hostmanager/hosts_file.rb
index 395b3e8..d35ae72 100644
--- a/lib/vagrant-hostmanager/hosts_file.rb
+++ b/lib/vagrant-hostmanager/hosts_file.rb
@@ -1,175 +1,85 @@
+require 'tempfile'
+
 module VagrantPlugins
   module HostManager
     module HostsFile
-      def update_guests(machine, provider)
-        machines = []
-
-        env = machine.env
-        # create the temporary hosts file
-        path = env.tmp_path
-
-        #fetch hosts file from each machine
-        #for each machine, ensure all machine entries are updated
-        # add a hosts entry for each active machine matching the provider
+      def update_guests(env, provider)
+        entries = get_entries(env, provider)
         env.active_machines.each do |name, p|
           if provider == p
-            machines << machine = env.machine(name, provider)
-            machine.communicate.download('/etc/hosts',path.join("hosts.#{name}"))
-          end
-        end
-        env.active_machines.each do |name, p|
-            if provider == p
-                machines.each do |m|
-                    @logger.info "Adding entry for #{m.name} to hosts.#{name}"
-                    update_entry(m,path.join("hosts.#{name}"))
-                end
-            end
-            env.machine(name,p).communicate.upload(path.join("hosts.#{name}"), '/tmp/hosts')
-            env.machine(name,p).communicate.sudo("mv /tmp/hosts /etc/hosts")
-            FileUtils.rm(path.join("hosts.#{name}"))
-        end
+            target = env.machine(name, p)
+            next unless target.communicate.ready?
 
-      end
-
-      # delete victim machine from all guests
-      def delete_guests(victim, provider)
-        machines = []
-
-        env = victim.env
-        # create the temporary hosts file
-        path = env.tmp_path
-
-        #fetch hosts file from each machine
-        #for each machine, ensure all machine entries are updated
-        # add a hosts entry for each active machine matching the provider
-        env.active_machines.each do |name, p|
-          if provider == p
-            machines << machine = env.machine(name, provider)
-            machine.communicate.download('/etc/hosts',path.join("hosts.#{name}"))
-            delete_entry(victim,path.join("hosts.#{name}"))               
-            if machine.communicate.ready?
-                machine.env.ui.info I18n.t('vagrant_hostmanager.action.update_guest', {
-                    :name => machine.name
-                })
-                machine.communicate.upload(path.join("hosts.#{name}"), '/tmp/hosts')
-                machine.communicate.sudo("mv /tmp/hosts /etc/hosts")
-            end
-            FileUtils.rm(path.join("hosts.#{name}"))
+            file = env.tmp_path.join("hosts.#{name}")
+            target.communicate.download('/etc/hosts', file)
+            update_file(file, entries, env.tmp_path)
+            target.communicate.upload(file, '/tmp/hosts')
+            target.communicate.sudo('mv /tmp/hosts /etc/hosts')
+            FileUtils.rm(file)
           end
         end
       end
 
-
-        # define a lambda for looking up a machine's ip address
-      def  get_ip_address(machine)
-          ip = nil
-          if machine.config.hostmanager.ignore_private_ip != true
-            machine.config.vm.networks.each do |network|
-              key, options = network[0], network[1]
-              ip = options[:ip] if key == :private_network
-              next if ip
-            end
-          end
-          ip || (machine.ssh_info ? machine.ssh_info[:host] : nil)
-        end
-
-      def update_entry(machine,file_name,sudo=false)
-         delete_entry(machine,file_name,sudo)
-         
-         host = machine.config.vm.hostname || name
-         id = machine.id
-         ip = get_ip_address(machine)
-         host_aliases = machine.config.hostmanager.aliases.join("\s").chomp
-         host_entry = "#{ip}\t#{host}\s#{host_aliases}\s# VAGRANT: #{id}\n" 
-         @logger.info "Adding /etc/hosts entry: #{host_entry}"
-         temp_file_name = Dir::Tmpname.make_tmpname(File.join(machine.env.tmp_path,'hostmanager'), nil) 
-         FileUtils.cp(file_name, temp_file_name)
-         File.open(temp_file_name,'a') do |tempfile|
-             @logger.info "writing #{host_entry} to #{tempfile.path}"
-             tempfile << host_entry
-         end
-
-         if sudo == false
-            @logger.info "copy #{temp_file_name} #{file_name}"
-            FileUtils.cp(temp_file_name,file_name)
-         else
-            machine.env.ui.info I18n.t('vagrant_hostmanager.action.run_sudo')
-            @logger.warn "Running sudo to replace local hosts file, enter your local password if prompted..."
-            @logger.info `sudo cp -v #{temp_file_name} #{file_name}`
-         end
-         FileUtils.rm(temp_file_name)
-      end
-
-      def delete_entry(machine,file_name,sudo=false)
-          host = machine.config.vm.hostname || name
-          temp_file_name = Dir::Tmpname.make_tmpname(File.join(machine.env.tmp_path,'hostmanager'), nil) 
-          if not machine.id.nil?
-              tempfile = File.open(temp_file_name,'w') do |f| 
-                File.open(file_name,'r').each_line do |line|
-                  if line.match(/#{machine.id}$/).nil?
-                     f << line
-                  else
-                      @logger.info "Matched #{machine.id}"
-                  end
-                end
-              end
-              if sudo == false
-                @logger.info "copy #{temp_file_name} #{file_name}"
-                    FileUtils.cp(temp_file_name,file_name)
-              else
-                  machine.env.ui.info I18n.t('vagrant_hostmanager.action.run_sudo')
-                  @logger.info `sudo cp -v #{temp_file_name} #{file_name}`
-              end
-              FileUtils.rm(temp_file_name)
-          else
-              @logger.warn "Machine id to delete was empty, skipping..."
-          end
-      end
-
-      def update_local(machine)
-         return if machine.id.nil?
-         update_entry(machine,'/etc/hosts',true)
-      end
-
-      def delete_local(machine)
-          return if machine.id.nil?
-          delete_entry(machine,'/etc/hosts',true)
-      end
-
-      def publish_local(tempfile)
-          @logger.info `sudo cp -v #{tempfile} /etc/hosts`
-      end
-
-
-      # Copy the temporary hosts file to the specified machine overwritting
-      # the existing /etc/hosts file.
-      def update(machine)
-        path = machine.env.tmp_path.join('hosts')
-        if machine.communicate.ready?
-          machine.env.ui.info I18n.t('vagrant_hostmanager.action.update_guest', {
-            :name => machine.name
-          })
-          machine.communicate.download(path, '/etc/hosts')
-          machine.communicate.upload(path, '/tmp/hosts')
-          machine.communicate.sudo("mv /tmp/hosts /etc/hosts")
-        end
+      def update_host(env, provider)
+        entries = get_entries(env, provider)
+        file = env.tmp_path.join('hosts.local')
+        FileUtils.cp('/etc/hosts', file)
+        update_file(file, entries, env.tmp_path)
+        `sudo cp #{file} /etc/hosts`
       end
 
       private
-      # Either use the active machines, or loop over all available machines and
-      # get those with the same provider (aka, ignore boxes that throw MachineNotFound errors).
-      #
-      # Returns an array with the same structure as env.active_machines:
-      # [ [:machine, :virtualbox], [:foo, :virtualbox] ]
+
+      def update_file(file, entries, tmp_path)
+        tmp_file = Tempfile.open('hostmanager', tmp_path, 'a')
+        begin
+          File.open(file).each_line do |line|
+            tmp_file << line unless line =~ /# VAGRANT ID:/
+          end
+          entries.each { |entry| tmp_file << entry }
+        ensure
+          tmp_file.close
+          FileUtils.cp(tmp_file, file)
+          tmp_file.unlink
+        end
+      end
+
+      def get_entries(env, provider)
+        entries = []
+        get_machines(env, provider).each do |name, p|
+          if provider == p
+            machine = env.machine(name, p)
+            host = machine.config.vm.hostname || name
+            id = machine.id
+            ip = get_ip_address(machine)
+            aliases = machine.config.hostmanager.aliases.join(' ').chomp
+            entries <<  "#{ip}\t#{host} #{aliases}\t# VAGRANT ID: #{id}\n"
+          end
+        end
+
+        entries
+      end
+
+      def get_ip_address(machine)
+        ip = nil
+        if machine.config.hostmanager.ignore_private_ip != true
+          machine.config.vm.networks.each do |network|
+            key, options = network[0], network[1]
+            ip = options[:ip] if key == :private_network
+            next if ip
+          end
+        end
+        ip || (machine.ssh_info ? machine.ssh_info[:host] : nil)
+      end
+
       def get_machines(env, provider)
         if env.config_global.hostmanager.include_offline?
           machines = []
           env.machine_names.each do |name|
             begin
-              m = env.machine(name, provider)
+              env.machine(name, provider)
               machines << [name, provider]
-            rescue Vagrant::Errors::MachineNotFound => ex
-              # ignore this box.
+            rescue Vagrant::Errors::MachineNotFound
             end
           end
           machines
diff --git a/lib/vagrant-hostmanager/plugin.rb b/lib/vagrant-hostmanager/plugin.rb
index d77a737..f73ce6e 100644
--- a/lib/vagrant-hostmanager/plugin.rb
+++ b/lib/vagrant-hostmanager/plugin.rb
@@ -1,7 +1,4 @@
-require 'vagrant-hostmanager/action/delete_guests'
-require 'vagrant-hostmanager/action/update_guests'
-require 'vagrant-hostmanager/action/update_local_entry'
-require 'vagrant-hostmanager/action/delete_local_entry'
+require 'vagrant-hostmanager/action/update_hosts_file'
 
 module VagrantPlugins
   module HostManager
@@ -19,18 +16,12 @@ module VagrantPlugins
         Config
       end
 
-      action_hook(self::ALL_ACTIONS) do |hook|
-        hook.after(VagrantPlugins::ProviderVirtualBox::Action::Boot, Action::UpdateGuests)
-      end
-
       action_hook(:hostmanager, :machine_action_up) do |hook|
-        hook.prepend(Action::UpdateGuests)
-        hook.prepend(Action::UpdateLocalEntry)
+        hook.prepend(Action::UpdateHostsFile)
       end
 
       action_hook(:hostmanager, :machine_action_destroy) do |hook|
-        hook.prepend(Action::DeleteLocalEntry)
-        hook.prepend(Action::DeleteGuests)
+        hook.append(Action::UpdateHostsFile)
       end
 
       provisioner(:hostmanager) do
diff --git a/lib/vagrant-hostmanager/provisioner.rb b/lib/vagrant-hostmanager/provisioner.rb
index 65d43db..d726636 100644
--- a/lib/vagrant-hostmanager/provisioner.rb
+++ b/lib/vagrant-hostmanager/provisioner.rb
@@ -4,8 +4,10 @@ module VagrantPlugins
       include HostsFile
 
       def provision
-        generate(@machine.env, @machine.box.provider.to_sym)
-        update(@machine)
+        update_guests(@machine.env, @machine.provider_name)
+        if @machine.env.config_global.hostmanager.manage_host?
+          update_host(@machine.env, @machine.provider_name)
+        end
       end
     end
   end
diff --git a/lib/vagrant-hostmanager/version.rb b/lib/vagrant-hostmanager/version.rb
index bf8a05a..58ffcee 100644
--- a/lib/vagrant-hostmanager/version.rb
+++ b/lib/vagrant-hostmanager/version.rb
@@ -1,5 +1,5 @@
 module VagrantPlugins
   module HostManager
-    VERSION = '0.4.0'
+    VERSION = '1.0.0'
   end
 end
diff --git a/test/Vagrantfile b/test/Vagrantfile
index f1923cf..011c54e 100644
--- a/test/Vagrantfile
+++ b/test/Vagrantfile
@@ -4,14 +4,17 @@
 Vagrant.require_plugin('vagrant-hostmanager')
 
 Vagrant.configure('2') do |config|
-  config.vm.box = 'precise64-chef11.2'
-  config.vm.box_url = 'https://opscode-vm.s3.amazonaws.com/vagrant/opscode_ubuntu-12.04_chef-11.2.0.box'
+  config.vm.box = 'precise64'
+  config.vm.box_url = 'http://cloud-images.ubuntu.com/precise/current/precise-server-cloudimg-vagrant-amd64-disk1.box'
 
   config.hostmanager.enabled = true
+  config.hostmanager.manage_host = true
+  config.hostmanager.include_offline = true
 
   config.vm.define :server1 do |server|
     server.vm.hostname = 'fry'
     server.vm.network :private_network, :ip => '10.0.5.2'
+    server.hostmanager.aliases = %w(test-alias)
   end
 
   config.vm.define :server2 do |server|
diff --git a/vagrant-hostmanager.gemspec b/vagrant-hostmanager.gemspec
index 40131ef..74aa9f0 100644
--- a/vagrant-hostmanager.gemspec
+++ b/vagrant-hostmanager.gemspec
@@ -1,4 +1,5 @@
 # -*- encoding: utf-8 -*-
+
 lib = File.expand_path('../lib', __FILE__)
 $LOAD_PATH.unshift(lib) unless $LOAD_PATH.include?(lib)
 require 'vagrant-hostmanager/version'
@@ -14,4 +15,7 @@ Gem::Specification.new do |gem|
   gem.files         = `git ls-files`.split($/)
   gem.test_files    = gem.files.grep(%r{^(test|spec|features)/})
   gem.require_paths = ['lib']
+
+  gem.add_development_dependency 'bundler', '~> 1.3'
+  gem.add_development_dependency 'rake'
 end