Setup DigitalOcean plugin to use Metadata API

2 files changed, 95 insertions, 41 deletions

diff --git a/lib/ohai/mixin/do_metadata.rb b/lib/ohai/mixin/do_metadata.rb
new file mode 100644
index 00000000..ec4a08a8
--- /dev/null
+++ b/lib/ohai/mixin/do_metadata.rb
@@ -0,0 +1,73 @@
+
+# Author:: Dylan Page (<dpage@digitalocean.com>)
+# License:: Apache License, Version 2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+require 'net/http'
+require 'socket'
+
+module Ohai
+  module Mixin
+    module DOMetadata
+
+      DO_METADATA_ADDR = "169.254.269.254" unless defined?(DO_METADATA_ADDR)
+      DO_METADATA_URL = "/metadata/v1.json" unless defined?(DO_METADATA_URL)
+
+      def can_metadata_connect?(addr, port, timeout=2)
+        t = Socket.new(Socket::Constants::AF_INET, Socket::Constants::SOCK_STREAM, 0)
+        saddr = Socket.pack_sockaddr_in(port, addr)
+        connected = false
+
+        begin
+          t.connect_nonblock(saddr)
+        rescue Errno::EINPROGRESS
+          r,w,e = IO::select(nil,[t],nil,timeout)
+          if !w.nil?
+            connected = true
+          else
+            begin
+              t.connect_nonblock(saddr)
+            rescue Errno::EISCONN
+              t.close
+              connected = true
+            rescue SystemCallError
+            end
+          end
+        rescue SystemCallError
+        end
+        Ohai::Log.debug("DOMetadata mixin: can_metadata_connect? == #{connected}")
+        connected
+      end
+
+      def http_client
+        Net::HTTP.start(DO_METADATA_ADDR).tap {|h| h.read_timeout = 6}
+      end
+
+      def fetch_metadata()
+        uri = "#{DO_METADATA_URL}"
+        response = http_client.get(uri)
+        case response.code
+        when "200"
+          response.body
+        when "404"
+          Ohai::Log.debug("Encountered 404 response retreiving Digital Ocean metadata: #{uri} ; continuing.")
+          Hash.new
+        else
+          raise "Encountered error retrieving Digital Ocean metadata (#{uri} returned #{response.code} response)"
+        end
+      end
+
+    end
+  end
+end
diff --git a/lib/ohai/plugins/digital_ocean.rb b/lib/ohai/plugins/digital_ocean.rb
index a06c3c67..19cb087b 100644
--- a/lib/ohai/plugins/digital_ocean.rb
+++ b/lib/ohai/plugins/digital_ocean.rb
@@ -1,4 +1,5 @@
 #
+# Author:: Dylan Page (<dpage@digitalocean.com>)
 # Author:: Stafford Brunk (<stafford.brunk@gmail.com>)
 # License:: Apache License, Version 2.0
 #
@@ -14,65 +15,45 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-require "ohai/util/ip_helper"
+require 'ohai/mixin/do_metadata'
+require 'yaml'
 
 Ohai.plugin(:DigitalOcean) do
-  include Ohai::Util::IpHelper
+  include Ohai::Mixin::DOMetadata
 
-  DIGITALOCEAN_FILE = "/etc/digitalocean" unless defined?(DIGITALOCEAN_FILE)
+  DO_CLOUD_INIT_FILE = "/etc/cloud/cloud.cfg" unless defined?(DO_CLOUD_INIT_FILE)
 
   provides "digital_ocean"
-  depends "network/interfaces"
 
-  def extract_droplet_ip_addresses
-    addresses = Mash.new({ "v4" => [], "v6" => [] })
-    network[:interfaces].each_value do |iface|
-      iface[:addresses].each do |address, details|
-        next if details[:family] == "lladdr" || loopback?(address)
+  depends "network/interfaces"
 
-        ip = IPAddress(address)
-        type = digital_ocean_address_type(ip)
-        address_hash = build_address_hash(ip, details)
-        addresses[type] << address_hash
-      end
+  def has_do_init?
+    if File.exist?(DO_CLOUD_INIT_FILE)
+      datasource = YAML.load_file(DO_CLOUD_INIT_FILE)
+      puts datasource.inspect
+      Ohai::Log.debug("digital_ocean plugin: has_do_init? == true")
+      true
+    else
+      Ohai::Log.debug("digital_ocean plugin: has_do_init? == false")
+      false
     end
-    addresses
   end
 
-  def build_address_hash(ip, details)
-    address_hash = Mash.new({
-      "ip_address" => ip.address,
-      "type" => private_address?(ip.address) ? "private" : "public",
-    })
+  def looks_like_digital_ocean?
+    return true if hint?("digital_ocean")
 
-    if ip.ipv4?
-      address_hash["netmask"] = details[:netmask]
-    elsif ip.ipv6?
-      address_hash["cidr"] = ip.prefix
+    if has_do_init?
+     return true if can_metadata_connect?(Ohai::Mixin::DOMetadata::DO_METADATA_ADDR,80)
     end
-    address_hash
-  end
-
-  def digital_ocean_address_type(ip)
-    ip.ipv4? ? "v4" : "v6"
-  end
-
-  def looks_like_digital_ocean?
-    hint?("digital_ocean") || File.exist?(DIGITALOCEAN_FILE)
   end
 
   collect_data do
     if looks_like_digital_ocean?
       Ohai::Log.debug("Plugin Digitalocean: looks_like_digital_ocean? == true")
       digital_ocean Mash.new
-      hint = hint?("digital_ocean") || {}
-      hint.each { |k, v| digital_ocean[k] = v unless k == "ip_addresses" }
-
-      # Extract actual ip addresses
-      # The networks sub-hash is structured similarly to how
-      # Digital Ocean's v2 API structures things:
-      # https://developers.digitalocean.com/#droplets
-      digital_ocean[:networks] = extract_droplet_ip_addresses
+      fetch_metadata.each do |k, v|
+        digital_ocean[k] = v
+      end
     else
       Ohai::Log.debug("Plugin Digitalocean: No hints present for and doesn't look like digitalocean")
       false