Provide go implementation for 2fa_recovery_codes command

4 files changed, 203 insertions, 94 deletions

diff --git a/spec/gitlab_shell_authorized_keys_check_spec.rb b/spec/gitlab_shell_authorized_keys_check_spec.rb
index baaa560..7050604 100644
--- a/spec/gitlab_shell_authorized_keys_check_spec.rb
+++ b/spec/gitlab_shell_authorized_keys_check_spec.rb
@@ -1,20 +1,7 @@
 require_relative 'spec_helper'
 
 describe 'bin/gitlab-shell-authorized-keys-check' do
-  def original_root_path
-    ROOT_PATH
-  end
-
-  # All this test boilerplate is mostly copy/pasted between
-  # gitlab_shell_gitlab_shell_spec.rb and
-  # gitlab_shell_authorized_keys_check_spec.rb
-  def tmp_root_path
-    @tmp_root_path ||= File.realpath(Dir.mktmpdir)
-  end
-
-  def config_path
-    File.join(tmp_root_path, 'config.yml')
-  end
+  include_context 'gitlab shell'
 
   def tmp_socket_path
     # This has to be a relative path shorter than 100 bytes due to
@@ -22,12 +9,8 @@ describe 'bin/gitlab-shell-authorized-keys-check' do
     'tmp/gitlab-shell-authorized-keys-check-socket'
   end
 
-  before(:all) do
-    FileUtils.mkdir_p(File.dirname(tmp_socket_path))
-    FileUtils.touch(File.join(tmp_root_path, '.gitlab_shell_secret'))
-
-    @server = HTTPUNIXServer.new(BindAddress: tmp_socket_path)
-    @server.mount_proc('/api/v4/internal/authorized_keys') do |req, res|
+  def mock_server(server)
+    server.mount_proc('/api/v4/internal/authorized_keys') do |req, res|
       if req.query['key'] == 'known-rsa-key'
         res.status = 200
         res.content_type = 'application/json'
@@ -36,28 +19,14 @@ describe 'bin/gitlab-shell-authorized-keys-check' do
         res.status = 404
       end
     end
-
-    @webrick_thread = Thread.new { @server.start }
-
-    sleep(0.1) while @webrick_thread.alive? && @server.status != :Running
-    raise "Couldn't start stub GitlabNet server" unless @server.status == :Running
-
-    File.open(config_path, 'w') do |f|
-      f.write("---\ngitlab_url: http+unix://#{CGI.escape(tmp_socket_path)}\n")
-    end
-
-    copy_dirs = ['bin', 'lib']
-    FileUtils.rm_rf(copy_dirs.map { |d| File.join(tmp_root_path, d) })
-    FileUtils.cp_r(copy_dirs, tmp_root_path)
   end
 
-  after(:all) do
-    @server.shutdown if @server
-    @webrick_thread.join if @webrick_thread
-    FileUtils.rm_rf(tmp_root_path)
+  before(:all) do
+    write_config(
+      "gitlab_url" => "http+unix://#{CGI.escape(tmp_socket_path)}",
+    )
   end
 
-  let(:gitlab_shell_path) { File.join(tmp_root_path, 'bin', 'gitlab-shell') }
   let(:authorized_keys_check_path) { File.join(tmp_root_path, 'bin', 'gitlab-shell-authorized-keys-check') }
 
   it 'succeeds when a valid key is given' do
diff --git a/spec/gitlab_shell_gitlab_shell_spec.rb b/spec/gitlab_shell_gitlab_shell_spec.rb
index cb3fd9c..6d6e172 100644
--- a/spec/gitlab_shell_gitlab_shell_spec.rb
+++ b/spec/gitlab_shell_gitlab_shell_spec.rb
@@ -3,33 +3,10 @@ require_relative 'spec_helper'
 require 'open3'
 
 describe 'bin/gitlab-shell' do
-  def original_root_path
-    ROOT_PATH
-  end
-
-  # All this test boilerplate is mostly copy/pasted between
-  # gitlab_shell_gitlab_shell_spec.rb and
-  # gitlab_shell_authorized_keys_check_spec.rb
-  def tmp_root_path
-    @tmp_root_path ||= File.realpath(Dir.mktmpdir)
-  end
-
-  def config_path
-    File.join(tmp_root_path, 'config.yml')
-  end
-
-  def tmp_socket_path
-    # This has to be a relative path shorter than 100 bytes due to
-    # limitations in how Unix sockets work.
-    'tmp/gitlab-shell-socket'
-  end
-
-  before(:all) do
-    FileUtils.mkdir_p(File.dirname(tmp_socket_path))
-    FileUtils.touch(File.join(tmp_root_path, '.gitlab_shell_secret'))
+  include_context 'gitlab shell'
 
