Add :algorithm option (#41)

Author: LoicBersier

lib/nimble_totp.ex

@@ -65,8 +65,8 @@ defmodule NimbleTOTP do
       uri |> EQRCode.encode() |> EQRCode.svg()
       #=> "<?xml version=\\"1.0\\" standalone=\\"yes\\"?>\\n<svg version=\\"1.1\\" ...
        
-      uri = NimbleTOTP.otpauth_uri("Acme", "alice", secret, digits: 8)
-      #=> "otpauth://totp/Acme:alice?secret=MFRGGZA&issuer=Acme"
+      uri = NimbleTOTP.otpauth_uri("Acme", "alice", secret, digits: 8, algorithm: :sha512)
+      #=> "otpauth://totp/Acme:alice?secret=MFRGGZA&issuer=Acme&digits=8&algorithm=sha512"
       uri |> EQRCode.encode() |> EQRCode.svg()
       #=> "<?xml version=\\"1.0\\" standalone=\\"yes\\"?>\\n<svg version=\\"1.1\\" ...
 
@@ -138,6 +138,7 @@ defmodule NimbleTOTP do
   import Bitwise
   @default_digits 6
   @default_totp_period 30
+  @default_algorithm :sha
 
   @typedoc "Unix time in seconds, `t:DateTime.t/0` or `t:NaiveDateTime.t/0`."
   @type time() :: DateTime.t() | NaiveDateTime.t() | integer()
@@ -224,6 +225,8 @@ defmodule NimbleTOTP do
       If this option is given to `verification_code/2`, it must also be given to `valid?/3`.
     * `:digits` - The desired length of the totp. Default is 6.
       If this option is given to `verification_code/2`, it must also be given to `valid?/3` and `otpauth_uri/3`/`otpauth_uri/4`.
+    * `:algorithm` - The algorithm to use for the totp. Default is :sha.
+      If this option is given to `verification_code/2`, it must also be given to `valid?/3` and `otpauth_uri/3`/`otpauth_uri/4`.
 
   ## Examples
 
@@ -237,31 +240,37 @@ defmodule NimbleTOTP do
     time = opts |> Keyword.get_lazy(:time, fn -> System.os_time(:second) end) |> to_unix()
     period = Keyword.get(opts, :period, @default_totp_period)
     digits = Keyword.get(opts, :digits, @default_digits)
+    algorithm = Keyword.get(opts, :algorithm, @default_algorithm)
 
     digits not in 6..10 && raise ArgumentError, "digits must be between 6 and 10"
 
-    verification_code(secret, time, period, digits)
+    algorithm not in [:sha, :sha256, :sha512] &&
+      raise ArgumentError, "algorithm must be one of :sha, :sha256, :sha512"
+
+    verification_code(secret, time, period, digits, algorithm)
   end
 
-  @spec verification_code(binary(), integer(), pos_integer(), integer()) :: binary()
-  defp verification_code(secret, time, period, digits) do
+  @spec verification_code(binary(), integer(), pos_integer(), integer(), atom()) :: binary()
+  defp verification_code(secret, time, period, digits, algorithm) do
     secret
-    |> hmac(time, period)
+    |> hmac(time, period, algorithm)
     |> hmac_truncate()
     |> rem(Integer.pow(10, digits))
     |> to_string()
     |> String.pad_leading(digits, "0")
   end
 
-  defp hmac(secret, time, period) do
+  defp hmac(secret, time, period, algorithm) do
     moving_factor = <<Integer.floor_div(time, period)::64>>
-    hmac_sha(secret, moving_factor)
+    hmac_sha(secret, moving_factor, algorithm)
   end
 
-  defp hmac_sha(key, data), do: :crypto.mac(:hmac, :sha, key, data)
+  defp hmac_sha(key, data, algorithm), do: :crypto.mac(:hmac, algorithm, key, data)
 
   defp hmac_truncate(hmac) do
-    <<_::19-binary, _::4, offset::4>> = hmac
+    key_length = byte_size(hmac) - 1
+
+    <<_::size(key_length)-binary, _::4, offset::4>> = hmac
     <<_::size(offset)-binary, p::4-binary, _::binary>> = hmac
     <<_::1, bits::31>> = p
     bits
@@ -284,6 +293,9 @@ defmodule NimbleTOTP do
     * `:digits` - The desired length of the totp. Default is 6.
       If this option is given to `verification_code/2`, it must also be given to `valid?/3` and `otpauth_uri/3`/`otpauth_uri/4`.
 
+    * `:algorithm` - The algorithm to use for the totp. Default is :sha.
+      If this option is given to `verification_code/2`, it must also be given to `valid?/3` and `otpauth_uri/3`/`otpauth_uri/4`.
+
   ## Preventing TOTP code reuse
 
   The `:since` option can be used to prevent TOTP codes from being reused. When set
@@ -317,10 +329,14 @@ defmodule NimbleTOTP do
     time = opts |> Keyword.get(:time, System.os_time(:second)) |> to_unix()
     period = Keyword.get(opts, :period, @default_totp_period)
     digits = Keyword.get(opts, :digits, @default_digits)
+    algorithm = Keyword.get(opts, :algorithm, @default_algorithm)
 
     digits not in 6..10 && raise ArgumentError, "digits must be between 6 and 10"
 
-    code = verification_code(secret, time, period, digits)
+    algorithm not in [:sha, :sha256, :sha512] &&
+      raise ArgumentError, "algorithm must be one of :sha, :sha256, :sha512"
+
+    code = verification_code(secret, time, period, digits, algorithm)
 
     byte_size(code) == byte_size(otp) and validate_digits(code, otp) == 0 and
       not reused?(time, period, opts)

test/nimble_totp_test.exs

@@ -57,6 +57,21 @@ defmodule NimbleTOTPTest do
              """
     end
 
+    test "Generate the uri with extra params (deprecated, with algorithm option)" do
+      secret = Base.decode32!("PTEPUGZ7DUWTBGMW4WLKB6U63MGKKMCA")
+      app = "Bytepack App"
+
+      assert NimbleTOTP.otpauth_uri("#{app}:user@test.com", secret,
+               issuer: app,
+               algorithm: :sha512
+             ) == """
+             otpauth://totp/Bytepack%20App:user@test.com?\
+             secret=PTEPUGZ7DUWTBGMW4WLKB6U63MGKKMCA&\
+             issuer=Bytepack%20App&\
+             algorithm=sha512\
+             """
+    end
+
     test "raises error if issuer contains a colon (otpauth_uri/4)" do
       secret = Base.decode32!("PTEPUGZ7DUWTBGMW4WLKB6U63MGKKMCA")
 
@@ -137,6 +152,22 @@ defmodule NimbleTOTPTest do
 
       assert code1 == code2
     end
+
+    test "generate different codes with different algorithm" do
+      secret = NimbleTOTP.secret()
+      time1 = ~N[2020-04-08 17:49:59Z]
+      time2 = ~N[2020-04-08 17:50:00Z]
+      time3 = ~N[2020-04-08 17:50:30Z]
+
+      code1 = NimbleTOTP.verification_code(secret, time: time1, algorithm: :sha)
+      assert code1 == NimbleTOTP.verification_code(secret, time: time1, algorithm: :sha)
+
+      code2 = NimbleTOTP.verification_code(secret, time: time2, algorithm: :sha256)
+      assert code2 == NimbleTOTP.verification_code(secret, time: time2, algorithm: :sha256)
+
+      code3 = NimbleTOTP.verification_code(secret, time: time3, algorithm: :sha512)
+      assert code3 == NimbleTOTP.verification_code(secret, time: time3, algorithm: :sha512)
+    end
   end
 
   describe "valid?/2" do
@@ -233,6 +264,15 @@ defmodule NimbleTOTPTest do
       end
     end
 
+    test "rejects if given an invalid algorithm" do
+      time = System.os_time(:second)
+      secret = NimbleTOTP.secret()
+
+      assert_raise ArgumentError, "algorithm must be one of :sha, :sha256, :sha512", fn ->
+        NimbleTOTP.verification_code(secret, time: time, algorithm: :invalid)
+      end
+    end
+
     test "returns false if the otp is not a binary" do
       secret = NimbleTOTP.secret()
       refute NimbleTOTP.valid?(secret, 123_456, time: System.os_time(:second))