From 98741e5e4a4619879681edd446a6891aefe16016 Mon Sep 17 00:00:00 2001
From: Brian Cain <brian.cain@puppetlabs.com>
Date: Fri, 25 Sep 2015 14:03:47 -0700
Subject: [PATCH] Fix token auth for methods

Properly set the header to be X-AUTH-TOKEN for requests to the pooler.
---
 lib/vmfloaty.rb        |  6 ++++--
 lib/vmfloaty/http.rb   | 20 --------------------
 lib/vmfloaty/pooler.rb | 31 ++++++++++++++++++-------------
 vmfloaty.gemspec       |  2 +-
 4 files changed, 23 insertions(+), 36 deletions(-)

diff --git a/lib/vmfloaty.rb b/lib/vmfloaty.rb
index 3b17f38..33ce187 100644
--- a/lib/vmfloaty.rb
+++ b/lib/vmfloaty.rb
@@ -10,7 +10,7 @@ class Vmfloaty
   include Commander::Methods
 
   def run
-    program :version, '0.2.3'
+    program :version, '0.2.4'
     program :description, 'A CLI helper tool for Puppet Labs vmpooler to help you stay afloat'
 
     config = read_config
@@ -116,13 +116,15 @@ class Vmfloaty
       c.description = ''
       c.example 'Schedules the deletion of a host or hosts', 'floaty delete myhost1,myhost2 --url http://vmpooler.example.com'
       c.option '--verbose', 'Enables verbose output'
+      c.option '--token STRING', String, 'Token for vmpooler'
       c.option '--url STRING', String, 'URL of vmpooler'
       c.action do |args, options|
         verbose = options.verbose || config['verbose']
         hosts = args[0]
+        token = options.token || config['token']
         url = options.url ||= config['url']
 
-        Pooler.delete(verbose, url, hosts)
+        Pooler.delete(verbose, url, hosts, token)
       end
     end
 
diff --git a/lib/vmfloaty/http.rb b/lib/vmfloaty/http.rb
index ecc6b2f..bd1cc30 100644
--- a/lib/vmfloaty/http.rb
+++ b/lib/vmfloaty/http.rb
@@ -37,24 +37,4 @@ class Http
     return conn
   end
 
-  def self.get_conn_with_token(verbose, url, token)
-    if url.nil?
-      STDERR.puts "The url you provided was empty"
-      exit 1
-    end
-
-    if token.nil?
-      STDERR.puts "The token you provided was empty"
-      exit 1
-    end
-
-    conn = Faraday.new(:url => url, :ssl => {:verify => false}) do |faraday|
-      faraday.request :url_encoded
-      faraday.request :token_auth, token
-      faraday.response :logger if verbose
-      faraday.adapter Faraday.default_adapter
-    end
-
-    return conn
-  end
 end
diff --git a/lib/vmfloaty/pooler.rb b/lib/vmfloaty/pooler.rb
index 963f0fc..2c82e6c 100644
--- a/lib/vmfloaty/pooler.rb
+++ b/lib/vmfloaty/pooler.rb
@@ -20,11 +20,9 @@ class Pooler
 
   def self.retrieve(verbose, os_type, token, url)
     os = os_type.gsub(',','+')
-    if token.nil?
-      conn = Http.get_conn(verbose, url)
-    else
-      conn = Http.get_conn_with_token(verbose, url, token)
-      conn.headers['X-AUTH-TOKEN']
+    conn = Http.get_conn(verbose, url)
+    if token
+      conn.headers['X-AUTH-TOKEN'] = token
     end
 
     response = conn.post "/vm/#{os}"
@@ -42,15 +40,18 @@ class Pooler
       modify_body['tags'] = tags
     end
 
-    conn = Http.get_conn_with_token(verbose, url, token)
-    conn.headers['X-AUTH-TOKEN']
+    conn = Http.get_conn(verbose, url)
+    conn.headers['X-AUTH-TOKEN'] = token
 
-    response = conn.put "/vm/#{hostname}"
+    response = conn.put do |req|
+      req.url "/vm/#{hostname}"
+      req.body = modify_body
+    end
     res_body = JSON.parse(response.body)
     res_body
   end
 
-  def self.delete(verbose, url, hostnames)
+  def self.delete(verbose, url, hostnames, token)
     if hostnames.nil?
       STDERR.puts "You did not provide any hosts to delete"
       exit 1
@@ -59,6 +60,10 @@ class Pooler
     hosts = hostnames.split(',')
     conn = Http.get_conn(verbose, url)
 
+    if token
+      conn.headers['X-AUTH-TOKEN'] = token
+    end
+
     hosts.each do |host|
       puts "Scheduling host #{host} for deletion"
       response = conn.delete "/vm/#{host}"
@@ -97,8 +102,8 @@ class Pooler
   end
 
   def self.snapshot(verbose, url, hostname, token)
-    conn = Http.get_conn_with_token(verbose, url, token)
-    conn.headers['X-AUTH-TOKEN']
+    conn = Http.get_conn(verbose, url)
+    conn.headers['X-AUTH-TOKEN'] = token
 
     response = conn.post "/vm/#{hostname}/snapshot"
     res_body = JSON.parse(response.body)
@@ -106,8 +111,8 @@ class Pooler
   end
 
   def self.revert(verbose, url, hostname, token, snapshot_sha)
-    conn = Http.get_conn_with_token(verbose, url, token)
-    conn.headers['X-AUTH-TOKEN']
+    conn = Http.get_conn(verbose, url)
+    conn.headers['X-AUTH-TOKEN'] = token
 
     response = conn.post "/vm/#{hostname}/snapshot/#{snapshot}"
     res_body = JSON.parse(response.body)
diff --git a/vmfloaty.gemspec b/vmfloaty.gemspec
index 262e6ed..228b3a4 100644
--- a/vmfloaty.gemspec
+++ b/vmfloaty.gemspec
@@ -1,6 +1,6 @@
 Gem::Specification.new do |s|
   s.name = 'vmfloaty'
-  s.version = '0.2.3'
+  s.version = '0.2.4'
   s.authors = ['Brian Cain']
   s.email = ['brian.cain@puppetlabs.com']
   s.license = 'Apache'