-    @server = HTTPUNIXServer.new(BindAddress: tmp_socket_path)
-    @server.mount_proc('/api/v4/internal/discover') do |req, res|
+  def mock_server(server)
+    server.mount_proc('/api/v4/internal/discover') do |req, res|
       identifier = req.query['key_id'] || req.query['username'] || req.query['user_id']
       known_identifiers = %w(10 someuser 100)
       if known_identifiers.include?(identifier)
@@ -47,24 +24,16 @@ describe 'bin/gitlab-shell' do
         res.status = 500
       end
     end
-
-    @webrick_thread = Thread.new { @server.start }
-
-    sleep(0.1) while @webrick_thread.alive? && @server.status != :Running
-    raise "Couldn't start stub GitlabNet server" unless @server.status == :Running
-    system(original_root_path, 'bin/compile')
-    copy_dirs = ['bin', 'lib']
-    FileUtils.rm_rf(copy_dirs.map { |d| File.join(tmp_root_path, d) })
-    FileUtils.cp_r(copy_dirs, tmp_root_path)
   end
 
-  after(:all) do
-    @server.shutdown if @server
-    @webrick_thread.join if @webrick_thread
-    FileUtils.rm_rf(tmp_root_path)
-  end
+  def run!(args, env: {'SSH_CONNECTION' => 'fake'})
+    cmd = [
+      gitlab_shell_path,
+      args
+    ].flatten.compact.join(' ')
 
-  let(:gitlab_shell_path) { File.join(tmp_root_path, 'bin', 'gitlab-shell') }
+    Open3.capture3(env, cmd)
+  end
 
   shared_examples 'results with keys' do
     # Basic valid input
@@ -175,19 +144,4 @@ describe 'bin/gitlab-shell' do
       expect(status).not_to be_success
     end
   end
-
-  def run!(args, env: {'SSH_CONNECTION' => 'fake'})
-    cmd = [
-      gitlab_shell_path,
-      args
-    ].flatten.compact.join(' ')
-
-    Open3.capture3(env, cmd)
-  end
-
-  def write_config(config)
-    File.open(config_path, 'w') do |f|
-      f.write(config.to_yaml)
-    end
-  end
 end
