Commit d5b94dc0dfc0dab29c7844d968625fa8607e4751

Cléo Rebert 2023-01-06T10:10:01

Merge pull request #47 from timvisee/feature-steam Add Steam algorithm behind 'steam' feature

diff --git a/Cargo.toml b/Cargo.toml
index ddef606..605bebf 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -21,6 +21,7 @@ otpauth = ["url", "urlencoding"]
 qr = ["qrcodegen", "image", "base64", "otpauth"]
 serde_support = ["serde"]
 gen_secret = ["rand"]
+steam = []
 
 [dependencies]
 serde = { version = "1.0", features = ["derive"], optional = true }
diff --git a/README.md b/README.md
index b1393c4..ab8663f 100644
--- a/README.md
+++ b/README.md
@@ -19,6 +19,8 @@ With optional feature "serde_support", library-defined types `TOTP` and `Algorit
 With optional feature "gen_secret", a secret will be generated for you to store in database.
 ### zeroize
 Securely zero secret information when the TOTP struct is dropped.
+### steam
+Add support for Steam TOTP tokens.
 
 
 # Examples
diff --git a/src/custom_providers.rs b/src/custom_providers.rs
new file mode 100644
index 0000000..8f9487b
--- /dev/null
+++ b/src/custom_providers.rs
@@ -0,0 +1,43 @@
+#[cfg(feature = "steam")]
+use crate::{Algorithm, TOTP};
+
+#[cfg(feature = "steam")]
+impl TOTP {
+    #[cfg(feature = "otpauth")]
+    /// Will create a new instance of TOTP using the Steam algorithm with given parameters. See [the doc](struct.TOTP.html#fields) for reference as to how to choose those values
+    ///
+    /// # Description
+    /// * `secret`: expect a non-encoded value, to pass in base32 string use `Secret::Encoded(String)`
+    ///
+    /// ```rust
+    /// use totp_rs::{Secret, TOTP};
+    /// let secret = Secret::Encoded("ABCDEFGHIJKLMNOPQRSTUVWXYZ234567".to_string());
+    /// let totp = TOTP::new_steam(secret.to_bytes().unwrap(), Some("username".to_string()));
+    /// ```
+    pub fn new_steam(secret: Vec<u8>, account_name: String) -> TOTP {
+        Self::new_unchecked(
+            Algorithm::Steam,
+            5,
+            1,
+            30,
+            secret,
+            Some("Steam".into()),
+            account_name,
+        )
+    }
+
+    #[cfg(not(feature = "otpauth"))]
+    /// Will create a new instance of TOTP using the Steam algorithm with given parameters. See [the doc](struct.TOTP.html#fields) for reference as to how to choose those values
+    ///
+    /// # Description
+    /// * `secret`: expect a non-encoded value, to pass in base32 string use `Secret::Encoded(String)`
+    ///
+    /// ```rust
+    /// use totp_rs::{Secret, TOTP};
+    /// let secret = Secret::Encoded("ABCDEFGHIJKLMNOPQRSTUVWXYZ234567".to_string());
+    /// let totp = TOTP::new_steam(secret.to_bytes().unwrap());
+    /// ```
+    pub fn new_steam(secret: Vec<u8>) -> TOTP {
+        Self::new_unchecked(Algorithm::Steam, 5, 1, 30, secret)
+    }
+}
diff --git a/src/lib.rs b/src/lib.rs
index 55ba5d5..3b08937 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -47,6 +47,7 @@
 //! # }
 //! ```
 
+mod custom_providers;
 mod rfc;
 mod secret;
 mod url_error;
@@ -75,6 +76,10 @@ type HmacSha1 = hmac::Hmac<sha1::Sha1>;
 type HmacSha256 = hmac::Hmac<sha2::Sha256>;
 type HmacSha512 = hmac::Hmac<sha2::Sha512>;
 
+/// Alphabet for Steam tokens.
+#[cfg(feature = "steam")]
+const STEAM_CHARS: &str = "23456789BCDFGHJKMNPQRTVWXY";
+
 /// Algorithm enum holds the three standards algorithms for TOTP as per the [reference implementation](https://tools.ietf.org/html/rfc6238#appendix-A)
 #[derive(Debug, Copy, Clone, Eq, PartialEq)]
 #[cfg_attr(feature = "serde_support", derive(Serialize, Deserialize))]
@@ -82,6 +87,8 @@ pub enum Algorithm {
     SHA1,
     SHA256,
     SHA512,
+    #[cfg(feature = "steam")]
+    Steam,
 }
 
 impl std::default::Default for Algorithm {
@@ -96,6 +103,8 @@ impl fmt::Display for Algorithm {
             Algorithm::SHA1 => f.write_str("SHA1"),
             Algorithm::SHA256 => f.write_str("SHA256"),
             Algorithm::SHA512 => f.write_str("SHA512"),
+            #[cfg(feature = "steam")]
+            Algorithm::Steam => f.write_str("SHA1"),
         }
     }
 }
