Add option to recover 2FA via SSH

7 files changed, 224 insertions, 12 deletions

diff --git a/CHANGELOG b/CHANGELOG
index 250e822..bce8783 100644
--- a/CHANGELOG
+++ b/CHANGELOG
@@ -1,3 +1,6 @@
+v3.5.0
+  - Add option to recover 2FA via SSH
+
 v3.4.0
   - Redis Sentinel support
 
diff --git a/lib/gitlab_net.rb b/lib/gitlab_net.rb
index 35a8833..47bae95 100644
--- a/lib/gitlab_net.rb
+++ b/lib/gitlab_net.rb
@@ -72,6 +72,15 @@ class GitlabNet
     nil
   end
 
+  def two_factor_recovery_codes(key)
+    key_id = key.gsub('key-', '')
+    resp = post("#{host}/two_factor_recovery_codes", key_id: key_id)
+
+    JSON.parse(resp.body) if resp.code == '200'
+  rescue
+    {}
+  end
+
   def redis_client
     redis_config = config.redis
     database = redis_config['database'] || 0
diff --git a/lib/gitlab_shell.rb b/lib/gitlab_shell.rb
index b6c358e..1fdb9e5 100644
--- a/lib/gitlab_shell.rb
+++ b/lib/gitlab_shell.rb
@@ -8,9 +8,10 @@ class GitlabShell
   class InvalidRepositoryPathError < StandardError; end
 
   GIT_COMMANDS = %w(git-upload-pack git-receive-pack git-upload-archive git-annex-shell git-lfs-authenticate).freeze
+  API_COMMANDS = %w(2fa_recovery_codes)
   GL_PROTOCOL = 'ssh'.freeze
 
-  attr_accessor :key_id, :repo_name, :git_cmd
+  attr_accessor :key_id, :repo_name, :command
   attr_reader :repo_path
 
   def initialize(key_id)
@@ -30,7 +31,7 @@ class GitlabShell
     args = Shellwords.shellwords(origin_cmd)
     parse_cmd(args)
 
-    verify_access
+    verify_access if GIT_COMMANDS.include?(args.first)
 
     process_cmd(args)
 
@@ -58,12 +59,14 @@ class GitlabShell
   protected
 
   def parse_cmd(args)
-    @git_cmd = args.first
-    @git_access = @git_cmd
+    @command = args.first
+    @git_access = @command
 
-    raise DisallowedCommandError unless GIT_COMMANDS.include?(@git_cmd)
+    return if API_COMMANDS.include?(@command)
 
-    case @git_cmd
+    raise DisallowedCommandError unless GIT_COMMANDS.include?(@command)
+
+    case @command
     when 'git-annex-shell'
       raise DisallowedCommandError unless @config.git_annex_enabled?
 
@@ -94,7 +97,9 @@ class GitlabShell
   end
 
   def process_cmd(args)
-    if @git_cmd == 'git-annex-shell'
+    return self.send("api_#{@command}") if API_COMMANDS.include?(@command)
+
+    if @command == 'git-annex-shell'
       raise DisallowedCommandError unless @config.git_annex_enabled?
 
       # Make sure repository has git-annex enabled
@@ -113,8 +118,8 @@ class GitlabShell
       $logger.info "gitlab-shell: executing git-annex command <#{parsed_args.join(' ')}> for #{log_username}."
       exec_cmd(*parsed_args)
     else
-      $logger.info "gitlab-shell: executing git command <#{@git_cmd} #{repo_path}> for #{log_username}."
-      exec_cmd(@git_cmd, repo_path)
+      $logger.info "gitlab-shell: executing git command <#{@command} #{repo_path}> for #{log_username}."
+      exec_cmd(@command, repo_path)
     end
   end
 
@@ -181,6 +186,39 @@ class GitlabShell
 
   private
 
