# Gemfile
gem 'devise'
gem 'active_model_otp'
gem 'rqrcode'

# config/routes.rb
Rails.application.routes.draw do
  devise_for :users, controllers: { sessions: 'users/sessions' }

  resources :users do
    member do
      post :enable_multi_factor_authentication, to: 'users/multi_factor_authentication#verify_enable'
      post :disable_multi_factor_authentication, to: 'users/multi_factor_authentication#verify_disabled'
    end
  end

  get :protected, to: 'visitors#protected'  
  root 'visitors#index'
end

# application.css
.qr {
  border-width: 0;
  border-style: none;
  border-color: #0000ff;
  border-collapse: collapse;
}
.qr td {
  border-width: 0;
  border-style: none;
  border-color: #0000ff;
  border-collapse: collapse;
  padding: 0;
  margin: 0;
  width: 10px;
  height: 10px;
}
.qr td.black { background-color: #000; }
.qr td.white { background-color: #fff; }

# controllers/users/multi_factor_authentication_controller.rb
class Users::MultiFactorAuthenticationController < ApplicationController
  before_action :authenticate_user!
  before_action :set_user

  def verify_enable
    if current_user == @user && 
       current_user.authenticate_otp(params[:multi_factor_authentication][:otp_code_token], drift: 60)
      current_user.otp_module_enabled!
      redirect_to edit_user_registration_path, notice: 'Two Factor Authentication Enabled'
    else
      redirect_to edit_user_registration_path, alert: 'Two Factor Authentication could not be enabled'
    end
  end

  def verify_disabled
    if current_user == @user && 
       current_user.authenticate_otp(params[:multi_factor_authentication][:otp_code_token], drift: 60)
      current_user.otp_module_disabled!
      redirect_to edit_user_registration_path, notice: 'Two Factor Authentication Disabled'
    else
      redirect_to edit_user_registration_path, alert: 'Two Factor Authentication could not be disabled'
    end
  end

  private

  def set_user
    @user = User.find(params[:id])
  end
end

# controllers/users/sessions_controller.rb
class Users::SessionsController < Devise::SessionsController
  def create
    self.resource = warden.authenticate!(auth_options)

    if resource && resource.otp_module_disabled?
      continue_sign_in(resource, resource_name)

    elsif resource && resource.otp_module_enabled?

      if params[:user][:otp_code_token].size > 0
        if resource.authenticate_otp(params[:user][:otp_code_token], drift: 60)
          continue_sign_in(resource, resource_name)
        else
          sign_out resource
          redirect_to root_url, alert: 'Bad Credentials Supplied.'
        end
      else
        sign_out resource
        redirect_to root_url, alert: 'Your account needs to supply a token.'
      end

    end
  end

  private

  def continue_sign_in(resource, resource_name)
    set_flash_message!(:notice, :signed_in)
    sign_in(resource_name, resource)
    yield resource if block_given?
    respond_with resource, location: after_sign_in_path_for(resource)
  end
end

# models/user.rb
class User < ApplicationRecord
  # Include default devise modules. Others available are:
  # :confirmable, :lockable, :timeoutable and :omniauthable
  devise :database_authenticatable, :registerable,
         :recoverable, :rememberable, :trackable, :validatable

  has_one_time_password
  enum otp_module: { disabled: 0, enabled: 1 }, _prefix: true
  attr_accessor :otp_code_token
end

# migration file
# rails g migration add_otp_secret_key_to_users otp_secret_key:string otp_module:integer

class AddOtpSecretKeyToUsers < ActiveRecord::Migration[5.0]
  def change
    add_column :users, :otp_secret_key, :string
    add_column :users, :otp_module, :integer, default: 0
  end
end

# devise/registrations/edit.html.erb
<div class="row">
  <div class="col-sm-4 col-sm-offset-4">
    <h1>Edit <%= resource_name.to_s.humanize %></h1>
    <hr>
    <%= form_for(resource, as: resource_name, url: registration_path(resource_name), html: { method: :put }) do |f| %>
      <%= devise_error_messages! %>

      <div class="form-group">
        <%= f.label :email %><br />
        <%= f.email_field :email, class: 'form-control'  %>
      </div>

      <% if devise_mapping.confirmable? && resource.pending_reconfirmation? %>
        <div>Currently waiting confirmation for: <%= resource.unconfirmed_email %></div>
      <% end %>

      <div class="form-group">
        <%= f.label :password %> <i>(leave blank if you don't want to change it)</i><br />
        <%= f.password_field :password, autocomplete: "off", class: 'form-control'  %>
      </div>

      <div class="form-group">
        <%= f.label :password_confirmation %><br />
        <%= f.password_field :password_confirmation, autocomplete: "off", class: 'form-control'  %>
      </div>

      <div class="form-group">
        <%= f.label :current_password %> <i>(we need your current password to confirm your changes)</i><br />
        <%= f.password_field :current_password, autocomplete: "off", class: 'form-control'  %>
      </div>

      <div class="form-group">
        <%= f.submit "Update", class: 'btn btn-lg btn-block btn-primary' %>
        <%= link_to "#{@user.otp_module_enabled? ? 'Disable' : 'Enable'} Two Factor", 
                    '#two_factor', 
                    data: { toggle: :modal }, 
                    class: 'btn btn-lg btn-block btn-info' %>
      </div>
    <% end %>
  </div>
</div>

<div class="modal fade" id="two_factor">
  <% url = @user.otp_module_enabled? ? disable_multi_factor_authentication_user_path(@user) : enable_multi_factor_authentication_user_path(@user) %>
  <%= simple_form_for :multi_factor_authentication, url: url, html: { class: 'form-inline' }  do |f| %>
    <div class="modal-dialog">
      <div class="modal-content">
        <div class="modal-header">
          <button type="button" class="close" data-dismiss="modal" aria-hidden="true">&times;</button>
          <h4 class="modal-title"><%= @user.otp_module_enabled? ? 'Disable' : 'Enable' %> Two Factor Authentication</h4>
        </div>
        <div class="modal-body">
          <% unless @user.otp_module_enabled? %>
            <% qr = RQRCode::QRCode.new(resource.provisioning_uri, size: 10, level: :h ) %>
            <table class="qr" align="center">
              <% qr.modules.each_index do |x| %>
                  <tr>
                    <% qr.modules.each_index do |y| %>
                        <% if qr.dark?(x,y) %>
                            <td class="black"/>
                        <% else %>
                            <td class="white"/>
                        <% end %>
                    <% end %>
                  </tr>
              <% end %>
            </table>
            <hr>
          <% end %>
          <div class='form-group'>
            <div class='text-center'>
              <%= f.input_field :otp_code_token, placeholder: 'Verify Token', class: 'form-control input-lg' %>
            </div>
          </div>
        </div>
        <div class="modal-footer">
          <%= f.submit "Update", class: 'btn btn-lg btn-block btn-primary' %>
        </div>
      </div>
    </div>
  <% end %>
</div>

# devise/sessions/new.html.erb
<%= f.text_field :otp_code_token, placeholder: 'Token', class: 'form-control' %>