@@ -114,6 +123,8 @@ impl Algorithm {
             Algorithm::SHA1 => Algorithm::hash(HmacSha1::new_from_slice(key).unwrap(), data),
             Algorithm::SHA256 => Algorithm::hash(HmacSha256::new_from_slice(key).unwrap(), data),
             Algorithm::SHA512 => Algorithm::hash(HmacSha512::new_from_slice(key).unwrap(), data),
+            #[cfg(feature = "steam")]
+            Algorithm::Steam => Algorithm::hash(HmacSha1::new_from_slice(key).unwrap(), data),
         }
     }
 }
@@ -384,13 +395,28 @@ impl TOTP {
     pub fn generate(&self, time: u64) -> String {
         let result: &[u8] = &self.sign(time);
         let offset = (result.last().unwrap() & 15) as usize;
-        let result =
+        #[allow(unused_mut)]
+        let mut result =
             u32::from_be_bytes(result[offset..offset + 4].try_into().unwrap()) & 0x7fff_ffff;
-        format!(
-            "{1:00$}",
-            self.digits,
-            result % 10_u32.pow(self.digits as u32)
-        )
+
+        match self.algorithm {
+            Algorithm::SHA1 | Algorithm::SHA256 | Algorithm::SHA512 => format!(
+                "{1:00$}",
+                self.digits,
+                result % 10_u32.pow(self.digits as u32)
+            ),
+            #[cfg(feature = "steam")]
+            Algorithm::Steam => (0..self.digits)
+                .map(|_| {
+                    let c = STEAM_CHARS
+                        .chars()
+                        .nth(result as usize % STEAM_CHARS.len())
+                        .unwrap();
+                    result /= STEAM_CHARS.len() as u32;
+                    c
+                })
+                .collect(),
+        }
     }
 
     /// Returns the timestamp of the first second for the next step
@@ -476,14 +502,6 @@ impl TOTP {
     fn parts_from_url<S: AsRef<str>>(
         url: S,
     ) -> Result<(Algorithm, usize, u8, u64, Vec<u8>, Option<String>, String), TotpUrlError> {
-        let url = Url::parse(url.as_ref()).map_err(TotpUrlError::Url)?;
-        if url.scheme() != "otpauth" {
-            return Err(TotpUrlError::Scheme(url.scheme().to_string()));
-        }
-        if url.host() != Some(Host::Domain("totp")) {
-            return Err(TotpUrlError::Host(url.host().unwrap().to_string()));
-        }
-
         let mut algorithm = Algorithm::SHA1;
         let mut digits = 6;
         let mut step = 30;
@@ -491,6 +509,22 @@ impl TOTP {
         let mut issuer: Option<String> = None;
         let mut account_name: String;
 
+        let url = Url::parse(url.as_ref()).map_err(TotpUrlError::Url)?;
+        if url.scheme() != "otpauth" {
+            return Err(TotpUrlError::Scheme(url.scheme().to_string()));
+        }
+        match url.host() {
+            Some(Host::Domain("totp")) => {}
+            #[cfg(feature = "steam")]
+            Some(Host::Domain("steam")) => {
+                algorithm = Algorithm::Steam;
+                digits = 5;
+            }
+            _ => {
+                return Err(TotpUrlError::Host(url.host().unwrap().to_string()));
+            }
+        }
+
         let path = url.path().trim_start_matches('/');
         if path.contains(':') {
             let parts = path.split_once(':').unwrap();
@@ -510,6 +544,10 @@ impl TOTP {
 
         for (key, value) in url.query_pairs() {
             match key.as_ref() {
+                #[cfg(feature = "steam")]
+                "algorithm" if algorithm == Algorithm::Steam => {
+                    // Do not change used algorithm if this is Steam
+                }
                 "algorithm" => {
                     algorithm = match value.as_ref() {
                         "SHA1" => Algorithm::SHA1,
@@ -535,6 +573,12 @@ impl TOTP {
                     )
                     .ok_or_else(|| TotpUrlError::Secret(value.to_string()))?;
                 }
+                #[cfg(feature = "steam")]
+                "issuer" if value.to_lowercase() == "steam" => {
+                    algorithm = Algorithm::Steam;
+                    digits = 5;
+                    issuer = Some(value.into());
+                }
                 "issuer" => {
                     let param_issuer = value
                         .parse::<String>()
@@ -564,6 +608,12 @@ impl TOTP {
     /// Secret will be base 32'd without padding, as per RFC.
     #[cfg(feature = "otpauth")]
     pub fn get_url(&self) -> String {
+        #[allow(unused_mut)]
+        let mut host = "totp";
+        #[cfg(feature = "steam")]
+        if self.algorithm == Algorithm::Steam {
+            host = "steam";
+        }
         let account_name: String = urlencoding::encode(self.account_name.as_str()).to_string();
         let mut label: String = format!("{}?", account_name);
         if self.issuer.is_some() {
@@ -573,7 +623,8 @@ impl TOTP {
         }
 
         format!(
-            "otpauth://totp/{}secret={}&digits={}&algorithm={}",
+            "otpauth://{}/{}secret={}&digits={}&algorithm={}",
+            host,
             label,
             self.get_secret_base32(),
             self.digits,