keycloak-memoizeit

diff --git a/services/src/main/java/org/keycloak/services/models/nosql/keycloak/adapters/MongoDBSessionFactory.java b/services/src/main/java/org/keycloak/services/models/nosql/keycloak/adapters/MongoDBSessionFactory.java
index 73ba461..6f1d7f7 100644
--- a/services/src/main/java/org/keycloak/services/models/nosql/keycloak/adapters/MongoDBSessionFactory.java
+++ b/services/src/main/java/org/keycloak/services/models/nosql/keycloak/adapters/MongoDBSessionFactory.java
@@ -20,6 +20,7 @@ import org.keycloak.services.models.nosql.keycloak.data.SocialLinkData;
 import org.keycloak.services.models.nosql.keycloak.data.UserData;
 import org.keycloak.services.models.nosql.impl.MongoDBImpl;
 import org.keycloak.services.models.nosql.impl.MongoDBQueryBuilder;
+import org.keycloak.services.models.nosql.keycloak.data.credentials.OTPData;
 import org.keycloak.services.models.nosql.keycloak.data.credentials.PasswordData;
 
 /**
@@ -36,6 +37,7 @@ public class MongoDBSessionFactory implements KeycloakSessionFactory {
             RoleData.class,
             RequiredCredentialData.class,
             PasswordData.class,
+            OTPData.class,
             SocialLinkData.class,
             ApplicationData.class
     };

diff --git a/services/src/main/java/org/keycloak/services/models/nosql/keycloak/adapters/RealmAdapter.java b/services/src/main/java/org/keycloak/services/models/nosql/keycloak/adapters/RealmAdapter.java
index 24b0146..4d613d5 100644
--- a/services/src/main/java/org/keycloak/services/models/nosql/keycloak/adapters/RealmAdapter.java
+++ b/services/src/main/java/org/keycloak/services/models/nosql/keycloak/adapters/RealmAdapter.java
@@ -32,6 +32,8 @@ import org.keycloak.services.models.nosql.keycloak.data.RoleData;
 import org.keycloak.services.models.nosql.keycloak.data.SocialLinkData;
 import org.keycloak.services.models.nosql.keycloak.data.UserData;
 import org.picketlink.idm.credential.Credentials;
+import org.picketlink.idm.credential.Password;
+import org.picketlink.idm.credential.TOTPCredentials;
 
 /**
  * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
@@ -44,9 +46,9 @@ public class RealmAdapter implements RealmModel {
     protected volatile transient PublicKey publicKey;
     protected volatile transient PrivateKey privateKey;
 
-    // TODO: likely shouldn't be static. And setup is not called ATM, which means that it's not possible to configure stuff like PasswordEncoder etc.
-    private static PasswordCredentialHandler passwordCredentialHandler = new PasswordCredentialHandler();
-    private static TOTPCredentialHandler totpCredentialHandler = new TOTPCredentialHandler();
+    // TODO: likely shouldn't be static. And ATM, just empty map is passed -> It's not possible to configure stuff like PasswordEncoder etc.
+    private static PasswordCredentialHandler passwordCredentialHandler = new PasswordCredentialHandler(new HashMap<String, Object>());
+    private static TOTPCredentialHandler totpCredentialHandler = new TOTPCredentialHandler(new HashMap<String, Object>());
 
     public RealmAdapter(RealmData realmData, NoSQL noSQL) {
         this.realm = realmData;
@@ -659,7 +661,8 @@ public class RealmAdapter implements RealmModel {
 
     @Override
     public boolean validateTOTP(UserModel user, String password, String token) {
-        return false;  //To change body of implemented methods use File | Settings | File Templates.
+        Credentials.Status status = totpCredentialHandler.validate(noSQL, ((UserAdapter)user).getUser(), password, token, null);
+        return status == Credentials.Status.VALID;
     }
 
     @Override
@@ -667,10 +670,7 @@ public class RealmAdapter implements RealmModel {
         if (cred.getType().equals(CredentialRepresentation.PASSWORD)) {
             passwordCredentialHandler.update(noSQL, ((UserAdapter)user).getUser(), cred.getValue(), null, null);
         } else if (cred.getType().equals(CredentialRepresentation.TOTP)) {
-            // TODO
-//            TOTPCredential totp = new TOTPCredential(cred.getValue());
-//            totp.setDevice(cred.getDevice());
-//            idm.updateCredential(((UserAdapter)user).getUser(), totp);
+            totpCredentialHandler.update(noSQL, ((UserAdapter)user).getUser(), cred.getValue(), cred.getDevice(), null, null);
         } else if (cred.getType().equals(CredentialRepresentation.CLIENT_CERT)) {
             // TODO
 //            X509Certificate cert = null;

diff --git a/services/src/main/java/org/keycloak/services/models/nosql/keycloak/credentials/PasswordCredentialHandler.java b/services/src/main/java/org/keycloak/services/models/nosql/keycloak/credentials/PasswordCredentialHandler.java
index 7709987..52d706e 100644
--- a/services/src/main/java/org/keycloak/services/models/nosql/keycloak/credentials/PasswordCredentialHandler.java
+++ b/services/src/main/java/org/keycloak/services/models/nosql/keycloak/credentials/PasswordCredentialHandler.java
@@ -1,15 +1,11 @@
 package org.keycloak.services.models.nosql.keycloak.credentials;
 
-import java.security.NoSuchAlgorithmException;
-import java.security.SecureRandom;
 import java.util.Date;
 import java.util.Map;
 import java.util.UUID;
 
 import org.keycloak.services.models.nosql.api.NoSQL;
 import org.keycloak.services.models.nosql.api.query.NoSQLQuery;
-import org.keycloak.services.models.nosql.api.query.NoSQLQueryBuilder;
-import org.keycloak.services.models.nosql.impl.MongoDBQueryBuilder;
 import org.keycloak.services.models.nosql.keycloak.data.UserData;
 import org.keycloak.services.models.nosql.keycloak.data.credentials.PasswordData;
 import org.picketlink.idm.credential.Credentials;
@@ -34,7 +30,11 @@ public class PasswordCredentialHandler {
 
     private PasswordEncoder passwordEncoder = new SHAPasswordEncoder(512);;
 
-    public void setup(Map<String, Object> options) {
+    public PasswordCredentialHandler(Map<String, Object> options) {
+        setup(options);
+    }
+
+    private void setup(Map<String, Object> options) {
         if (options != null) {
             Object providedEncoder = options.get(PASSWORD_ENCODER);
 
@@ -64,6 +64,7 @@ public class PasswordCredentialHandler {
 
                 // If the stored hash is null we automatically fail validation
                 if (passwordData != null) {
+                    // TODO: Status.INVALID should have bigger priority than Status.EXPIRED?
                     if (!isCredentialExpired(passwordData.getExpiryDate())) {
 
                         boolean matches = this.passwordEncoder.verify(saltPassword(passwordToValidate, passwordData.getSalt()), passwordData.getEncodedHash());

diff --git a/services/src/main/java/org/keycloak/services/models/nosql/keycloak/credentials/TOTPCredentialHandler.java b/services/src/main/java/org/keycloak/services/models/nosql/keycloak/credentials/TOTPCredentialHandler.java
index a5dc8d2..84b5f0a 100644
--- a/services/src/main/java/org/keycloak/services/models/nosql/keycloak/credentials/TOTPCredentialHandler.java
+++ b/services/src/main/java/org/keycloak/services/models/nosql/keycloak/credentials/TOTPCredentialHandler.java
@@ -1,5 +1,21 @@
 package org.keycloak.services.models.nosql.keycloak.credentials;
 
+import java.util.Date;
+import java.util.Map;
+
+import org.keycloak.services.models.nosql.api.NoSQL;
+import org.keycloak.services.models.nosql.api.query.NoSQLQuery;
+import org.keycloak.services.models.nosql.keycloak.data.UserData;
+import org.keycloak.services.models.nosql.keycloak.data.credentials.OTPData;
+import org.picketlink.idm.credential.Credentials;
+import org.picketlink.idm.credential.util.TimeBasedOTP;
+
+import static org.picketlink.common.util.StringUtil.isNullOrEmpty;
+import static org.picketlink.idm.credential.util.TimeBasedOTP.DEFAULT_ALGORITHM;
+import static org.picketlink.idm.credential.util.TimeBasedOTP.DEFAULT_DELAY_WINDOW;
+import static org.picketlink.idm.credential.util.TimeBasedOTP.DEFAULT_INTERVAL_SECONDS;
+import static org.picketlink.idm.credential.util.TimeBasedOTP.DEFAULT_NUMBER_DIGITS;
+
 /**
  * Defacto forked from {@link org.picketlink.idm.credential.handler.TOTPCredentialHandler}
  *
@@ -7,5 +23,116 @@ package org.keycloak.services.models.nosql.keycloak.credentials;
  */
 public class TOTPCredentialHandler extends PasswordCredentialHandler {
 
-    // TODO: implement
+    public static final String ALGORITHM = "ALGORITHM";
+    public static final String INTERVAL_SECONDS = "INTERVAL_SECONDS";
+    public static final String NUMBER_DIGITS = "NUMBER_DIGITS";
+    public static final String DELAY_WINDOW = "DELAY_WINDOW";
+    public static final String DEFAULT_DEVICE = "DEFAULT_DEVICE";
+
+    private TimeBasedOTP totp;
+
+    public TOTPCredentialHandler(Map<String, Object> options) {
+        super(options);
+        setup(options);
+    }
+
+    private void setup(Map<String, Object> options) {
+        String algorithm = getConfigurationProperty(options, ALGORITHM, DEFAULT_ALGORITHM);
+        String intervalSeconds = getConfigurationProperty(options, INTERVAL_SECONDS, "" + DEFAULT_INTERVAL_SECONDS);
+        String numberDigits = getConfigurationProperty(options, NUMBER_DIGITS, "" + DEFAULT_NUMBER_DIGITS);
+        String delayWindow = getConfigurationProperty(options, DELAY_WINDOW, "" + DEFAULT_DELAY_WINDOW);
+
+        this.totp = new TimeBasedOTP(algorithm, Integer.parseInt(numberDigits), Integer.valueOf(intervalSeconds), Integer.valueOf(delayWindow));
+    }
+
+    public Credentials.Status validate(NoSQL noSQL, UserData user, String passwordToValidate, String token, String device) {
+        Credentials.Status status = super.validate(noSQL, user, passwordToValidate);
+
+        if (Credentials.Status.VALID != status) {
+            return status;
+        }
+
+        device = getDevice(device);
+
+        user = noSQL.loadObject(UserData.class, user.getId());
+
+        // If the user for the provided username cannot be found we fail validation
+        if (user != null) {
+            if (user.isEnabled()) {
+
+                // Try to find OTP based on userId and device (For now assume that this is unique combination)
+                NoSQLQuery query = noSQL.createQueryBuilder()
+                        .andCondition("userId", user.getId())
+                        .andCondition("device", device)
+                        .build();
+                OTPData otpData = noSQL.loadSingleObject(OTPData.class, query);
+
+                // If the stored OTP is null we automatically fail validation
+                if (otpData != null) {
+                    // TODO: Status.INVALID should have bigger priority than Status.EXPIRED?
+                    if (!PasswordCredentialHandler.isCredentialExpired(otpData.getExpiryDate())) {
+                        boolean isValid = this.totp.validate(token, otpData.getSecretKey().getBytes());
+                        if (!isValid) {
+                            status = Credentials.Status.INVALID;
+                        }
+                    }  else {
+                        status = Credentials.Status.EXPIRED;
+                    }
+                } else {
+                    status = Credentials.Status.UNVALIDATED;
+                }
+            } else {
+                status = Credentials.Status.ACCOUNT_DISABLED;
+            }
+        } else {
+            status = Credentials.Status.INVALID;
+        }
+
+        return status;
+    }
+
+    public void update(NoSQL noSQL, UserData user, String secret, String device, Date effectiveDate, Date expiryDate) {
+        device = getDevice(device);
+
+        // Try to look if user already has otp (Right now, supports just one OTP per user)
+        NoSQLQuery query = noSQL.createQueryBuilder()
+                .andCondition("userId", user.getId())
+                .andCondition("device", device)
+                .build();
+
+        OTPData otpData = noSQL.loadSingleObject(OTPData.class, query);
+        if (otpData == null) {
+            otpData = new OTPData();
+        }
+
+        otpData.setSecretKey(secret);
+        otpData.setDevice(device);
+
+        if (effectiveDate != null) {
+            otpData.setEffectiveDate(effectiveDate);
+        }
+
+        otpData.setExpiryDate(expiryDate);
+        otpData.setUserId(user.getId());
+
+        noSQL.saveObject(otpData);
+    }
+
+    private String getDevice(String device) {
+        if (isNullOrEmpty(device)) {
+            device = DEFAULT_DEVICE;
+        }
+
+        return device;
+    }
+
+    private String getConfigurationProperty(Map<String, Object> options, String key, String defaultValue) {
+        Object value = options.get(key);
+
+        if (value != null) {
+            return String.valueOf(value);
+        }
+
+        return defaultValue;
+    }
 }

