thodg/totp-rs

diff --git a/Cargo.toml b/Cargo.toml
index cc05ed1..16f5bad 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "totp-rs"
-version = "0.3.2"
+version = "0.4.0"
 authors = ["Cleo Rebert <cleo.rebert@gmail.com>"]
 edition = "2018"
 readme = "README.md"
diff --git a/README.md b/README.md
index f372113..c2b4d66 100644
--- a/README.md
+++ b/README.md
@@ -7,14 +7,13 @@ This library permits the creation of 2FA authentification tokens per TOTP, the v
 Add it to your `Cargo.toml`:
 ```toml
 [dependencies]
-totp-rs = "~0.3"
+totp-rs = "~0.4"
 ```
 You can then do something like:
 ```Rust
 use std::time::SystemTime;
 use totp_rs::{Algorithm, TOTP};
 
-let username = "example".to_owned();
 let totp = TOTP::new(
     Algorithm::SHA1,
     6,
@@ -25,7 +24,7 @@ let totp = TOTP::new(
 let time = SystemTime::now()
     .duration_since(SystemTime::UNIX_EPOCH).unwrap()
     .as_secs();
-let url = totp.get_url(format!("account:{}", username), "my-org.com".to_owned());
+let url = totp.get_url("user@example.com", "my-org.com");
 println!("{}", url);
 let token = totp.generate(time);
 println!("{}", token);
@@ -36,14 +35,13 @@ println!("{}", token);
 Add it to your `Cargo.toml`:
 ```toml
 [dependencies.totp-rs]
-version = "~0.3"
+version = "~0.4"
 features = ["qr"]
 ```
 You can then do something like:
 ```Rust
 use totp_rs::{Algorithm, TOTP};
 
