#!/usr/bin/env ruby
require 'rubygems'
require 'rotp'

# This is a script to use two factor authentication with public ssh keys (since you can't use PAM (g2fa) and public keys)
# You must refrence this script in your sshd config: ForceCommand="/usr/bin/two_factor_ssh"
# The script has to be executible by the user logging in

# get the username of the user logging in
user = ENV["USER"]
file = "/etc/google-authenticator/#{user}/auth_key"

# Ensure the user is configured with an auth_key
abort "User not configured for two-factor authentication" unless File.exist?(file)

# read in the users pre-shared key and ensure there are no hanging chars
# TODO check to see fi the file exists and the directory exists else abort with meaningful txt
authkey = File.open(file) {|f| f.readline}
authkey.chomp! 

# we'll pass in a secret to this script from the authorized_keys file
abort unless secret = authkey
 
# prompt the user for their validation code
 
STDERR.write "Enter the validation code: "
until validation_code = STDIN.gets.strip
  sleep 1
end
 
# check the validation code is correct
 
abort "Invalid" unless validation_code == ROTP::TOTP.new(secret).now.to_s
 
# user has validated so we'll give them their shell
 
Kernel.exec ENV['SSH_ORIGINAL_COMMAND'] || ENV['SHELL']