use 2fa directly instead of auto-preauth strategy

Gemfile.lock

@@ -30,7 +30,7 @@ GEM
     docile (1.1.5)
     excon (0.42.1)
     fakefs (0.5.4)
-    heroku-api (0.3.21)
+    heroku-api (0.3.22)
       excon (~> 0.38)
       multi_json (~> 1.8)
     json (1.8.1)

lib/heroku/auth.rb

@@ -82,16 +82,14 @@ def password    # :nodoc:
       get_credentials[1]
     end
 
-    def api_key(user=get_credentials[0], password=get_credentials[1], second_factor=nil)
-      params = default_params
-      if second_factor
-        params[:headers].merge!("Heroku-Two-Factor-Code" => second_factor)
-      end
-      api = Heroku::API.new(params)
-      api.post_login(user, password).body["api_key"]
+    def api_key(user=get_credentials[0], password=get_credentials[1])
+      @api ||= Heroku::API.new(default_params)
+      api_key = @api.post_login(user, password).body["api_key"]
+      @api = nil
+      api_key
     rescue Heroku::API::Errors::Forbidden => e
       if e.response.headers.has_key?("Heroku-Two-Factor-Required")
-        second_factor = ask_for_second_factor
+        ask_for_second_factor
         retry
       end
     rescue Heroku::API::Errors::Unauthorized => e
@@ -218,7 +216,7 @@ def ask_for_credentials
 
     def ask_for_second_factor
       $stderr.print "Two-factor code: "
-      ask
+      api.second_factor = ask
     end
 
     def preauth

lib/heroku/command.rb

@@ -9,6 +9,10 @@ class CommandFailed  < RuntimeError; end
 
     extend Heroku::Helpers
 
+    class << self
+      attr_accessor :requires_preauth
+    end
+
     def self.load
       Dir[File.join(File.dirname(__FILE__), "command", "*.rb")].each do |file|
         require file
@@ -244,7 +248,11 @@ def self.run(cmd, arguments=[])
       error "API request timed out. Please try again, or contact support@heroku.com if this issue persists."
     rescue Heroku::API::Errors::Forbidden => e
       if e.response.headers.has_key?("Heroku-Two-Factor-Required")
-        Heroku::Auth.preauth
+        if requires_preauth
+          Heroku::Auth.preauth
+        else
+          Heroku::Auth.ask_for_second_factor
+        end
         retry
       else
         error extract_error(e.response.body)

lib/heroku/command/base.rb

@@ -294,6 +294,10 @@ def git_remotes(base_dir=Dir.pwd)
   def escape(value)
     heroku.escape(value)
   end
+
+  def requires_preauth
+    Heroku::Command.requires_preauth = true
+  end
 end
 
 module Heroku::Command

lib/heroku/command/config.rb

@@ -65,6 +65,7 @@ def index
   # B: two
   #
   def set
+    requires_preauth
     unless args.size > 0 and args.all? { |a| a.include?('=') }
       error("Usage: heroku config:set KEY1=VALUE1 [KEY2=VALUE2 ...]\nMust specify KEY and VALUE to set.")
     end

lib/heroku/command/pg.rb

@@ -25,6 +25,7 @@ def set_commands(shorthand)
   # list databases for an app
   #
   def index
+    requires_preauth
     validate_arguments!
 
     if hpg_databases_with_info.empty?
@@ -47,6 +48,7 @@ def index
   def info
     db = shift_argument
     validate_arguments!
+    requires_preauth
 
     if db
       @resolver = generate_resolver
@@ -64,6 +66,7 @@ def info
   # defaults to DATABASE_URL databases if no DATABASE is specified
   # if REPORT_ID is specified instead, a previous report is displayed
   def diagnose
+    requires_preauth
     db_id = shift_argument
     run_diagnose(db_id)
   end
@@ -73,6 +76,7 @@ def diagnose
   # sets DATABASE as your DATABASE_URL
   #
   def promote
+    requires_preauth
     unless db = shift_argument
       error("Usage: heroku pg:promote DATABASE\nMust specify DATABASE to promote.")
     end
@@ -94,6 +98,7 @@ def promote
   # defaults to DATABASE_URL databases if no DATABASE is specified
   #
   def psql
+    requires_preauth
     attachment = generate_resolver.resolve(shift_argument, "DATABASE_URL")
     validate_arguments!
 
@@ -123,6 +128,7 @@ def psql
   # delete all data in DATABASE
   #
   def reset
+    requires_preauth
     unless db = shift_argument
       error("Usage: heroku pg:reset DATABASE\nMust specify DATABASE to reset.")
     end
@@ -144,6 +150,7 @@ def reset
   # stop a replica from following and make it a read/write database
   #
   def unfollow
+    requires_preauth
     unless db = shift_argument
       error("Usage: heroku pg:unfollow REPLICA\nMust specify REPLICA to unfollow.")
     end
@@ -177,6 +184,7 @@ def unfollow
   # defaults to all databases if no DATABASE is specified
   #
   def wait
+    requires_preauth
     db = shift_argument
     validate_arguments!
 
@@ -196,6 +204,7 @@ def wait
   #   --reset  # Reset credentials on the specified database.
   #
   def credentials
+    requires_preauth
     unless db = shift_argument
       error("Usage: heroku pg:credentials DATABASE\nMust specify DATABASE to display credentials.")
     end
@@ -228,6 +237,7 @@ def credentials
   # view active queries with execution time
   #
   def ps
+    requires_preauth
     sql = %Q(
     SELECT
       #{pid_column},
@@ -260,6 +270,7 @@ def ps
   # -f,--force  # terminates the connection in addition to cancelling the query
   #
   def kill
+    requires_preauth
     procpid = shift_argument
     output_with_bang "procpid to kill is required" unless procpid && procpid.to_i != 0
     procpid = procpid.to_i
@@ -275,6 +286,7 @@ def kill
   # terminates ALL connections
   #
   def killall
+    requires_preauth
     db = args.first
     attachment = generate_resolver.resolve(db, "DATABASE_URL")
     client = hpg_client(attachment)
@@ -299,6 +311,7 @@ def killall
   # push from LOCAL_SOURCE_DATABASE to REMOTE_TARGET_DATABASE
   # REMOTE_TARGET_DATABASE must be empty.
   def push
+    requires_preauth
     local, remote = shift_argument, shift_argument
     unless [remote, local].all?
       Heroku::Command.run(current_command, ['--help'])
@@ -324,6 +337,7 @@ def push
   # pull from REMOTE_SOURCE_DATABASE to LOCAL_TARGET_DATABASE
   # LOCAL_TARGET_DATABASE must not already exist.
   def pull
+    requires_preauth
     remote, local = shift_argument, shift_argument
     unless [remote, local].all?
       Heroku::Command.run(current_command, ['--help'])
@@ -354,6 +368,7 @@ def pull
   #  window="<window>"  # set weekly UTC maintenance window for DATABASE
   #                     # eg: `heroku pg:maintenance window="Sunday 14:30"`
   def maintenance
+    requires_preauth
     mode_with_argument = shift_argument || ''
     mode, mode_argument = mode_with_argument.split('=')
 
@@ -397,6 +412,7 @@ def maintenance
   # unfollow a database and upgrade it to the latest PostgreSQL version
   #
   def upgrade
+    requires_preauth
     unless db = shift_argument
       error("Usage: heroku pg:upgrade REPLICA\nMust specify REPLICA to upgrade.")
     end