diff --git a/spec/gitlab_shell_two_factor_recovery_spec.rb b/spec/gitlab_shell_two_factor_recovery_spec.rb
new file mode 100644
index 0000000..19999e5
--- /dev/null
+++ b/spec/gitlab_shell_two_factor_recovery_spec.rb
@@ -0,0 +1,128 @@
+require_relative 'spec_helper'
+
+require 'open3'
+
+describe 'bin/gitlab-shell 2fa_recovery_codes' do
+  include_context 'gitlab shell'
+
+  def mock_server(server)
+    server.mount_proc('/api/v4/internal/two_factor_recovery_codes') do |req, res|
+      res.content_type = 'application/json'
+      res.status = 200
+
+      key_id = req.query['key_id'] || req.query['user_id']
+
+      unless key_id
+        body = JSON.parse(req.body)
+        key_id = body['key_id'] || body['user_id'].to_s
+      end
+
+      if key_id == '100'
+        res.body = '{"success":true, "recovery_codes": ["1", "2"]}'
+      else
+        res.body = '{"success":false, "message": "Forbidden!"}'
+      end
+    end
+
+    server.mount_proc('/api/v4/internal/discover') do |req, res|
+      res.status = 200
+      res.content_type = 'application/json'
+      res.body = '{"id":100, "name": "Some User", "username": "someuser"}'
+    end
+  end
+
+  shared_examples 'dialog for regenerating recovery keys' do
+    context 'when the user agrees to regenerate keys' do
+      def verify_successful_regeneration!(cmd)
+        Open3.popen2(env, cmd) do |stdin, stdout|
+          expect(stdout.gets).to eq("Are you sure you want to generate new two-factor recovery codes?\n")
+          expect(stdout.gets).to eq("Any existing recovery codes you saved will be invalidated. (yes/no)\n")
+
+          stdin.puts('yes')
+
+          expect(stdout.flush.read).to eq(
+            "\nYour two-factor authentication recovery codes are:\n\n" \
+            "1\n2\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.\n"
+          )
+        end
+      end
+
+      context 'when key is provided' do
+        let(:cmd) { "#{gitlab_shell_path} key-100" }
+
+        it 'the recovery keys are regenerated' do
+          verify_successful_regeneration!(cmd)
+        end
+      end
+
+      context 'when username is provided' do
+        let(:cmd) { "#{gitlab_shell_path} username-someone" }
+
+        it 'the recovery keys are regenerated' do
+          verify_successful_regeneration!(cmd)
+        end
+      end
+    end
+
+    context 'when the user disagrees to regenerate keys' do
+      let(:cmd) { "#{gitlab_shell_path} key-100" }
+
+      it 'the recovery keys are not regenerated' do
+        Open3.popen2(env, cmd) do |stdin, stdout|
+          expect(stdout.gets).to eq("Are you sure you want to generate new two-factor recovery codes?\n")
+          expect(stdout.gets).to eq("Any existing recovery codes you saved will be invalidated. (yes/no)\n")
+
+          stdin.puts('no')
+
+          expect(stdout.flush.read).to eq(
+            "\nNew recovery codes have *not* been generated. Existing codes will remain valid.\n"
+          )
+        end
+      end
+    end
+
+    context 'when API error occurs' do
+      let(:cmd) { "#{gitlab_shell_path} key-101" }
+
+      context 'when the user agrees to regenerate keys' do
+        it 'the recovery keys are regenerated' do
+          Open3.popen2(env, cmd) do |stdin, stdout|
+            expect(stdout.gets).to eq("Are you sure you want to generate new two-factor recovery codes?\n")
+            expect(stdout.gets).to eq("Any existing recovery codes you saved will be invalidated. (yes/no)\n")
+
+            stdin.puts('yes')
+
+            expect(stdout.flush.read).to eq("\nAn error occurred while trying to generate new recovery codes.\nForbidden!\n")
+          end
+        end
+      end
+    end
+  end
+
+  let(:env) { {'SSH_CONNECTION' => 'fake', 'SSH_ORIGINAL_COMMAND' => '2fa_recovery_codes' } }
+
+  describe 'without go features' do
+    before(:context) do
+      write_config(
+        "gitlab_url" => "http+unix://#{CGI.escape(tmp_socket_path)}",
+      )
+    end
+
+    it_behaves_like 'dialog for regenerating recovery keys'
+  end
+
+  describe 'with go features' do
+    before(:context) do
+      write_config(
+        "gitlab_url" => "http+unix://#{CGI.escape(tmp_socket_path)}",
+        "migration" => { "enabled" => true,
+                        "features" => ["2fa_recovery_codes"] }
+      )
+    end
+
+    it_behaves_like 'dialog for regenerating recovery keys'
+  end
+end
diff --git a/spec/support/gitlab_shell_setup.rb b/spec/support/gitlab_shell_setup.rb
new file mode 100644
index 0000000..eddd2d1
--- /dev/null
+++ b/spec/support/gitlab_shell_setup.rb
@@ -0,0 +1,58 @@
+RSpec.shared_context 'gitlab shell', shared_context: :metadata do
+  def original_root_path
+    ROOT_PATH
+  end
+
+  def config_path
+    File.join(tmp_root_path, 'config.yml')
+  end
+
+  def write_config(config)
+    File.open(config_path, 'w') do |f|
+      f.write(config.to_yaml)
+    end
+  end
+
+  def tmp_root_path
+    @tmp_root_path ||= File.realpath(Dir.mktmpdir)
+  end
+
+  def mock_server(server)
+    raise NotImplementedError.new(
+      'mock_server method must be implemented in order to include gitlab shell context'
+    )
+  end
+
+  # This has to be a relative path shorter than 100 bytes due to
+  # limitations in how Unix sockets work.
+  def tmp_socket_path
+    'tmp/gitlab-shell-socket'
+  end
+
+  let(:gitlab_shell_path) { File.join(tmp_root_path, 'bin', 'gitlab-shell') }
+
+  before(:all) do
+    FileUtils.mkdir_p(File.dirname(tmp_socket_path))
+    FileUtils.touch(File.join(tmp_root_path, '.gitlab_shell_secret'))
+
+    @server = HTTPUNIXServer.new(BindAddress: tmp_socket_path)
+
+    mock_server(@server)
+
+    @webrick_thread = Thread.new { @server.start }
+
+    sleep(0.1) while @webrick_thread.alive? && @server.status != :Running
+    raise "Couldn't start stub GitlabNet server" unless @server.status == :Running
+    system(original_root_path, 'bin/compile')
+
+    copy_dirs = ['bin', 'lib']
+    FileUtils.rm_rf(copy_dirs.map { |d| File.join(tmp_root_path, d) })
+    FileUtils.cp_r(copy_dirs, tmp_root_path)
+  end
+
+  after(:all) do
+    @server.shutdown if @server
+    @webrick_thread.join if @webrick_thread
+    FileUtils.rm_rf(tmp_root_path)
+  end
+end