Merge pull request #907 from Fodoj/master

Fetch AWS Account ID from metadata

3 files changed, 60 insertions, 1 deletions

diff --git a/lib/ohai/mixin/ec2_metadata.rb b/lib/ohai/mixin/ec2_metadata.rb
index eb4e4e6a..1797105d 100644
--- a/lib/ohai/mixin/ec2_metadata.rb
+++ b/lib/ohai/mixin/ec2_metadata.rb
@@ -200,6 +200,20 @@ module Ohai
         response.code == "200" ? response.body : nil
       end
 
+      def fetch_dynamic_data
+        @fetch_dynamic_data ||= begin
+          api_version = best_api_version
+          return {} if api_version.nil?
+          response = http_client.get("/#{api_version}/dynamic/instance-identity/document/")
+
+          if json?(response.body) && response.code == "200"
+            FFI_Yajl::Parser.parse(response.body)
+          else
+            {}
+          end
+        end
+      end
+
       private
 
       def expand_path(file_name)
diff --git a/lib/ohai/plugins/ec2.rb b/lib/ohai/plugins/ec2.rb
index c7f78568..583480a4 100644
--- a/lib/ohai/plugins/ec2.rb
+++ b/lib/ohai/plugins/ec2.rb
@@ -96,6 +96,7 @@ Ohai.plugin(:EC2) do
         ec2[k] = v
       end
       ec2[:userdata] = fetch_userdata
+      ec2[:account_id] = fetch_dynamic_data["accountId"]
       # ASCII-8BIT is equivalent to BINARY in this case
       if ec2[:userdata] && ec2[:userdata].encoding.to_s == "ASCII-8BIT"
         Ohai::Log.debug("Plugin EC2: Binary UserData Found. Storing in base64")
diff --git a/spec/unit/plugins/ec2_spec.rb b/spec/unit/plugins/ec2_spec.rb
index df0f00d0..fc83c8cc 100644
--- a/spec/unit/plugins/ec2_spec.rb
+++ b/spec/unit/plugins/ec2_spec.rb
@@ -49,7 +49,7 @@ describe Ohai::System, "plugin ec2" do
       allow(t).to receive(:connect_nonblock).and_raise(Errno::EINPROGRESS)
       allow(Socket).to receive(:new).and_return(t)
       expect(@http_client).to receive(:get).
-        with("/").twice.
+        with("/").exactly(3).times.
         and_return(double("Net::HTTP Response", :body => "2012-01-12", :code => "200"))
     end
 
@@ -71,6 +71,9 @@ describe Ohai::System, "plugin ec2" do
         expect(@http_client).to receive(:get).
           with("/2012-01-12/user-data/").
           and_return(double("Net::HTTP Response", :body => "By the pricking of my thumb...", :code => "200"))
+        expect(@http_client).to receive(:get).
+          with("/2012-01-12/dynamic/instance-identity/document/").
+          and_return(double("Net::HTTP Response", :body => "{\"accountId\":\"4815162342\"}", :code => "200"))
 
         plugin.run
 
@@ -89,6 +92,9 @@ describe Ohai::System, "plugin ec2" do
         expect(@http_client).to receive(:get).
           with("/2012-01-12/user-data/").
           and_return(double("Net::HTTP Response", :body => "^_<8B>^H^H<C7>U^@^Csomething^@KT<C8><C9>,)<C9>IU(I-.I<CB><CC>I<E5>^B^@^Qz<BF><B0>^R^@^@^@", :code => "200"))
+        expect(@http_client).to receive(:get).
+          with("/2012-01-12/dynamic/instance-identity/document/").
+          and_return(double("Net::HTTP Response", :body => "{\"accountId\":\"4815162342\"}", :code => "200"))
 
         plugin.run
 
@@ -98,6 +104,29 @@ describe Ohai::System, "plugin ec2" do
         expect(plugin[:ec2]["security_groups"]).to eql %w{group1 group2}
         expect(plugin[:ec2]["userdata"]).to eq(Base64.decode64("Xl88OEI+XkheSDxDNz5VXkBeQ3NvbWV0aGluZ15AS1Q8Qzg+PEM5PiwpPEM5\nPklVKEktLkk8Q0I+PENDPkk8RTU+XkJeQF5RejxCRj48QjA+XlJeQF5AXkA="))
       end
