add amp4e computers command.

7 files changed, 133 insertions, 2 deletions

diff --git a/lib/ats.rb b/lib/ats.rb
index 22e77c4..ccd264f 100644
--- a/lib/ats.rb
+++ b/lib/ats.rb
@@ -1,12 +1,15 @@
+require 'base64'
 require 'json'
 require 'logger'
 require 'net/http'
 require 'yaml'
 
-require 'ats/version'
-
 require 'ats/configuration'
 require 'ats/http_api'
+require 'ats/version'
+
+require 'ats/amp4e/api'
+require 'ats/amp4e/computers'
 require 'ats/threat_grid/api'
 require 'ats/threat_grid/organizations'
 require 'ats/threat_grid/samples'
diff --git a/lib/ats/amp4e/api.rb b/lib/ats/amp4e/api.rb
new file mode 100644
index 0000000..202df1a
--- /dev/null
+++ b/lib/ats/amp4e/api.rb
@@ -0,0 +1,55 @@
+module ATS
+  module AMP4E
+    class API
+      HEADERS = {
+        'Content-Type' => 'application/json',
+        'Accept' => 'application/json',
+        'User-Agent' => "RubyGems/ATS #{ATS::VERSION}",
+      }.freeze
+
+      attr_reader :http, :profile, :configuration
+
+      def initialize(
+        api: HttpAPI.new(headers: HEADERS),
+        profile: :default,
+        configuration: ATS.configuration
+      )
+        @http = api
+        @profile = profile.to_s
+        @configuration = configuration
+      end
+
+      def computers
+        ATS::AMP4E::Computers.new(self)
+      end
+
+      def get(url, params: {}, version: 1)
+        http.get(build_uri(url, version: version), headers: headers, body: params) do |request, response|
+          JSON.parse(response.body, symbolize_names: true)
+        end
+      end
+
+      private
+
+      def build_uri(relative_url, version:)
+        URI.parse("#{api_host}/v#{version}/#{relative_url}")
+      end
+
+      def client_id
+        configuration[profile]['amp4e']['client_id']
+      end
+
+      def client_secret
+        configuration[profile]['amp4e']['client_secret']
+      end
+
+      def api_host
+        configuration[profile]['amp4e']['api_host']
+      end
+
+      def headers
+        { AUTHORIZATION: "Basic #{Base64.strict_encode64("#{client_id}:#{client_secret}")}" }
+      end
+    end
+  end
+end
diff --git a/lib/ats/amp4e/computers.rb b/lib/ats/amp4e/computers.rb
new file mode 100644
index 0000000..c1f49a1
--- /dev/null
+++ b/lib/ats/amp4e/computers.rb
@@ -0,0 +1,15 @@
+module ATS
+  module AMP4E
+    class Computers
+      attr_reader :api
+
+      def initialize(api)
+        @api = api
+      end
+
+      def list
+        api.get("computers")
+      end
+    end
+  end
+end
diff --git a/lib/ats/cli.rb b/lib/ats/cli.rb
index c8d360d..682bc50 100644
--- a/lib/ats/cli.rb
+++ b/lib/ats/cli.rb
@@ -2,6 +2,9 @@ require 'ats'
 
 require 'thor'
 require 'ats/cli/threat_grid'
+require 'ats/cli/amp4e/command'
+require 'ats/cli/amp4e/computers'
+require 'ats/cli/amp4e/application'
 
 module ATS
   module CLI
@@ -11,6 +14,9 @@ module ATS
       desc 'threatgrid SUBCOMMAND ...ARGS', 'interact with the Threat Grid API'
       subcommand 'threatgrid', ThreatGrid::Application
 
+      desc 'amp4e SUBCOMMAND ...ARGS', 'interact with the AMP for Endpoints API'
+      subcommand 'amp4e', AMP4E::Application
+
       desc 'version', 'Display the current version'
       def version
         say ATS::VERSION
diff --git a/lib/ats/cli/amp4e/application.rb b/lib/ats/cli/amp4e/application.rb
new file mode 100644
index 0000000..34e5693
--- /dev/null
+++ b/lib/ats/cli/amp4e/application.rb
@@ -0,0 +1,12 @@
+module ATS
+  module CLI
+    module AMP4E
+      class Application < Thor
+        class_option :profile, default: :default, required: false
+
+        desc 'computers SUBCOMMAND ...ARGS', 'interact with the Threat Grid API'
+        subcommand :computers, ATS::CLI::AMP4E::Computers
+      end
+    end
+  end
+end
diff --git a/lib/ats/cli/amp4e/command.rb b/lib/ats/cli/amp4e/command.rb
new file mode 100644
index 0000000..9e6a93a
--- /dev/null
+++ b/lib/ats/cli/amp4e/command.rb
@@ -0,0 +1,26 @@
+module ATS
+  module CLI
+    module AMP4E
+      class Command < Thor
+        class_option :profile, default: :default, required: false
+
+        def self.printable_commands(*args)
+          super.map do |x|
+            x[0] = x[0].gsub(/^ats/, 'ats amp4e')
+            x
+          end
+        end
+
+        protected
+
+        def api
+          ATS::AMP4E::API.new(profile: options['profile'])
+        end
+
+        def print_json(json)
+          say JSON.pretty_generate(json), :green
+        end
+      end
+    end
+  end
+end
diff --git a/lib/ats/cli/amp4e/computers.rb b/lib/ats/cli/amp4e/computers.rb
new file mode 100644
index 0000000..741584f
--- /dev/null
+++ b/lib/ats/cli/amp4e/computers.rb
@@ -0,0 +1,14 @@
+module ATS
+  module CLI
+    module AMP4E
+      class Computers < Command
+        class_option :profile, default: :default, required: false
+
+        desc 'list', 'list'
+        def list
+          print_json api.computers.list
+        end
+      end
+    end
+  end
+end