diff --git a/services/src/main/java/org/keycloak/services/models/nosql/keycloak/data/credentials/OTPData.java b/services/src/main/java/org/keycloak/services/models/nosql/keycloak/data/credentials/OTPData.java
new file mode 100644
index 0000000..cf4fef8
--- /dev/null
+++ b/services/src/main/java/org/keycloak/services/models/nosql/keycloak/data/credentials/OTPData.java
@@ -0,0 +1,66 @@
+package org.keycloak.services.models.nosql.keycloak.data.credentials;
+
+import java.util.Date;
+
+import org.keycloak.services.models.nosql.api.AbstractNoSQLObject;
+import org.keycloak.services.models.nosql.api.NoSQLCollection;
+import org.keycloak.services.models.nosql.api.NoSQLField;
+
+/**
+ * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
+ */
+@NoSQLCollection(collectionName = "otpCredentials")
+public class OTPData extends AbstractNoSQLObject {
+
+    private Date effectiveDate = new Date();
+    private Date expiryDate;
+    private String secretKey;
+    private String device;
+
+    private String userId;
+
+    @NoSQLField
+    public Date getEffectiveDate() {
+        return effectiveDate;
+    }
+
+    public void setEffectiveDate(Date effectiveDate) {
+        this.effectiveDate = effectiveDate;
+    }
+
+    @NoSQLField
+    public Date getExpiryDate() {
+        return expiryDate;
+    }
+
+    public void setExpiryDate(Date expiryDate) {
+        this.expiryDate = expiryDate;
+    }
+
+    @NoSQLField
+    public String getSecretKey() {
+        return secretKey;
+    }
+
+    public void setSecretKey(String secretKey) {
+        this.secretKey = secretKey;
+    }
+
+    @NoSQLField
+    public String getDevice() {
+        return device;
+    }
+
+    public void setDevice(String device) {
+        this.device = device;
+    }
+
+    @NoSQLField
+    public String getUserId() {
+        return userId;
+    }
+
+    public void setUserId(String userId) {
+        this.userId = userId;
+    }
+}