+
+      it "fetches AWS account id" do
+        paths.each do |name, body|
+          expect(@http_client).to receive(:get).
+            with("/2012-01-12/#{name}").
+            and_return(double("Net::HTTP Response", :body => body, :code => "200"))
+        end
+        expect(@http_client).to receive(:get).
+          with("/2012-01-12/user-data/").
+          and_return(double("Net::HTTP Response", :body => "^_<8B>^H^H<C7>U^@^Csomething^@KT<C8><C9>,)<C9>IU(I-.I<CB><CC>I<E5>^B^@^Qz<BF><B0>^R^@^@^@", :code => "200"))
+        expect(@http_client).to receive(:get).
+          with("/2012-01-12/dynamic/instance-identity/document/").
+          and_return(double("Net::HTTP Response", :body => "{\"accountId\":\"4815162342\"}", :code => "200"))
+
+        plugin.run
+
+        expect(plugin[:ec2]).not_to be_nil
+        expect(plugin[:ec2]["instance_type"]).to eq("c1.medium")
+        expect(plugin[:ec2]["ami_id"]).to eq("ami-5d2dc934")
+        expect(plugin[:ec2]["security_groups"]).to eql %w{group1 group2}
+        expect(plugin[:ec2]["account_id"]).to eq("4815162342")
+      end
+
     end
 
     it "parses ec2 network/ directory as a multi-level hash" do
@@ -122,6 +151,9 @@ describe Ohai::System, "plugin ec2" do
       expect(@http_client).to receive(:get).
         with("/2012-01-12/user-data/").
         and_return(double("Net::HTTP Response", :body => "By the pricking of my thumb...", :code => "200"))
+      expect(@http_client).to receive(:get).
+        with("/2012-01-12/dynamic/instance-identity/document/").
+        and_return(double("Net::HTTP Response", :body => "{\"accountId\":\"4815162342\"}", :code => "200"))
 
       plugin.run
 
@@ -150,6 +182,9 @@ describe Ohai::System, "plugin ec2" do
         expect(@http_client).to receive(:get).
           with("/2012-01-12/user-data/").
           and_return(double("Net::HTTP Response", :body => "By the pricking of my thumb...", :code => "200"))
+        expect(@http_client).to receive(:get).
+          with("/2012-01-12/dynamic/instance-identity/document/").
+          and_return(double("Net::HTTP Response", :body => "{\"accountId\":\"4815162342\"}", :code => "200"))
 
         plugin.run
 
@@ -180,6 +215,9 @@ describe Ohai::System, "plugin ec2" do
         expect(@http_client).to receive(:get).
           with("/2012-01-12/user-data/").
           and_return(double("Net::HTTP Response", :body => "By the pricking of my thumb...", :code => "200"))
+        expect(@http_client).to receive(:get).
+          with("/2012-01-12/dynamic/instance-identity/document/").
+          and_return(double("Net::HTTP Response", :body => "{\"accountId\":\"4815162342\"}", :code => "200"))
 
         plugin.run
 
@@ -219,6 +257,9 @@ describe Ohai::System, "plugin ec2" do
       expect(@http_client).to receive(:get).
         with("/2012-01-12/user-data/").
         and_return(double("Net::HTTP Response", :body => "By the pricking of my thumb...", :code => "200"))
+      expect(@http_client).to receive(:get).
+        with("/2012-01-12/dynamic/instance-identity/document/").
+        and_return(double("Net::HTTP Response", :body => "{\"accountId\":\"4815162342\"}", :code => "200"))
 
       plugin.run
 
@@ -238,6 +279,9 @@ describe Ohai::System, "plugin ec2" do
       expect(@http_client).to receive(:get).
         with("/2012-01-12/user-data/").
         and_return(double("Net::HTTP Response", :body => "By the pricking of my thumb...", :code => "200"))
+      expect(@http_client).to receive(:get).
+        with("/2012-01-12/dynamic/instance-identity/document/").
+        and_return(double("Net::HTTP Response", :body => "{\"accountId\":\"4815162342\"}", :code => "200"))
 
       plugin.run