-let username = "example".to_owned();
 let totp = TOTP::new(
     Algorithm::SHA1,
     6,
@@ -51,6 +49,6 @@ let totp = TOTP::new(
     30,
     "supersecret".to_owned().into_bytes(),
 );
-let code = totp.get_qr(format!("account:{}", username), "my-org.com".to_owned())?;
+let code = totp.get_qr("user@example.com", "my-org.com")?;
 println!("{}", code);
 ```
diff --git a/src/lib.rs b/src/lib.rs
index d2beeff..1c7b54f 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -6,7 +6,6 @@
 //! use std::time::SystemTime;
 //! use totp_rs::{Algorithm, TOTP};
 //! 
-//! let username = "example".to_owned();
 //! let totp = TOTP::new(
 //!     Algorithm::SHA1,
 //!     6,
@@ -17,7 +16,7 @@
 //! let time = SystemTime::now()
 //!     .duration_since(SystemTime::UNIX_EPOCH).unwrap()
 //!     .as_secs();
-//! let url = totp.get_url(format!("account:{}", username), "my-org.com".to_owned());
+//! let url = totp.get_url("user@example.com", "my-org.com");
 //! println!("{}", url);
 //! let token = totp.generate(time);
 //! println!("{}", token);
@@ -26,7 +25,6 @@
 //! ```rust
 //! use totp_rs::{Algorithm, TOTP};
 //!
-//! let username = "example".to_owned();
 //! let totp = TOTP::new(
 //!     Algorithm::SHA1,
 //!     6,
@@ -34,7 +32,7 @@
 //!     30,
 //!     "supersecret".to_owned().into_bytes(),
 //! );
-//! let code = totp.get_qr(format!("account:{}", username), "my-org.com".to_owned())?;
+//! let code = totp.get_qr("user@example.com", "my-org.com").unwrap();
 //! println!("{}", code);
 //! ```
 
@@ -127,11 +125,10 @@ impl TOTP {
     }
 
     /// Will check if token is valid by current time, accounting [skew](struct.TOTP.html#structfield.skew)
-    pub fn check(&self, token: String, time: u64) -> bool {
+    pub fn check(&self, token: &str, time: u64) -> bool {
         let basestep = time / self.step - (self.skew as u64);
         for i in 0..self.skew * 2 + 1 {
             let step_time = (basestep + (i as u64)) * (self.step as u64);
-            println!("{}", self.generate(step_time));
             if self.generate(step_time) == token {
                 return true;
             }
@@ -139,20 +136,24 @@ impl TOTP {
         false
     }
 
+    /// Will return the base32 representation of the secret, which might be useful when users want to manually add the secret to their authenticator
+    pub fn get_secret_base32(&self) -> String {
+        base32::encode(base32::Alphabet::RFC4648 { padding: false }, &self.secret)
+    }
+
     /// Will generate a standard URL used to automatically add TOTP auths. Usually used with qr codes
-    pub fn get_url(&self, label: String, issuer: String) -> String {
-        let algorithm: String;
-        match self.algorithm {
-            Algorithm::SHA1 => algorithm = "SHA1".to_owned(),
-            Algorithm::SHA256 => algorithm = "SHA256".to_owned(),
-            Algorithm::SHA512 => algorithm = "SHA512".to_owned(),
-        }
+    pub fn get_url(&self, label: &str, issuer: &str) -> String {
+        let algorithm = match self.algorithm {
+            Algorithm::SHA1 => "SHA1",
+            Algorithm::SHA256 => "SHA256",
+            Algorithm::SHA512 => "SHA512",
+        };
         format!(
             "otpauth://totp/{}?secret={}&issuer={}&digits={}&algorithm={}",
             label,
-            base32::encode(base32::Alphabet::RFC4648 { padding: false }, &self.secret),
+            self.get_secret_base32(),
             issuer,
-            self.digits.to_string(),
+            self.digits,
             algorithm,
         )
     }
@@ -167,8 +168,8 @@ impl TOTP {
     #[cfg(feature = "qr")]
     pub fn get_qr(
         &self,
-        label: String,
-        issuer: String,
+        label: &str,
+        issuer: &str,
     ) -> Result<String, Box<dyn std::error::Error>> {
         let url = self.get_url(label, issuer);
         let code = QrCode::new(&url)?;
@@ -184,3 +185,64 @@ impl TOTP {
         Ok(base64::encode(vec))
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn url_for_secret_matches() {
+        let totp = TOTP::new(Algorithm::SHA1, 6, 1, 1, String::from("TestSecret").into_bytes());
+        let url = totp.get_url("test_url", "totp-rs");
+        assert_eq!(url.as_str(), "otpauth://totp/test_url?secret=KRSXG5CTMVRXEZLU&issuer=totp-rs&digits=6&algorithm=SHA1");
+    }
+
+    #[test]
+    fn returns_base32() {
+        let totp = TOTP::new(Algorithm::SHA1, 6, 1, 1, String::from("TestSecret").into_bytes());
+        assert_eq!(totp.get_secret_base32().as_str(), "KRSXG5CTMVRXEZLU");
+    }
+
+    #[test]
+    fn generates_token() {
+        let totp = TOTP::new(Algorithm::SHA1, 6, 1, 1, String::from("TestSecret").into_bytes());
+        assert_eq!(totp.generate(1000).as_str(), "718996");
+    }
+
+    #[test]
+    fn generates_token_sha256() {
+        let totp = TOTP::new(Algorithm::SHA256, 6, 1, 1, String::from("TestSecret").into_bytes());
+        assert_eq!(totp.generate(1000).as_str(), "423657");
+    }
+
+    #[test]
+    fn generates_token_sha512() {
+        let totp = TOTP::new(Algorithm::SHA512, 6, 1, 1, String::from("TestSecret").into_bytes());
+        assert_eq!(totp.generate(1000).as_str(), "416767");
+    }
+
+    #[test]
+    fn checks_token() {
+        let totp = TOTP::new(Algorithm::SHA1, 6, 1, 1, String::from("TestSecret").into_bytes());
+        assert!(totp.check("718996", 1000));
+    }
+
+    #[test]
+    fn checks_token_with_skew() {
+        let totp = TOTP::new(Algorithm::SHA1, 6, 1, 1, String::from("TestSecret").into_bytes());
+        assert!(totp.check("527544", 2000) && totp.check("712039", 2000) && totp.check("714250", 2000));
+    }
+
+    #[test]
+    #[cfg(feature = "qr")]
+    fn generates_qr() {
+        use sha1::{Sha1, Digest};
+
+        let totp = TOTP::new(Algorithm::SHA1, 6, 1, 1, String::from("TestSecret").into_bytes());
+        let qr = totp.get_qr("test_url", "totp-rs").unwrap();
+
+        // Create hash from image
+        let hash_digest = Sha1::digest(qr.as_bytes());
+        assert_eq!(format!("{:x}", hash_digest).as_str(), "3abc0127e7a2b1013fb25c97ef14422c1fe9e878");
+    }
+}
\ No newline at end of file