diff --git a/lib/heroku/auth.rb b/lib/heroku/auth.rb index 5ec04b2d9..5387a5d22 100644 --- a/lib/heroku/auth.rb +++ b/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 diff --git a/spec/heroku/auth_spec.rb b/spec/heroku/auth_spec.rb index 9da71e573..5e50b1614 100644 --- a/spec/heroku/auth_spec.rb +++ b/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 diff --git a/spec/heroku/command/keys_spec.rb b/spec/heroku/command/keys_spec.rb index d479f17ab..bb4d5d210 100644 --- a/spec/heroku/command/keys_spec.rb +++ b/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