Offer TOTP Autofill for OTP fields (#899)

10 files changed, 107 insertions, 37 deletions

diff --git a/app/src/main/java/com/zeapo/pwdstore/autofill/oreo/AutofillHelper.kt b/app/src/main/java/com/zeapo/pwdstore/autofill/oreo/AutofillHelper.kt
index d417484b..2a8443ba 100644
--- a/app/src/main/java/com/zeapo/pwdstore/autofill/oreo/AutofillHelper.kt
+++ b/app/src/main/java/com/zeapo/pwdstore/autofill/oreo/AutofillHelper.kt
@@ -86,7 +86,7 @@ val AssistStructure.ViewNode.webOrigin: String?
         "$scheme://$domain"
     }
 
-data class Credentials(val username: String?, val password: String) {
+data class Credentials(val username: String?, val password: String, val otp: String?) {
     companion object {
         fun fromStoreEntry(
             context: Context,
@@ -98,7 +98,7 @@ data class Credentials(val username: String?, val password: String) {
             val username = entry.username
                 ?: directoryStructure.getUsernameFor(file)
                 ?: context.getDefaultUsername()
-            return Credentials(username, entry.password)
+            return Credentials(username, entry.password, entry.calculateTotpCode())
         }
     }
 }
diff --git a/app/src/main/java/com/zeapo/pwdstore/autofill/oreo/AutofillScenario.kt b/app/src/main/java/com/zeapo/pwdstore/autofill/oreo/AutofillScenario.kt
index f1514851..8e209a60 100644
--- a/app/src/main/java/com/zeapo/pwdstore/autofill/oreo/AutofillScenario.kt
+++ b/app/src/main/java/com/zeapo/pwdstore/autofill/oreo/AutofillScenario.kt
@@ -29,6 +29,7 @@ sealed class AutofillScenario<out T : Any> {
     companion object {
         const val BUNDLE_KEY_USERNAME_ID = "usernameId"
         const val BUNDLE_KEY_FILL_USERNAME = "fillUsername"
+        const val BUNDLE_KEY_OTP_ID = "otpId"
         const val BUNDLE_KEY_CURRENT_PASSWORD_IDS = "currentPasswordIds"
         const val BUNDLE_KEY_NEW_PASSWORD_IDS = "newPasswordIds"
         const val BUNDLE_KEY_GENERIC_PASSWORD_IDS = "genericPasswordIds"
@@ -38,6 +39,7 @@ sealed class AutofillScenario<out T : Any> {
                 Builder<AutofillId>().apply {
                     username = clientState.getParcelable(BUNDLE_KEY_USERNAME_ID)
                     fillUsername = clientState.getBoolean(BUNDLE_KEY_FILL_USERNAME)
+                    otp = clientState.getParcelable(BUNDLE_KEY_OTP_ID)
                     currentPassword.addAll(
                         clientState.getParcelableArrayList(
                             BUNDLE_KEY_CURRENT_PASSWORD_IDS
@@ -64,6 +66,7 @@ sealed class AutofillScenario<out T : Any> {
     class Builder<T : Any> {
         var username: T? = null
         var fillUsername = false
+        var otp: T? = null
         val currentPassword = mutableListOf<T>()
         val newPassword = mutableListOf<T>()
         val genericPassword = mutableListOf<T>()
@@ -74,6 +77,7 @@ sealed class AutofillScenario<out T : Any> {
                 ClassifiedAutofillScenario(
                     username = username,
                     fillUsername = fillUsername,
+                    otp = otp,
                     currentPassword = currentPassword,
                     newPassword = newPassword
                 )
@@ -81,6 +85,7 @@ sealed class AutofillScenario<out T : Any> {
                 GenericAutofillScenario(
                     username = username,
                     fillUsername = fillUsername,
+                    otp = otp,
                     genericPassword = genericPassword
                 )
             }
@@ -89,6 +94,7 @@ sealed class AutofillScenario<out T : Any> {
 
     abstract val username: T?
     abstract val fillUsername: Boolean
+    abstract val otp: T?
     abstract val allPasswordFields: List<T>
     abstract val passwordFieldsToFillOnMatch: List<T>
     abstract val passwordFieldsToFillOnSearch: List<T>
@@ -99,19 +105,19 @@ sealed class AutofillScenario<out T : Any> {
         get() = listOfNotNull(username) + passwordFieldsToSave
 
     val allFields
-        get() = listOfNotNull(username) + allPasswordFields
+        get() = listOfNotNull(username, otp) + allPasswordFields
 
     fun fieldsToFillOn(action: AutofillAction): List<T> {
-        val passwordFieldsToFill = when (action) {
-            AutofillAction.Match -> passwordFieldsToFillOnMatch
-            AutofillAction.Search -> passwordFieldsToFillOnSearch
+        val credentialFieldsToFill = when (action) {
+            AutofillAction.Match -> passwordFieldsToFillOnMatch + listOfNotNull(otp)
+            AutofillAction.Search -> passwordFieldsToFillOnSearch + listOfNotNull(otp)
             AutofillAction.Generate -> passwordFieldsToFillOnGenerate
         }
         return when {
-            passwordFieldsToFill.isNotEmpty() -> {
+            credentialFieldsToFill.isNotEmpty() -> {
                 // If the current action would fill into any password field, we also fill into the
                 // username field if possible.
-                listOfNotNull(username.takeIf { fillUsername }) + passwordFieldsToFill
+                listOfNotNull(username.takeIf { fillUsername }) + credentialFieldsToFill
             }
             allPasswordFields.isEmpty() && action != AutofillAction.Generate -> {
                 // If there no password fields at all, we still offer to fill the username, e.g. in
@@ -127,6 +133,7 @@ sealed class AutofillScenario<out T : Any> {
 data class ClassifiedAutofillScenario<T : Any>(
     override val username: T?,
     override val fillUsername: Boolean,
+    override val otp: T?,
     val currentPassword: List<T>,
     val newPassword: List<T>
 ) : AutofillScenario<T>() {
@@ -147,6 +154,7 @@ data class ClassifiedAutofillScenario<T : Any>(
 data class GenericAutofillScenario<T : Any>(
     override val username: T?,
     override val fillUsername: Boolean,
+    override val otp: T?,
     val genericPassword: List<T>
 ) : AutofillScenario<T>() {
 
@@ -183,14 +191,15 @@ fun Dataset.Builder.fillWith(
 ) {
     val credentialsToFill = credentials ?: Credentials(
         "USERNAME",
-        "PASSWORD"
+        "PASSWORD",
+        "OTP"
     )
     for (field in scenario.fieldsToFillOn(action)) {
-        val value = if (field == scenario.username) {
-            credentialsToFill.username
-        } else {
-            credentialsToFill.password
-        } ?: continue
+        val value = when (field) {
+            scenario.username -> credentialsToFill.username
+            scenario.otp -> credentialsToFill.otp
+            else -> credentialsToFill.password
+        }
         setValue(field, AutofillValue.forText(value))
     }
 }
@@ -209,6 +218,7 @@ inline fun <T : Any, S : Any> AutofillScenario<T>.map(transform: (T) -> S): Auto
     val builder = AutofillScenario.Builder<S>()
     builder.username = username?.let(transform)
     builder.fillUsername = fillUsername
+    builder.otp = otp?.let(transform)
     when (this) {
         is ClassifiedAutofillScenario -> {
             builder.currentPassword.addAll(currentPassword.map(transform))
@@ -225,9 +235,10 @@ inline fun <T : Any, S : Any> AutofillScenario<T>.map(transform: (T) -> S): Auto
 @JvmName("toBundleAutofillId")
 private fun AutofillScenario<AutofillId>.toBundle(): Bundle = when (this) {
     is ClassifiedAutofillScenario<AutofillId> -> {
-        Bundle(4).apply {
+        Bundle(5).apply {
             putParcelable(AutofillScenario.BUNDLE_KEY_USERNAME_ID, username)
             putBoolean(AutofillScenario.BUNDLE_KEY_FILL_USERNAME, fillUsername)
+            putParcelable(AutofillScenario.BUNDLE_KEY_OTP_ID, otp)
             putParcelableArrayList(
                 AutofillScenario.BUNDLE_KEY_CURRENT_PASSWORD_IDS, ArrayList(currentPassword)
             )
@@ -237,9 +248,10 @@ private fun AutofillScenario<AutofillId>.toBundle(): Bundle = when (this) {
         }
     }
     is GenericAutofillScenario<AutofillId> -> {
-        Bundle(3).apply {
+        Bundle(4).apply {
             putParcelable(AutofillScenario.BUNDLE_KEY_USERNAME_ID, username)
             putBoolean(AutofillScenario.BUNDLE_KEY_FILL_USERNAME, fillUsername)
+            putParcelable(AutofillScenario.BUNDLE_KEY_OTP_ID, otp)
             putParcelableArrayList(
                 AutofillScenario.BUNDLE_KEY_GENERIC_PASSWORD_IDS, ArrayList(genericPassword)
             )
diff --git a/app/src/main/java/com/zeapo/pwdstore/autofill/oreo/AutofillStrategy.kt b/app/src/main/java/com/zeapo/pwdstore/autofill/oreo/AutofillStrategy.kt
index 6f3b4ff5..790a72e1 100644
--- a/app/src/main/java/com/zeapo/pwdstore/autofill/oreo/AutofillStrategy.kt
+++ b/app/src/main/java/com/zeapo/pwdstore/autofill/oreo/AutofillStrategy.kt
@@ -166,6 +166,13 @@ val autofillStrategy = strategy {
         }
     }
 
+    // Match a single focused OTP field.
+    rule(applyInSingleOriginMode = true) {
+        otp {
+            takeSingle { otpCertainty >= Likely && isFocused }
+        }
+    }
+
     // Match a single focused username field without a password field.
     rule(applyInSingleOriginMode = true) {
         username {
diff --git a/app/src/main/java/com/zeapo/pwdstore/autofill/oreo/AutofillStrategyDsl.kt b/app/src/main/java/com/zeapo/pwdstore/autofill/oreo/AutofillStrategyDsl.kt
index 3b648234..5e6f460e 100644
--- a/app/src/main/java/com/zeapo/pwdstore/autofill/oreo/AutofillStrategyDsl.kt
+++ b/app/src/main/java/com/zeapo/pwdstore/autofill/oreo/AutofillStrategyDsl.kt
@@ -164,7 +164,7 @@ class AutofillRule private constructor(
     )
 
     enum class FillableFieldType {
-        Username, CurrentPassword, NewPassword, GenericPassword,
+        Username, Otp, CurrentPassword, NewPassword, GenericPassword,
     }
 
     @AutofillDsl
@@ -192,6 +192,18 @@ class AutofillRule private constructor(
             )
         }
 
+        fun otp(optional: Boolean = false, block: SingleFieldMatcher.Builder.() -> Unit) {
+            require(matchers.none { it.type == FillableFieldType.Otp }) { "Every rule block can only have at most one otp block" }
+            matchers.add(
+                AutofillRuleMatcher(
+                    type = FillableFieldType.Otp,
+                    matcher = SingleFieldMatcher.Builder().apply(block).build(),
+                    optional = optional,
+                    matchHidden = false
+                )
+            )
+        }
+
         fun currentPassword(optional: Boolean = false, matchHidden: Boolean = false, block: FieldMatcher.Builder.() -> Unit) {
             require(matchers.none { it.type == FillableFieldType.GenericPassword }) { "Every rule block can only have either genericPassword or {current,new}Password blocks" }
             matchers.add(
@@ -247,6 +259,7 @@ class AutofillRule private constructor(
     fun match(
         allPassword: List<FormField>,
         allUsername: List<FormField>,
+        allOtp: List<FormField>,
         singleOriginMode: Boolean,
         isManualRequest: Boolean
     ): AutofillScenario<FormField>? {
@@ -264,6 +277,7 @@ class AutofillRule private constructor(
         for ((type, matcher, optional, matchHidden) in matchers) {
             val fieldsToMatchOn = when (type) {
                 FillableFieldType.Username -> allUsername
+                FillableFieldType.Otp -> allOtp
                 else -> allPassword
             }.filter { matchHidden || it.isVisible }
             val matchResult = matcher.match(fieldsToMatchOn, alreadyMatched) ?: if (optional) {
@@ -281,6 +295,10 @@ class AutofillRule private constructor(
                     // Hidden username fields should be saved but not filled.
                     scenarioBuilder.fillUsername = scenarioBuilder.username!!.isVisible == true
                 }
+                FillableFieldType.Otp -> {
+                    check(matchResult.size == 1 && scenarioBuilder.otp == null)
+                    scenarioBuilder.otp = matchResult.single()
+                }
                 FillableFieldType.CurrentPassword -> scenarioBuilder.currentPassword.addAll(
                     matchResult
                 )
@@ -338,12 +356,16 @@ class AutofillStrategy private constructor(private val rules: List<AutofillRule>
         val possibleUsernameFields =
             fields.filter { it.usernameCertainty >= CertaintyLevel.Possible }
         d { "Possible username fields: ${possibleUsernameFields.size}" }
+        val possibleOtpFields =
+            fields.filter { it.otpCertainty >= CertaintyLevel.Possible }
+        d { "Possible otp fields: ${possibleOtpFields.size}" }
         // Return the result of the first rule that matches
         d { "Rules: ${rules.size}" }
         for (rule in rules) {
             return rule.match(
                 possiblePasswordFields,
                 possibleUsernameFields,
+                possibleOtpFields,
                 singleOriginMode = singleOriginMode,
                 isManualRequest = isManualRequest
             )
diff --git a/app/src/main/java/com/zeapo/pwdstore/autofill/oreo/FormField.kt b/app/src/main/java/com/zeapo/pwdstore/autofill/oreo/FormField.kt
index 0c96b587..3c1e3a0a 100644
--- a/app/src/main/java/com/zeapo/pwdstore/autofill/oreo/FormField.kt
+++ b/app/src/main/java/com/zeapo/pwdstore/autofill/oreo/FormField.kt
@@ -10,6 +10,7 @@ import android.text.InputType
 import android.view.View
 import android.view.autofill.AutofillId
 import androidx.annotation.RequiresApi
+import androidx.autofill.HintConstants
 import java.util.Locale
 
 enum class CertaintyLevel {
@@ -31,14 +32,21 @@ class FormField(
     companion object {
 
         @RequiresApi(Build.VERSION_CODES.O)
-        private val HINTS_USERNAME = listOf(View.AUTOFILL_HINT_USERNAME)
+        private val HINTS_USERNAME = listOf(HintConstants.AUTOFILL_HINT_USERNAME)
 
         @RequiresApi(Build.VERSION_CODES.O)
-        private val HINTS_PASSWORD = listOf(View.AUTOFILL_HINT_PASSWORD)
+        private val HINTS_PASSWORD = listOf(HintConstants.AUTOFILL_HINT_PASSWORD)
 
         @RequiresApi(Build.VERSION_CODES.O)
-        private val HINTS_FILLABLE = HINTS_USERNAME + HINTS_PASSWORD + listOf(
-            View.AUTOFILL_HINT_EMAIL_ADDRESS, View.AUTOFILL_HINT_NAME, View.AUTOFILL_HINT_PHONE
+        private val HINTS_OTP = listOf(HintConstants.AUTOFILL_HINT_SMS_OTP)
+
+        @RequiresApi(Build.VERSION_CODES.O)
+        private val HINTS_FILLABLE = HINTS_USERNAME + HINTS_PASSWORD + HINTS_OTP + listOf(
+            HintConstants.AUTOFILL_HINT_EMAIL_ADDRESS,
+            HintConstants.AUTOFILL_HINT_NAME,
+            HintConstants.AUTOFILL_HINT_PERSON_NAME,
+            HintConstants.AUTOFILL_HINT_PHONE,
+            HintConstants.AUTOFILL_HINT_PHONE_NUMBER
         )
 
         private val ANDROID_TEXT_FIELD_CLASS_NAMES = listOf(
@@ -67,11 +75,12 @@ class FormField(
 
         private val HTML_INPUT_FIELD_TYPES_USERNAME = listOf("email", "tel", "text")
         private val HTML_INPUT_FIELD_TYPES_PASSWORD = listOf("password")
+        private val HTML_INPUT_FIELD_TYPES_OTP = listOf("tel", "text")
         private val HTML_INPUT_FIELD_TYPES_FILLABLE =
-            HTML_INPUT_FIELD_TYPES_USERNAME + HTML_INPUT_FIELD_TYPES_PASSWORD
+            (HTML_INPUT_FIELD_TYPES_USERNAME + HTML_INPUT_FIELD_TYPES_PASSWORD + HTML_INPUT_FIELD_TYPES_OTP).toSet().toList()
 
         @RequiresApi(Build.VERSION_CODES.O)
-        private fun isSupportedHint(hint: String) = hint in HINTS_USERNAME + HINTS_PASSWORD
+        private fun isSupportedHint(hint: String) = hint in HINTS_FILLABLE
 
         private val EXCLUDED_TERMS = listOf(
             "url_bar", // Chrome/Edge/Firefox address bar
@@ -85,6 +94,9 @@ class FormField(
         private val USERNAME_HEURISTIC_TERMS = listOf(
             "alias", "e-mail", "email", "login", "user"
         )
+        private val OTP_HEURISTIC_TERMS = listOf(
+            "code", "otp"
+        )
     }
 
     val autofillId: AutofillId = node.autofillId!!
@@ -120,6 +132,7 @@ class FormField(
         htmlAttributes.entries.joinToString { "${it.key}=${it.value}" }
     private val htmlInputType = htmlAttributes["type"]
     private val htmlName = htmlAttributes["name"] ?: ""
+    private val htmlMaxLength = htmlAttributes["maxlength"]?.toIntOrNull()
     private val isHtmlField = htmlTag == "input"
     private val isHtmlPasswordField =
         isHtmlField && htmlInputType in HTML_INPUT_FIELD_TYPES_PASSWORD
@@ -140,6 +153,7 @@ class FormField(
         if (autofillHints.isEmpty()) false else autofillHints.intersect(HINTS_FILLABLE).isEmpty()
     private val hasAutofillHintPassword = autofillHints.intersect(HINTS_PASSWORD).isNotEmpty()
     private val hasAutofillHintUsername = autofillHints.intersect(HINTS_USERNAME).isNotEmpty()
+    private val hasAutofillHintOtp = autofillHints.intersect(HINTS_OTP).isNotEmpty()
 
     // W3C autocomplete hint detection for HTML fields
     private val htmlAutocomplete = htmlAttributes["autocomplete"]
@@ -151,6 +165,7 @@ class FormField(
     val hasAutocompleteHintNewPassword = htmlAutocomplete == "new-password"
     private val hasAutocompleteHintPassword =
         hasAutocompleteHintCurrentPassword || hasAutocompleteHintNewPassword
+    val hasAutocompleteHintOtp = htmlAutocomplete == "one-time-code"
 
     // Basic autofill exclusion checks
     private val hasAutofillTypeText = node.autofillType == View.AUTOFILL_TYPE_TEXT
@@ -193,8 +208,18 @@ class FormField(
     val passwordCertainty =
         if (isCertainPasswordField) CertaintyLevel.Certain else if (isLikelyPasswordField) CertaintyLevel.Likely else if (isPossiblePasswordField) CertaintyLevel.Possible else CertaintyLevel.Impossible
 
+    // OTP field heuristics (based only on the current field)
+    private val isPossibleOtpField = notExcluded && !isPossiblePasswordField && isTextField
+    private val isCertainOtpField =
+        isPossibleOtpField && (hasAutofillHintOtp || hasAutocompleteHintOtp || htmlMaxLength in 6..8)
+    private val isLikelyOtpField = isPossibleOtpField && (isCertainOtpField || OTP_HEURISTIC_TERMS.any {
+        fieldId.contains(it) || hint.contains(it) || htmlName.contains(it)
+    })
+    val otpCertainty =
+        if (isCertainOtpField) CertaintyLevel.Certain else if (isLikelyOtpField) CertaintyLevel.Likely else if (isPossibleOtpField) CertaintyLevel.Possible else CertaintyLevel.Impossible
+
     // Username field heuristics (based only on the current field)
-    private val isPossibleUsernameField = notExcluded && !isPossiblePasswordField
+    private val isPossibleUsernameField = notExcluded && !isPossiblePasswordField && !isCertainOtpField && isTextField
     private val isCertainUsernameField =
         isPossibleUsernameField && (hasAutofillHintUsername || hasAutocompleteHintUsername)
     private val isLikelyUsernameField = isPossibleUsernameField && (isCertainUsernameField || (USERNAME_HEURISTIC_TERMS.any {
@@ -224,8 +249,8 @@ class FormField(
     override fun toString(): String {
         val field = if (isHtmlTextField) "$htmlTag[type=$htmlInputType]" else className
         val description =
-            "\"$hint\", \"$fieldId\"${if (isFocused) ", focused" else ""}${if (isVisible) ", visible" else ""}, $webOrigin, $htmlAttributesDebug"
-        return "$field ($description): password=$passwordCertainty, username=$usernameCertainty"
+            "\"$hint\", \"$fieldId\"${if (isFocused) ", focused" else ""}${if (isVisible) ", visible" else ""}, $webOrigin, $htmlAttributesDebug, $autofillHints"
+        return "$field ($description): password=$passwordCertainty, username=$usernameCertainty, otp=$otpCertainty"
     }
 
     override fun equals(other: Any?): Boolean {
diff --git a/app/src/main/java/com/zeapo/pwdstore/autofill/oreo/OreoAutofillService.kt b/app/src/main/java/com/zeapo/pwdstore/autofill/oreo/OreoAutofillService.kt
index 350d187b..cdf9a8ff 100644
--- a/app/src/main/java/com/zeapo/pwdstore/autofill/oreo/OreoAutofillService.kt
+++ b/app/src/main/java/com/zeapo/pwdstore/autofill/oreo/OreoAutofillService.kt
@@ -106,7 +106,7 @@ class OreoAutofillService : AutofillService() {
         callback.onSuccess(
             AutofillSaveActivity.makeSaveIntentSender(
                 this,
-                credentials = Credentials(username, password),
+                credentials = Credentials(username, password, null),
                 formOrigin = formOrigin
             )
         )
diff --git a/app/src/main/java/com/zeapo/pwdstore/autofill/oreo/ui/AutofillDecryptActivity.kt b/app/src/main/java/com/zeapo/pwdstore/autofill/oreo/ui/AutofillDecryptActivity.kt
index b66d068b..d7d8daaf 100644
--- a/app/src/main/java/com/zeapo/pwdstore/autofill/oreo/ui/AutofillDecryptActivity.kt
+++ b/app/src/main/java/com/zeapo/pwdstore/autofill/oreo/ui/AutofillDecryptActivity.kt
@@ -100,7 +100,7 @@ class AutofillDecryptActivity : Activity(), CoroutineScope {
         directoryStructure = AutofillPreferences.directoryStructure(this)
         d { action.toString() }
         launch {
-            val credentials = decryptUsernameAndPassword(File(filePath))
+            val credentials = decryptCredential(File(filePath))
             if (credentials == null) {
                 setResult(RESULT_CANCELED)
             } else {
@@ -153,7 +153,7 @@ class AutofillDecryptActivity : Activity(), CoroutineScope {
         }
     }
 
-    private suspend fun decryptUsernameAndPassword(
+    private suspend fun decryptCredential(
         file: File,
         resumeIntent: Intent? = null
     ): Credentials? {
@@ -178,6 +178,7 @@ class AutofillDecryptActivity : Activity(), CoroutineScope {
             OpenPgpApi.RESULT_CODE_SUCCESS -> {
                 try {
                     val entry = withContext(Dispatchers.IO) {
+                        @Suppress("BlockingMethodInNonBlockingContext")
                         PasswordEntry(decryptedOutput)
                     }
                     Credentials.fromStoreEntry(this, file, entry, directoryStructure)
@@ -203,7 +204,7 @@ class AutofillDecryptActivity : Activity(), CoroutineScope {
                             )
                         }
                     }
-                    decryptUsernameAndPassword(file, intentToResume)
+                    decryptCredential(file, intentToResume)
                 } catch (e: Exception) {
                     e(e) { "OpenPgpApi ACTION_DECRYPT_VERIFY failed with user interaction" }
                     null
diff --git a/app/src/main/java/com/zeapo/pwdstore/autofill/oreo/ui/AutofillSaveActivity.kt b/app/src/main/java/com/zeapo/pwdstore/autofill/oreo/ui/AutofillSaveActivity.kt
index fad13ec8..b5bd9e38 100644
--- a/app/src/main/java/com/zeapo/pwdstore/autofill/oreo/ui/AutofillSaveActivity.kt
+++ b/app/src/main/java/com/zeapo/pwdstore/autofill/oreo/ui/AutofillSaveActivity.kt
@@ -130,7 +130,7 @@ class AutofillSaveActivity : Activity() {
                         finish()
                         return
                     }
-                val credentials = Credentials(username, password)
+                val credentials = Credentials(username, password, null)
                 val fillInDataset = FillableForm.makeFillInDataset(
                     this,
                     credentials,
diff --git a/app/src/main/java/com/zeapo/pwdstore/crypto/DecryptActivity.kt b/app/src/main/java/com/zeapo/pwdstore/crypto/DecryptActivity.kt
index ad0b16f7..52353318 100644
--- a/app/src/main/java/com/zeapo/pwdstore/crypto/DecryptActivity.kt
+++ b/app/src/main/java/com/zeapo/pwdstore/crypto/DecryptActivity.kt
@@ -200,12 +200,7 @@ class DecryptActivity : BasePgpActivity(), OpenPgpServiceConnection.OnBound {
                                         }
                                         launch(Dispatchers.IO) {
                                             repeat(Int.MAX_VALUE) {
-                                                val code = Otp.calculateCode(
-                                                    entry.totpSecret!!,
-                                                    Date().time / (1000 * entry.totpPeriod),
-                                                    entry.totpAlgorithm,
-                                                    entry.digits
-                                                ) ?: "Error"
+                                                val code = entry.calculateTotpCode() ?: "Error"
                                                 withContext(Dispatchers.Main) {
                                                     otpText.setText(code)
                                                 }
diff --git a/app/src/main/java/com/zeapo/pwdstore/model/PasswordEntry.kt b/app/src/main/java/com/zeapo/pwdstore/model/PasswordEntry.kt
index da2d57c6..27a0c584 100644
--- a/app/src/main/java/com/zeapo/pwdstore/model/PasswordEntry.kt
+++ b/app/src/main/java/com/zeapo/pwdstore/model/PasswordEntry.kt
@@ -4,10 +4,12 @@
  */
 package com.zeapo.pwdstore.model
 
+import com.zeapo.pwdstore.utils.Otp
 import com.zeapo.pwdstore.utils.TotpFinder
 import com.zeapo.pwdstore.utils.UriTotpFinder
 import java.io.ByteArrayOutputStream
 import java.io.UnsupportedEncodingException
+import java.util.Date
 
 /**
  * A single entry in password store. [totpFinder] is an implementation of [TotpFinder] that let's us
@@ -50,6 +52,12 @@ class PasswordEntry(content: String, private val totpFinder: TotpFinder = UriTot
         return username != null
     }
 
+    fun calculateTotpCode(): String? {
+        if (totpSecret == null)
+            return null
+        return Otp.calculateCode(totpSecret, Date().time / (1000 * totpPeriod), totpAlgorithm, digits)
+    }
+
     val extraContentWithoutAuthData by lazy {
         extraContent.splitToSequence("\n").filter { line ->
             return@filter when {