diff --git a/services/src/test/java/org/keycloak/services/managers/AuthenticationManagerTest.java b/services/src/test/java/org/keycloak/services/managers/AuthenticationManagerTest.java
index ce3ed3d..a2049fd 100755
--- a/services/src/test/java/org/keycloak/services/managers/AuthenticationManagerTest.java
+++ b/services/src/test/java/org/keycloak/services/managers/AuthenticationManagerTest.java
@@ -18,24 +18,20 @@ import org.keycloak.models.UserCredentialModel;
 import org.keycloak.models.UserModel;
 import org.keycloak.models.UserModel.RequiredAction;
 import org.keycloak.services.resources.KeycloakApplication;
+import org.keycloak.test.common.AbstractKeycloakTest;
+import org.keycloak.test.common.SessionFactoryTestContext;
 import org.picketlink.idm.credential.util.TimeBasedOTP;
 
-public class AuthenticationManagerTest {
+public class AuthenticationManagerTest extends AbstractKeycloakTest {
 
-    private RealmManager adapter;
     private AuthenticationManager am;
-    private KeycloakSessionFactory factory;
     private MultivaluedMap<String, String> formData;
-    private KeycloakSession identitySession;
     private TimeBasedOTP otp;
     private RealmModel realm;
     private UserModel user;
 
-    @After
-    public void after() throws Exception {
-        identitySession.getTransaction().commit();
-        identitySession.close();
-        factory.close();
+    public AuthenticationManagerTest(SessionFactoryTestContext testContext) {
+        super(testContext);
     }
 
     @Test
@@ -134,12 +130,8 @@ public class AuthenticationManagerTest {
 
     @Before
     public void before() throws Exception {
-        factory = KeycloakApplication.buildSessionFactory();
-        identitySession = factory.createSession();
-        identitySession.getTransaction().begin();
-        adapter = new RealmManager(identitySession);
-
-        realm = adapter.createRealm("Test");
+        super.before();
+        realm = getRealmManager().createRealm("Test");
         realm.setAccessCodeLifespan(100);
         realm.setCookieLoginAllowed(true);
         realm.setEnabled(true);