From 4c4f6370755bac7877893c5c96ad2eba5a881321 Mon Sep 17 00:00:00 2001
From: Scott Schneider <sschneider@puppetlabs.com>
Date: Fri, 14 Nov 2014 12:50:19 -0800
Subject: [PATCH] Allow authorized_keys to be overridden at VM checkout time

---
 lib/vmpooler.rb       |  2 ++
 lib/vmpooler/api.rb   | 15 +++++++++++++++
 vmpooler.yaml.example |  6 ++++++
 3 files changed, 23 insertions(+)

diff --git a/lib/vmpooler.rb b/lib/vmpooler.rb
index ef820da..5f0aa2c 100644
--- a/lib/vmpooler.rb
+++ b/lib/vmpooler.rb
@@ -2,10 +2,12 @@ require 'rubygems' unless defined?(Gem)
 
 module Vmpooler
   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.rb b/lib/vmpooler/api.rb
index 40c9156..4e6d1b9 100644
--- a/lib/vmpooler/api.rb
+++ b/lib/vmpooler/api.rb
@@ -190,6 +190,10 @@ module Vmpooler
           jdata = JSON.parse(request.body.read)
 
           jdata.each do |template, count|
+            if (template == 'key')
+              result['key'] = count
+            end
+
             if ( $redis.scard('vmpooler__ready__'+template) < count.to_i )
               available = 0
             end
@@ -199,6 +203,10 @@ module Vmpooler
             result['ok'] = true
 
             jdata.each do |template, count|
+              if (template == 'key')
+                next
+              end
+
               result[template] ||= {}
 
               result[template]['ok'] = true ##
@@ -214,6 +222,13 @@ module Vmpooler
 
                   result[template]['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[template]['hostname'] )
                     result[template]['hostname'] = [result[template]['hostname']] if ! result[template]['hostname'].is_a?(Array)
                     result[template]['hostname'].push(vm)
diff --git a/vmpooler.yaml.example b/vmpooler.yaml.example
index a8ffc8c..87b2ddc 100644
--- a/vmpooler.yaml.example
+++ b/vmpooler.yaml.example
@@ -101,6 +101,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:
 
@@ -111,6 +116,7 @@
   vm_checktime: 15
   vm_lifetime: 12
   domain: 'company.com'
+  ssh_key: '/root/.ssh/id_rsa'
 
 # :pools:
 #