From 2433d39df5aa37c8927b0583ef1dc1c9a2e9c67f Mon Sep 17 00:00:00 2001
From: Stepan Fedorko-Bartos <step7750@gmail.com>
Date: Thu, 15 Nov 2018 18:54:53 -0700
Subject: [PATCH] Allows Custom Yubico OTP Server

---
 .env                       | 3 ++-
 src/api/core/two_factor.rs | 9 ++++++++-
 src/main.rs                | 2 ++
 3 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/.env b/.env
index d960f8e..0f73893 100644
--- a/.env
+++ b/.env
@@ -43,9 +43,10 @@
 ## Yubico (Yubikey) Settings
 ## Set your Client ID and Secret Key for Yubikey OTP
 ## You can generate it here: https://upgrade.yubico.com/getapikey/
-## TODO: Allow choosing custom YubiCloud server
+## You can optionally specify a custom OTP server
 # YUBICO_CLIENT_ID=11111
 # YUBICO_SECRET_KEY=AAAAAAAAAAAAAAAAAAAAAAAA
+# YUBICO_SERVER=http://yourdomain.com/wsapi/2.0/verify
 
 ## Rocket specific settings, check Rocket documentation to learn more
 # ROCKET_ENV=staging
diff --git a/src/api/core/two_factor.rs b/src/api/core/two_factor.rs
index 2ba03fc..7e09e22 100644
--- a/src/api/core/two_factor.rs
+++ b/src/api/core/two_factor.rs
@@ -561,7 +561,14 @@ fn verify_yubikey_otp(otp: String) -> JsonResult {
     let yubico = Yubico::new();
     let config = Config::default().set_client_id(CONFIG.yubico_client_id.to_owned()).set_key(CONFIG.yubico_secret_key.to_owned());
 
-    let result = yubico.verify(otp, config);
+    let result;
+
+    if CONFIG.yubico_server.is_some() {
+        result = yubico.verify(otp, config.set_api_hosts(vec![CONFIG.yubico_server.to_owned().unwrap()]));
+    }
+    else {
+        result = yubico.verify(otp, config);
+    }
 
     match result {
         Ok(_answer) => Ok(Json(json!({}))),
diff --git a/src/main.rs b/src/main.rs
index f4c6d34..feb241a 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -249,6 +249,7 @@ pub struct Config {
     yubico_cred_set: bool,
     yubico_client_id: String,
     yubico_secret_key: String,
+    yubico_server: Option<String>,
 
     mail: Option<MailConfig>,
 }
@@ -294,6 +295,7 @@ impl Config {
             yubico_cred_set: yubico_client_id.is_some() && yubico_secret_key.is_some(),
             yubico_client_id: yubico_client_id.unwrap_or("00000".into()),
             yubico_secret_key: yubico_secret_key.unwrap_or("AAAAAAA".into()),
+            yubico_server: get_env("YUBICO_SERVER"),
 
             mail: MailConfig::load(),
         }