+  def continue?(question)
+    puts "#{question} (yes/no)"
+    STDOUT.flush # Make sure the question gets output before we wait for input
+    continue = STDIN.gets.chomp
+    puts '' # Add a buffer in the output
+    continue == 'yes'
+  end
+
+  def api_2fa_recovery_codes
+    continue = continue?(
+      "Are you sure you want to generate new two-factor recovery codes?\n" \
+      "Any existing recovery codes you saved will be invalidated."
+    )
+
+    unless continue
+      puts 'New recovery codes have *not* been generated. Existing codes will remain valid.'
+      return
+    end
+
+    resp = api.two_factor_recovery_codes(key_id)
+    if resp['success']
+      codes = resp['recovery_codes'].join("\n")
+      puts "Your two-factor authentication recovery codes are:\n\n" \
+           "#{codes}\n\n" \
+           "During sign in, use one of the codes above when prompted for\n" \
+           "your two-factor code. Then, visit your Profile Settings and add\n" \
+           "a new device so you do not lose access to your account again."
+    else
+      puts "An error occurred while trying to generate new recovery codes.\n" \
+           "#{resp['message']}"
+    end
+  end
+
   def repo_path=(repo_path)
     raise ArgumentError, "Repository path not provided. Please make sure you're using GitLab v8.10 or later." unless repo_path
     raise InvalidRepositoryPathError if File.absolute_path(repo_path) != repo_path
diff --git a/spec/gitlab_net_spec.rb b/spec/gitlab_net_spec.rb
index d4585d2..bcd0d79 100644
--- a/spec/gitlab_net_spec.rb
+++ b/spec/gitlab_net_spec.rb
@@ -106,6 +106,24 @@ describe GitlabNet, vcr: true do
     end
   end
 
+  describe '#two_factor_recovery_codes' do
+    it 'returns two factor recovery codes' do
+      VCR.use_cassette('two-factor-recovery-codes') do
+        result = gitlab_net.two_factor_recovery_codes('key-1')
+        expect(result['success']).to be_true
+        expect(result['recovery_codes']).to eq(['f67c514de60c4953','41278385fc00c1e0'])
+      end
+    end
+
+    it 'returns false when recovery codes cannot be generated' do
+      VCR.use_cassette('two-factor-recovery-codes-fail') do
+        result = gitlab_net.two_factor_recovery_codes('key-1')
+        expect(result['success']).to be_false
+        expect(result['message']).to eq('Could not find the given key')
+      end
+    end
+  end
+
   describe :check_access do
     context 'ssh key with access to project' do
       it 'should allow pull access for dev.gitlab.org' do
diff --git a/spec/gitlab_shell_spec.rb b/spec/gitlab_shell_spec.rb
index 0b0a817..ea11652 100644
--- a/spec/gitlab_shell_spec.rb
+++ b/spec/gitlab_shell_spec.rb
@@ -23,6 +23,10 @@ describe GitlabShell do
     double(GitlabNet).tap do |api|
       api.stub(discover: { 'name' => 'John Doe' })
       api.stub(check_access: GitAccessStatus.new(true, 'ok', repo_path))
+      api.stub(two_factor_recovery_codes: {
+        'success' => true,
+        'recovery_codes' => ['f67c514de60c4953', '41278385fc00c1e0']
+      })
     end
   end
 
@@ -53,7 +57,7 @@ describe GitlabShell do
         end
 
         its(:repo_name) { should == 'gitlab-ci.git' }
-        its(:git_cmd) { should == 'git-upload-pack' }
+        its(:command) { should == 'git-upload-pack' }
       end
 
       context 'namespace' do
@@ -65,7 +69,7 @@ describe GitlabShell do
         end
 
         its(:repo_name) { should == 'dmitriy.zaporozhets/gitlab-ci.git' }
-        its(:git_cmd) { should == 'git-upload-pack' }
+        its(:command) { should == 'git-upload-pack' }
       end
 
       context 'with an invalid number of arguments' do
@@ -75,6 +79,24 @@ describe GitlabShell do
           expect { subject.send :parse_cmd, ssh_args }.to raise_error(GitlabShell::DisallowedCommandError)
         end
       end
+
+      context 'with an API command' do
+        before do
+          subject.send :parse_cmd, ssh_args
+        end
+
+        context 'when generating recovery codes' do
+          let(:ssh_args) { %w(2fa_recovery_codes) }
+
+          it 'sets the correct command' do
+            expect(subject.command).to eq('2fa_recovery_codes')
+          end
+
+          it 'does not set repo name' do
+            expect(subject.repo_name).to be_nil
+          end
+        end
+      end
     end
 
     describe 'git-annex' do
@@ -88,7 +110,7 @@ describe GitlabShell do
       end
 
       its(:repo_name) { should == 'dzaporozhets/gitlab.git' }
-      its(:git_cmd) { should == 'git-annex-shell' }
+      its(:command) { should == 'git-annex-shell' }
     end
   end
 
