From 50a12d672efc5186b97fbb6a4528aa504da646b0 Mon Sep 17 00:00:00 2001
From: Scott Schneider <sschneider@puppetlabs.com>
Date: Wed, 11 Mar 2015 13:44:50 -0700
Subject: [PATCH] Allow authorized_keys to be overridden on VM checkout

---
 README.md              |  1 +
 lib/vmpooler.rb        |  2 ++
 lib/vmpooler/api/v1.rb | 15 +++++++++++++++
 vmpooler.yaml.example  |  6 ++++++
 4 files changed, 24 insertions(+)

diff --git a/README.md b/README.md
index 4588384..fecc243 100644
--- a/README.md
+++ b/README.md
@@ -17,6 +17,7 @@ At [Puppet Labs](http://puppetlabs.com) we run acceptance tests on thousands of
 vmpooler requires the following Ruby gems be installed:
 
 - [json](http://rubygems.org/gems/json)
+- [net/scp](http://rubygems.org/gems/net-scp)
 - [rbvmomi](http://rubygems.org/gems/rbvmomi)
 - [redis](http://rubygems.org/gems/redis)
 - [sinatra](http://rubygems.org/gems/sinatra)
diff --git a/lib/vmpooler.rb b/lib/vmpooler.rb
index 97831de..9f04125 100644
--- a/lib/vmpooler.rb
+++ b/lib/vmpooler.rb
@@ -3,10 +3,12 @@ require 'rubygems' unless defined?(Gem)
 module Vmpooler
   require 'date'
   require 'json'
+  require 'net/scp'
   require 'open-uri'
   require 'rbvmomi'
   require 'redis'
   require 'sinatra/base'
+  require 'stringio'
   require 'time'
   require 'timeout'
   require 'yaml'
diff --git a/lib/vmpooler/api/v1.rb b/lib/vmpooler/api/v1.rb
index 0134b8c..e7a22c4 100644
--- a/lib/vmpooler/api/v1.rb
+++ b/lib/vmpooler/api/v1.rb
@@ -321,6 +321,10 @@ module Vmpooler
       jdata = JSON.parse(request.body.read)
 
       jdata.each do |key, val|
+        if key == 'key'
+          result['key'] = val
+        end
+
         if $redis.scard('vmpooler__ready__' + key) < val.to_i
           available = 0
         end
@@ -330,6 +334,10 @@ module Vmpooler
         result['ok'] = true
 
         jdata.each do |key, val|
+          if key == 'key'
+            next
+          end
+
           result[key] ||= {}
 
           result[key]['ok'] = true ##
@@ -345,6 +353,13 @@ module Vmpooler
 
               result[key]['ok'] = true ##
 
+              if ( result['key'] and $config[:config]['ssh_key'] )
+                Net::SCP.upload!(
+                  vm, 'root', StringIO.new(result['key']), '/root/.ssh/authorized_keys',
+                  :ssh => { :keys => [ $config[:config]['ssh_key'] ] }
+                )
+              end
+
               if result[key]['hostname']
                 result[key]['hostname'] = [result[key]['hostname']] unless result[key]['hostname'].is_a?(Array)
                 result[key]['hostname'].push(vm)
diff --git a/vmpooler.yaml.example b/vmpooler.yaml.example
index 489f014..2375a35 100644
--- a/vmpooler.yaml.example
+++ b/vmpooler.yaml.example
@@ -105,6 +105,11 @@
 #
 #   - domain
 #     If set, returns a top-level 'domain' JSON key in POST requests
+#
+#   - ssh_key
+#     The path to a private key able to log into pooled VMs.  If set, allows a
+#     'key' parameter to be passed in POST requests to '/vm', overriding root's
+#     default authorized_keys file
 
 # Example:
 
@@ -116,6 +121,7 @@
   vm_checktime: 15
   vm_lifetime: 12
   domain: 'company.com'
+  ssh_key: '/root/.ssh/id_rsa'
 
 # :pools:
 #