From 921e9f96b9bec5e2bf8633947792d6991956507f Mon Sep 17 00:00:00 2001
From: Harsh Shandilya <me@msfjarvis.dev>
Date: Sat, 17 Jul 2021 03:13:16 +0530
Subject: Refactor TOTP implementation and expand SteamGuard hacks (#1460)

* UriTotpFinder: commonize query parameter handling

* gitignore: add more IDEA files

* TotpFinder: add `findIssuer`

* PasswordEntry: don't eagerly fetch TOTP related fields

* format-common: expand SteamGuard workaround

* CHANGELOG: add SteamGuard workaround
---
 format-common/api/format-common.api                |  1 +
 .../msfjarvis/aps/data/passfile/PasswordEntry.kt   | 24 +++++++++++++---------
 .../main/kotlin/dev/msfjarvis/aps/util/totp/Otp.kt | 14 +++++++++----
 .../dev/msfjarvis/aps/util/totp/TotpFinder.kt      |  3 +++
 4 files changed, 28 insertions(+), 14 deletions(-)

(limited to 'format-common')

diff --git a/format-common/api/format-common.api b/format-common/api/format-common.api
index 09ff55a2..6d2377f9 100644
--- a/format-common/api/format-common.api
+++ b/format-common/api/format-common.api
@@ -21,6 +21,7 @@ public abstract interface class dev/msfjarvis/aps/util/totp/TotpFinder {
 	public static final field Companion Ldev/msfjarvis/aps/util/totp/TotpFinder$Companion;
 	public abstract fun findAlgorithm (Ljava/lang/String;)Ljava/lang/String;
 	public abstract fun findDigits (Ljava/lang/String;)Ljava/lang/String;
+	public abstract fun findIssuer (Ljava/lang/String;)Ljava/lang/String;
 	public abstract fun findPeriod (Ljava/lang/String;)J
 	public abstract fun findSecret (Ljava/lang/String;)Ljava/lang/String;
 }
diff --git a/format-common/src/main/kotlin/dev/msfjarvis/aps/data/passfile/PasswordEntry.kt b/format-common/src/main/kotlin/dev/msfjarvis/aps/data/passfile/PasswordEntry.kt
index 4fc9667d..0e8f6d2e 100644
--- a/format-common/src/main/kotlin/dev/msfjarvis/aps/data/passfile/PasswordEntry.kt
+++ b/format-common/src/main/kotlin/dev/msfjarvis/aps/data/passfile/PasswordEntry.kt
@@ -68,10 +68,7 @@ constructor(
    * and usernames stripped.
    */
   public val extraContentWithoutAuthData: String
-  private val digits: String
   private val totpSecret: String?
-  private val totpPeriod: Long
-  private val totpAlgorithm: String
 
   init {
     val (foundPassword, passContent) = findAndStripPassword(content.split("\n".toRegex()))
@@ -80,17 +77,18 @@ constructor(
     extraContentWithoutAuthData = generateExtraContentWithoutAuthData()
     extraContent = generateExtraContentPairs()
     username = findUsername()
-    digits = totpFinder.findDigits(content)
     totpSecret = totpFinder.findSecret(content)
-    totpPeriod = totpFinder.findPeriod(content)
-    totpAlgorithm = totpFinder.findAlgorithm(content)
     if (totpSecret != null) {
       scope.launch {
-        updateTotp(clock.millis())
+        val digits = totpFinder.findDigits(content)
+        val totpPeriod = totpFinder.findPeriod(content)
+        val totpAlgorithm = totpFinder.findAlgorithm(content)
+        val issuer = totpFinder.findIssuer(content)
         val remainingTime = totpPeriod - (clock.millis() % totpPeriod)
+        updateTotp(clock.millis(), totpPeriod, totpAlgorithm, digits, issuer)
         delay(Duration.seconds(remainingTime))
         repeat(Int.MAX_VALUE) {
-          updateTotp(clock.millis())
+          updateTotp(clock.millis(), totpPeriod, totpAlgorithm, digits, issuer)
           delay(Duration.seconds(totpPeriod))
         }
       }
@@ -186,9 +184,15 @@ constructor(
     return null
   }
 
-  private fun updateTotp(millis: Long) {
+  private fun updateTotp(
+    millis: Long,
+    totpPeriod: Long,
+    totpAlgorithm: String,
+    digits: String,
+    issuer: String?,
+  ) {
     if (totpSecret != null) {
-      Otp.calculateCode(totpSecret, millis / (1000 * totpPeriod), totpAlgorithm, digits)
+      Otp.calculateCode(totpSecret, millis / (1000 * totpPeriod), totpAlgorithm, digits, issuer)
         .mapBoth({ code -> _totp.value = code }, { throwable -> throw throwable })
     }
   }
diff --git a/format-common/src/main/kotlin/dev/msfjarvis/aps/util/totp/Otp.kt b/format-common/src/main/kotlin/dev/msfjarvis/aps/util/totp/Otp.kt
index 65284441..f1f71e00 100644
--- a/format-common/src/main/kotlin/dev/msfjarvis/aps/util/totp/Otp.kt
+++ b/format-common/src/main/kotlin/dev/msfjarvis/aps/util/totp/Otp.kt
@@ -23,8 +23,13 @@ internal object Otp {
     check(STEAM_ALPHABET.size == 26)
   }
 
-  fun calculateCode(secret: String, counter: Long, algorithm: String, digits: String) =
-      runCatching {
+  fun calculateCode(
+    secret: String,
+    counter: Long,
+    algorithm: String,
+    digits: String,
+    issuer: String?,
+  ) = runCatching {
     val algo = "Hmac${algorithm.uppercase(Locale.ROOT)}"
     val decodedSecret = BASE_32.decode(secret)
     val secretKey = SecretKeySpec(decodedSecret, algo)
@@ -40,8 +45,9 @@ internal object Otp {
     code[0] = (0x7f and code[0].toInt()).toByte()
     val codeInt = ByteBuffer.wrap(code).int
     check(codeInt > 0)
-    if (digits == "s") {
-      // Steam
+    // SteamGuard is a horrible OTP implementation that generates non-standard 5 digit OTPs as well
+    // as uses a custom character set.
+    if (digits == "s" || issuer == "Steam") {
       var remainingCodeInt = codeInt
       buildString {
         repeat(5) {
diff --git a/format-common/src/main/kotlin/dev/msfjarvis/aps/util/totp/TotpFinder.kt b/format-common/src/main/kotlin/dev/msfjarvis/aps/util/totp/TotpFinder.kt
index 64f22065..1cb7de97 100644
--- a/format-common/src/main/kotlin/dev/msfjarvis/aps/util/totp/TotpFinder.kt
+++ b/format-common/src/main/kotlin/dev/msfjarvis/aps/util/totp/TotpFinder.kt
@@ -20,6 +20,9 @@ public interface TotpFinder {
   /** Get the algorithm for the TOTP secret. */
   public fun findAlgorithm(content: String): String
 
+  /** Get the issuer for the TOTP secret, if any. */
+  public fun findIssuer(content: String): String?
+
   public companion object {
     public val TOTP_FIELDS: Array<String> = arrayOf("otpauth://totp", "totp:")
   }
-- 
cgit v1.2.3