@@ -233,6 +255,44 @@ describe GitlabShell do
         end
       end
     end
+
+    context 'with an API command' do
+      before do
+        allow(subject).to receive(:continue?).and_return(true)
+      end
+
+      context 'when generating recovery codes' do
+        let(:ssh_cmd) { '2fa_recovery_codes' }
+        after do
+          subject.exec(ssh_cmd)
+        end
+
+        it 'does not call verify_access' do
+          expect(subject).not_to receive(:verify_access)
+        end
+
+        it 'calls the corresponding method' do
+          expect(subject).to receive(:api_2fa_recovery_codes)
+        end
+
+        it 'outputs recovery codes' do
+          expect($stdout).to receive(:puts)
+            .with(/f67c514de60c4953\n41278385fc00c1e0/)
+        end
+
+        context 'when the process is unsuccessful' do
+          it 'displays the error to the user' do
+            api.stub(two_factor_recovery_codes: {
+              'success' => false,
+              'message' => 'Could not find the given key'
+            })
+
+            expect($stdout).to receive(:puts)
+              .with(/Could not find the given key/)
+          end
+        end
+      end
+    end
   end
 
   describe :validate_access do
diff --git a/spec/vcr_cassettes/two-factor-recovery-codes-fail.yml b/spec/vcr_cassettes/two-factor-recovery-codes-fail.yml
new file mode 100644
index 0000000..4d5d4c8
--- /dev/null
+++ b/spec/vcr_cassettes/two-factor-recovery-codes-fail.yml
@@ -0,0 +1,42 @@
+---
+http_interactions:
+- request:
+    method: post
+    uri: https://dev.gitlab.org/api/v3/internal/two_factor_recovery_codes
+    body:
+      encoding: US-ASCII
+      string: username=user-1&secret_token=a123
+    headers:
+      Accept-Encoding:
+      - gzip;q=1.0,deflate;q=0.6,identity;q=0.3
+      Accept:
+      - "*/*"
+      User-Agent:
+      - Ruby
+      Content-Type:
+      - application/x-www-form-urlencoded
+  response:
+    status:
+      code: 200
+      message: OK
+    headers:
+      Server:
+      - nginx
+      Date:
+      - Tue, 16 Aug 2016 22:10:11 GMT
+      Content-Type:
+      - application/json
+      Connection:
+      - keep-alive
+      Status:
+      - 200 OK
+      X-Request-Id:
+      - 4467029d-51c6-41bc-af5f-6da279dbb238
+      X-Runtime:
+      - '0.004589'
+    body:
+      encoding: UTF-8
+      string: '{ "success": false, "message": "Could not find the given key" }'
+    http_version:
+  recorded_at: Tue, 16 Aug 2016 22:10:11 GMT
+recorded_with: VCR 2.4.0
diff --git a/spec/vcr_cassettes/two-factor-recovery-codes.yml b/spec/vcr_cassettes/two-factor-recovery-codes.yml
new file mode 100644
index 0000000..2f42166
--- /dev/null
+++ b/spec/vcr_cassettes/two-factor-recovery-codes.yml
@@ -0,0 +1,42 @@
+---
+http_interactions:
+- request:
+    method: post
+    uri: https://dev.gitlab.org/api/v3/internal/two_factor_recovery_codes
+    body:
+      encoding: US-ASCII
+      string: username=user-1&secret_token=a123
+    headers:
+      Accept-Encoding:
+      - gzip;q=1.0,deflate;q=0.6,identity;q=0.3
+      Accept:
+      - "*/*"
+      User-Agent:
+      - Ruby
+      Content-Type:
+      - application/x-www-form-urlencoded
+  response:
+    status:
+      code: 200
+      message: OK
+    headers:
+      Server:
+      - nginx
+      Date:
+      - Tue, 16 Aug 2016 22:10:11 GMT
+      Content-Type:
+      - application/json
+      Connection:
+      - keep-alive
+      Status:
+      - 200 OK
+      X-Request-Id:
+      - 4467029d-51c6-41bc-af5f-6da279dbb238
+      X-Runtime:
+      - '0.004589'
+    body:
+      encoding: UTF-8
+      string: '{ "success": true, "recovery_codes": ["f67c514de60c4953","41278385fc00c1e0"] }'
+    http_version:
+  recorded_at: Tue, 16 Aug 2016 22:10:11 GMT
+recorded_with: VCR 2.4.0