#!/usr/bin/ruby

# You will need to install the rest-client gem and support package
# on linux:
# sudo apt-get install ruby-dev
##
# The above is to install the header files for ruby, so the gem
# install below can compile the native parts.
##
# sudo gem install rest-client

require 'rest-client'
require 'json'

#
# Usage:
# * Add the following trigger entries:

Triggers =<<EOF
   2fa	auth-pre-2fa	auth	"ruby #{__FILE__} list %user% %host%"
   2fa	auth-init-2fa	auth	"ruby #{__FILE__} init %user% %host% %method%"
   2fa	auth-check-2fa	auth	"ruby #{__FILE__} check %user% %host% %scheme% %token%"
EOF

#
# * Update the Okta credentials below this line:

@OKTA_KEY = 'XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX'
@OKTA_URL = 'https://dev-XXXXXX.oktapreview.com/'
@OKTA_DOMAIN = '@example.com'

def error( message="" )
	puts message
	exit 1
end

def done( scheme, status, message="", otp="", token="")
	obj = {
		"status" => status
	}
	if scheme.size > 0
	          obj["scheme"] = scheme
	end
	if message.size > 0
	          obj["message"] = message
	end
	if otp.size > 0
	          obj["challenge"] = otp
	end
	if token.size > 0
	          obj["token"] = token
	end
	puts JSON.generate( obj )
	exit 0
end

def done_list( methods, status, message="" )
	obj = {
	    "status" => status,
	    "methodlist" => []
	}
	if message.size > 0
	    obj["message"] = message
	end
	methods.each { |k,v|
	    obj["methodlist"].push( [ k, v ] )
	}
	puts JSON.generate( obj )
	exit 0
end

def done_skip( message="" )
	# done( scheme, status, message)
	# 'done' does the exit for us.
	done( "", 2, message )
end

def waiting( message="" )
	#  Return 2 to indicate that we're waiting on the user to accept or deny the 2FA prompt
	done( "", 2, message )
end

def getOktaUser( username )
	res = nil;
	begin
		res = RestClient.get(@OKTA_URL + '/api/v1/users/' + username, 'Authorization' => 'SSWS ' + @OKTA_KEY )
	rescue RestClient::ExceptionWithResponse => err
		error( JSON.parse(err.response)['errorSummary'] )
	end
	body = JSON.parse(res.body)
	error( body['errorSummary'] ) unless res.code == 200

	begin
	  res = RestClient.get(@OKTA_URL + '/api/v1/users/' + body['id'] + '/factors', 'Authorization' => 'SSWS ' + @OKTA_KEY )
	rescue RestClient::ExceptionWithResponse => err
	  error( JSON.parse(err.response)['errorSummary'] )
	end
	body2 = JSON.parse(res.body)
	error( body2['errorSummary'] ) unless res.code == 200

	body['factors'] = body2

	body
end

def verifyOktaOTP( url, passcode )
	res = nil
	begin
	  res = RestClient.post(url, passcode, 'Authorization' => 'SSWS ' + @OKTA_KEY, 'Content-Type' => 'application/json', 'Accepts' => 'application/json' )
	rescue RestClient::ExceptionWithResponse => err
	  error( JSON.parse(err.response)['errorSummary'] )
	end
	body = JSON.parse(res.body)
	error( body['errorSummary'] ) unless res.code.between?(200,299)

	body
end

def pollOktaPush( url )
	res = nil
	begin
		res = RestClient.get(url, 'Authorization' => 'SSWS ' + @OKTA_KEY, 'Content-Type' => 'application/json', 'Accepts' => 'application/json' )
	rescue RestClient::ExceptionWithResponse => err
		error( JSON.parse(err.response)['errorSummary'] )
	end
	body = JSON.parse(res.body)
	error( body['errorSummary'] ) unless res.code.between?(200,299)

	waiting( "Waiting...") if body['factorResult'] == "WAITING"
	error( body['factorResult'] ) if body['factorResult'] != "SUCCESS"
end



@schemes = [ 'otp-generated', 'otp-requested', 'challenge', 'external' ]
@actions = [ :list, :init, :check, :triggers ]

if ARGV.length  < 3 && ARGV[0] != "triggers"
  puts "Not enough arguments!"
  exit 255
end

action = ARGV[0];
user   = ARGV[1];
host   = ARGV[2];
scheme = ARGV[3]; # or method if doing init
token  = ARGV[4];

if !@actions.include?(action.to_sym)
  puts "Unknown action #{action}!"
  exit 255
end

case action.to_sym
when :triggers
	puts( Triggers )
	exit( 0 )
when :list
	userObj = getOktaUser( user + @OKTA_DOMAIN )
	enrolled = userObj['factors'].reject{|f| f["status"] != "ACTIVE"}.map {|f| {'scheme'=>f['factorType'], 'id'=>f['id'], 'provider'=>f['provider'], '_links'=>f['_links'], 'profile'=>f['profile']}}

	methods = {}
	enrolled.each do |m|
		methods[ m['id'] ] = m['provider'] + "-" + m['scheme']
	end
	done_list( methods, 0 )

when :init
	# Get the user+factors (or die)
	userObj = getOktaUser( user + @OKTA_DOMAIN )
	enrolled = userObj['factors'].reject{|f| f["status"] != "ACTIVE"}.map {|f| {'scheme'=>f['factorType'], 'id'=>f['id'], '_links'=>f['_links'], 'profile'=>f['profile']}}

	# Pick a scheme the user is enrolled in
	itt = enrolled.index {|f| f['id'] == scheme }
	error("Unknown method #{scheme}!") if itt == nil
	factor = enrolled[itt]

	if factor['scheme'] == 'push'
		res = verifyOktaOTP( factor['_links']['verify']['href'], '{}' )
		done( "external", 0, "Check your auth app!", "", res['_links']['poll']['href'])
	elsif factor['scheme'] == 'token:software:totp'
		done( "otp-generated", 0, "Use your time based code!", "",  factor['_links']['verify']['href'])
	elsif factor['scheme'] == 'sms'
		res = verifyOktaOTP( factor['_links']['verify']['href'], '{}' )
		done( "otp-requested", 0, "Check your SMS!", "", res['_links']['verify']['href'])
	elsif factor['scheme'] == 'question'
		done( "challenge", 0, "", factor['profile']['questionText'], factor['_links']['self']['href'] +"/verify")
	end


when :check
    if !@schemes.include?(scheme)
	puts "Unknown scheme #{scheme}!"
	exit 255
    end

	case scheme
	when 'otp-generated'
		otp = $stdin.read.strip
		# Check token (or die)
		verifyOktaOTP( token, JSON.generate({'passCode' => otp}) )
		done( "otp-generated", 0, "Okta says yes!")

	when 'otp-requested'
		otp = $stdin.read.strip
		# Check token (or die)
		verifyOktaOTP( token, JSON.generate({'passCode' => otp}) )
		done( "otp-requested", 0, "Okta says yes!")

	when 'challenge'
		otp = $stdin.read.strip
		# Check answer (or die)
		verifyOktaOTP( token, JSON.generate({'answer' => otp}) )
		done( "challenge", 0, "Okta says yes!")

	when 'external'
		pollOktaPush( token )
		done( "external", 0, "Okta says yes!")
	end
end

exit 1