prompt for creating ~/.ssh/id_rsa

lib/heroku/auth.rb

@@ -255,25 +255,37 @@ def ask_for_and_save_credentials
 
     def check_for_associated_ssh_key
       if api.get_keys.body.empty?
+        display "Your Heroku account does not have a public ssh key uploaded."
         associate_or_generate_ssh_key
       end
     end
 
     def associate_or_generate_ssh_key
-      public_keys = Dir.glob("#{home_directory}/.ssh/*.pub").sort
-
-      case public_keys.length
-      when 0 then
-        display "Could not find an existing public key."
+      unless File.exists?("#{home_directory}/.ssh/id_rsa.pub")
+        display "Could not find an existing public key at ~/.ssh/id_rsa.pub"
         display "Would you like to generate one? [Yn] ", false
-        unless ask.strip.downcase == "n"
+        unless ask.strip.downcase =~ /^n/
           display "Generating new SSH public key."
-          generate_ssh_key("id_rsa")
+          generate_ssh_key("#{home_directory}/.ssh/id_rsa")
           associate_key("#{home_directory}/.ssh/id_rsa.pub")
+          return
         end
-      when 1 then
-        display "Found existing public key: #{public_keys.first}"
-        associate_key(public_keys.first)
+      end
+
+      chosen = ssh_prompt
+      associate_key(chosen) if chosen
+    end
+
+    def ssh_prompt
+      public_keys = Dir.glob("#{home_directory}/.ssh/*.pub").sort
+      case public_keys.length
+      when 0
+        error("No SSH keys found")
+        return nil
+      when 1
+        display "Found an SSH public key at #{public_keys.first}"
+        display "Would you like to upload it to Heroku? [Yn] ", false
+        return ask.strip.downcase =~ /^n/ ? nil : public_keys.first
       else
         display "Found the following SSH public keys:"
         public_keys.each_with_index do |key, index|
@@ -285,19 +297,14 @@ def associate_or_generate_ssh_key
         if choice == -1 || chosen.nil?
           error("Invalid choice")
         end
-        associate_key(chosen)
+        return chosen
       end
     end
 
     def generate_ssh_key(keyfile)
-      ssh_dir = File.join(home_directory, ".ssh")
-      unless File.exists?(ssh_dir)
-        FileUtils.mkdir_p ssh_dir
-        unless running_on_windows?
-          File.chmod(0700, ssh_dir)
-        end
-      end
-      output = `ssh-keygen -t rsa -N "" -f \"#{home_directory}/.ssh/#{keyfile}\" 2>&1`
+      ssh_dir = File.dirname(keyfile)
+      FileUtils.mkdir_p ssh_dir, :mode => 0700
+      output = `ssh-keygen -t rsa -N "" -f \"#{keyfile}\" 2>&1`
       if ! $?.success?
         error("Could not generate key: #{output}")
       end

spec/heroku/auth_spec.rb

@@ -185,6 +185,7 @@ module Heroku
 
     describe "automatic key uploading" do
       before(:each) do
+        allow(@cli).to receive(:home_directory).and_return(Heroku::Helpers.home_directory)
         FileUtils.mkdir_p("#{@cli.home_directory}/.ssh")
         allow(@cli).to receive(:ask_for_credentials).and_return("username", "apikey")
       end
@@ -216,31 +217,20 @@ module Heroku
         describe "with zero public keys" do
           it "should ask to generate a key" do
             expect(@cli).to receive(:ask).and_return("y")
-            expect(@cli).to receive(:generate_ssh_key).with("id_rsa")
-            expect(@cli).to receive(:associate_key).with("#{@cli.home_directory}/.ssh/id_rsa.pub")
-            @cli.check_for_associated_ssh_key
-          end
-        end
-
-        describe "with one public key" do
-          before(:each) { FileUtils.touch("#{@cli.home_directory}/.ssh/id_rsa.pub") }
-          after(:each)  { FileUtils.rm("#{@cli.home_directory}/.ssh/id_rsa.pub") }
-
-          it "should upload the key" do
+            expect(@cli).to receive(:generate_ssh_key).with("#{@cli.home_directory}/.ssh/id_rsa")
             expect(@cli).to receive(:associate_key).with("#{@cli.home_directory}/.ssh/id_rsa.pub")
             @cli.check_for_associated_ssh_key
           end
         end
 
         describe "with many public keys" do
-          before(:each) do
+          before :each do
             FileUtils.touch("#{@cli.home_directory}/.ssh/id_rsa.pub")
             FileUtils.touch("#{@cli.home_directory}/.ssh/id_rsa2.pub")
           end
 
-          after(:each) do
-            FileUtils.rm("#{@cli.home_directory}/.ssh/id_rsa.pub")
-            FileUtils.rm("#{@cli.home_directory}/.ssh/id_rsa2.pub")
+          after :each do
+            FileUtils.rm_rf(@cli.home_directory)
           end
 
           it "should ask which key to upload" do

spec/heroku/command/keys_spec.rb

@@ -7,27 +7,20 @@ module Heroku::Command
 
     before(:each) do
       stub_core
+      allow(Heroku::Auth).to receive(:home_directory).and_return(Heroku::Helpers.home_directory)
     end
 
     context("add") do
-
-      after(:each) do
-        api.delete_key("pedro@heroku")
-      end
-
       it "tries to find a key if no key filename is supplied" do
         expect(Heroku::Auth).to receive(:ask).and_return("y")
-        expect(Heroku::Auth).to receive(:generate_ssh_key)
-        expect(File).to receive(:exists?).with('.git').and_return(false)
-        expect(File).to receive(:exists?).with('/.ssh/id_rsa.pub').and_return(true)
-        expect(File).to receive(:read).with('/.ssh/id_rsa.pub').and_return(KEY)
         stderr, stdout = execute("keys:add")
         expect(stderr).to eq("")
         expect(stdout).to eq <<-STDOUT
-Could not find an existing public key.
+Could not find an existing public key at ~/.ssh/id_rsa.pub
 Would you like to generate one? [Yn] Generating new SSH public key.
-Uploading SSH public key /.ssh/id_rsa.pub... done
+Uploading SSH public key #{Heroku::Auth.home_directory}/.ssh/id_rsa.pub... done
 STDOUT
+        api.delete_key(`whoami`.strip + '@' + `hostname`.strip)
       end
 
       it "adds a key from a specified keyfile path" do
@@ -39,8 +32,8 @@ module Heroku::Command
         expect(stdout).to eq <<-STDOUT
 Uploading SSH public key /my/key.pub... done
 STDOUT
+        api.delete_key("pedro@heroku")
       end
-
     end
 
     context("index") do