From b4f479a00f00fe540f742435d29b09148ae2e5ec Mon Sep 17 00:00:00 2001
From: Dominik Kwiatek <dom@dom.dog>
Date: Sun, 1 Nov 2020 17:55:31 +0100
Subject: [PATCH] Generate recovery keys on TOTP setup

---
 app/controllers/user_controller.rb            |  7 +++---
 app/models/totp_recovery_code.rb              |  3 +++
 .../settings/security/recovery_keys.haml      |  6 +++++
 ...201101155648_create_totp_recovery_codes.rb |  9 ++++++++
 db/schema.rb                                  |  8 ++++++-
 lib/core_ext/secure_random.rb                 | 23 +++++++++++++++++++
 spec/models/totp_recovery_code_spec.rb        |  5 ++++
 7 files changed, 57 insertions(+), 4 deletions(-)
 create mode 100644 app/models/totp_recovery_code.rb
 create mode 100644 app/views/settings/security/recovery_keys.haml
 create mode 100644 db/migrate/20201101155648_create_totp_recovery_codes.rb
 create mode 100644 lib/core_ext/secure_random.rb
 create mode 100644 spec/models/totp_recovery_code_spec.rb

diff --git a/app/controllers/user_controller.rb b/app/controllers/user_controller.rb
index b2fc4b7c..00d3903e 100644
--- a/app/controllers/user_controller.rb
+++ b/app/controllers/user_controller.rb
@@ -190,13 +190,14 @@ class UserController < ApplicationController
     current_user.otp_module = :enabled
 
     if current_user.authenticate_otp(req_params[:otp_validation], drift: APP_CONFIG.fetch(:otp_drift_period, 30).to_i)
-      flash[:success] = t('views.auth.2fa.setup.success')
+      @recovery_keys = TotpRecoveryCode.create!(Array.new(10) { {user: current_user, code: SecureRandom.base58(8).downcase} })
       current_user.save!
+
+      render 'settings/security/recovery_keys'
     else
       flash[:error] = t('views.auth.2fa.errors.invalid_code')
+      redirect_to edit_user_security_path
     end
-
-    redirect_to edit_user_security_path
   end
 
   def destroy_2fa
diff --git a/app/models/totp_recovery_code.rb b/app/models/totp_recovery_code.rb
new file mode 100644
index 00000000..a83df0e4
--- /dev/null
+++ b/app/models/totp_recovery_code.rb
@@ -0,0 +1,3 @@
+class TotpRecoveryCode < ApplicationRecord
+  belongs_to :user
+end
diff --git a/app/views/settings/security/recovery_keys.haml b/app/views/settings/security/recovery_keys.haml
new file mode 100644
index 00000000..066facf0
--- /dev/null
+++ b/app/views/settings/security/recovery_keys.haml
@@ -0,0 +1,6 @@
+%p= t('views.auth.2fa.setup.success')
+
+%ul
+  - @recovery_keys.each do |key|
+    %li
+      %code= key.code
diff --git a/db/migrate/20201101155648_create_totp_recovery_codes.rb b/db/migrate/20201101155648_create_totp_recovery_codes.rb
new file mode 100644
index 00000000..3b506256
--- /dev/null
+++ b/db/migrate/20201101155648_create_totp_recovery_codes.rb
@@ -0,0 +1,9 @@
+class CreateTotpRecoveryCodes < ActiveRecord::Migration[5.2]
+  def change
+    create_table :totp_recovery_codes do |t|
+      t.bigint :user_id
+      t.string :code, limit: 8
+    end
+    add_index :totp_recovery_codes, [:user_id, :code]
+  end
+end
diff --git a/db/schema.rb b/db/schema.rb
index 4ad56134..a297997d 100644
--- a/db/schema.rb
+++ b/db/schema.rb
@@ -10,7 +10,7 @@
 #
 # It's strongly recommended that you check this file into your version control system.
 
-ActiveRecord::Schema.define(version: 2020_10_18_090453) do
+ActiveRecord::Schema.define(version: 2020_11_01_155648) do
 
   # These are extensions that must be enabled in order to support this database
   enable_extension "plpgsql"
@@ -219,6 +219,12 @@ ActiveRecord::Schema.define(version: 2020_10_18_090453) do
     t.index ["user_id", "created_at"], name: "index_themes_on_user_id_and_created_at"
   end
 
+  create_table "totp_recovery_codes", force: :cascade do |t|
+    t.bigint "user_id"
+    t.string "code", limit: 8
+    t.index ["user_id", "code"], name: "index_totp_recovery_codes_on_user_id_and_code"
+  end
+
   create_table "users", id: :bigint, default: -> { "gen_timestamp_id('users'::text)" }, force: :cascade do |t|
     t.string "email", default: "", null: false
     t.string "encrypted_password", default: "", null: false
diff --git a/lib/core_ext/secure_random.rb b/lib/core_ext/secure_random.rb
new file mode 100644
index 00000000..5cc3c363
--- /dev/null
+++ b/lib/core_ext/secure_random.rb
@@ -0,0 +1,23 @@
+# From Rails 6
+# Remove once upgraded
+module SecureRandom
+  BASE36_ALPHABET = ("0".."9").to_a + ("a".."z").to_a
+  # SecureRandom.base36 generates a random base36 string in lowercase.
+  #
+  # The argument _n_ specifies the length of the random string to be generated.
+  #
+  # If _n_ is not specified or is +nil+, 16 is assumed. It may be larger in the future.
+  # This method can be used over +base58+ if a deterministic case key is necessary.
+  #
+  # The result will contain alphanumeric characters in lowercase.
+  #
+  #   p SecureRandom.base36 # => "4kugl2pdqmscqtje"
+  #   p SecureRandom.base36(24) # => "77tmhrhjfvfdwodq8w7ev2m7"
+  def self.base36(n = 16)
+    SecureRandom.random_bytes(n).unpack("C*").map do |byte|
+      idx = byte % 64
+      idx = SecureRandom.random_number(36) if idx >= 36
+      BASE36_ALPHABET[idx]
+    end.join
+  end
+end
\ No newline at end of file
diff --git a/spec/models/totp_recovery_code_spec.rb b/spec/models/totp_recovery_code_spec.rb
new file mode 100644
index 00000000..66f2040a
--- /dev/null
+++ b/spec/models/totp_recovery_code_spec.rb
@@ -0,0 +1,5 @@
+require 'rails_helper'
+
+RSpec.describe TotpRecoveryCode, type: :model do
+  pending "add some examples to (or delete) #{__FILE__}"
+end