stupid_2fa/
lib.rs

1use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
2use hmac::{Hmac, Mac};
3use sha2::Sha256;
4use std::env;
5use std::time::{SystemTime, UNIX_EPOCH};
6use uuid::Uuid;
7
8type HmacSha256 = Hmac<Sha256>;
9
10const SECRET_KEY_ENV_NAME: &'static str = "STUPID_2FA_PRIVATE_KEY";
11
12fn get_secret_key() -> String {
13    env::var(SECRET_KEY_ENV_NAME).expect("Secret key is not defined")
14}
15
16pub fn generate_client_code() -> String {
17    let device_id = Uuid::new_v4().to_string().replace("-", "");
18    let timestamp = SystemTime::now()
19        .duration_since(UNIX_EPOCH)
20        .unwrap()
21        .as_secs();
22    format!("{}-{}", device_id, timestamp)
23}
24
25pub fn generate_unlock_code(lock_code: &str, subscription_days: i64) -> String {
26    let secret_key = get_secret_key();
27    let message = format!("{}-{}", lock_code, subscription_days);
28    let mut mac =
29        HmacSha256::new_from_slice(secret_key.as_bytes()).expect("HMAC can take key of any size");
30    mac.update(message.as_bytes());
31    let result = mac.finalize().into_bytes();
32    URL_SAFE_NO_PAD.encode(result)
33}
34
35pub fn validate_unlock_code(lock_code: &str, unlock_code: &str, subscription_days: i64) -> bool {
36    let secret_key = get_secret_key();
37    let message = format!("{}-{}", lock_code, subscription_days);
38    let mut mac =
39        HmacSha256::new_from_slice(secret_key.as_bytes()).expect("HMAC can take key of any size");
40    mac.update(message.as_bytes());
41    let expected_result = mac.finalize().into_bytes();
42    let expected_unlock_code = URL_SAFE_NO_PAD.encode(expected_result);
43
44    expected_unlock_code == unlock_code
45}
46
47#[cfg(test)]
48mod tests {
49    use super::*;
50    use chrono::{Duration, Utc};
51    use std::env::set_var;
52    use std::sync::Once;
53
54    static INIT: Once = Once::new();
55
56    pub fn initialize() {
57        INIT.call_once(|| {
58            set_var(SECRET_KEY_ENV_NAME, "VALUE");
59        });
60    }
61
62    #[test]
63    fn test_generate_lock_code() {
64        initialize();
65        let lock_code = generate_client_code();
66        assert!(lock_code.contains('-'));
67        let parts: Vec<&str> = lock_code.split('-').collect();
68        assert_eq!(parts.len(), 2);
69        assert!(Uuid::parse_str(parts[0]).is_ok());
70        assert!(parts[1].parse::<u64>().is_ok());
71    }
72
73    #[test]
74    fn test_generate_unlock_code() {
75        initialize();
76        let lock_code = generate_client_code();
77        let subscription_days = 30;
78        let unlock_code = generate_unlock_code(&lock_code, subscription_days);
79        assert!(!unlock_code.is_empty());
80    }
81
82    #[test]
83    fn test_validate_unlock_code_valid() {
84        initialize();
85        let lock_code = generate_client_code();
86        let subscription_days = 30;
87        let unlock_code = generate_unlock_code(&lock_code, subscription_days);
88        assert!(validate_unlock_code(
89            &lock_code,
90            &unlock_code,
91            subscription_days,
92        ));
93    }
94
95    #[test]
96    fn test_validate_unlock_code_invalid() {
97        initialize();
98        let lock_code = generate_client_code();
99        let invalid_unlock_code = "invalid_unlock_code";
100        let subscription_days = 30;
101        assert!(!validate_unlock_code(
102            &lock_code,
103            invalid_unlock_code,
104            subscription_days,
105        ));
106    }
107
108    #[test]
109    fn test_validate_unlock_code_expired() {
110        initialize();
111        let device_id = Uuid::new_v4().to_string();
112        let timestamp = (Utc::now() - Duration::days(31)).timestamp();
113        let lock_code = format!("{}-{}", device_id, timestamp);
114        let subscription_days = 30;
115        let unlock_code = generate_unlock_code(&lock_code, subscription_days);
116        assert!(validate_unlock_code(
117            &lock_code,
118            &unlock_code,
119            subscription_days,
120        ));
121    }
122
123    #[test]
124    fn test_validate_unlock_code_not_expired() {
125        initialize();
126        let device_id = Uuid::new_v4().to_string();
127        let timestamp = Utc::now().timestamp();
128        let lock_code = format!("{}-{}", device_id, timestamp);
129        let subscription_days = 30;
130        let unlock_code = generate_unlock_code(&lock_code, subscription_days);
131        assert!(validate_unlock_code(
132            &lock_code,
133            &unlock_code,
134            subscription_days,
135        ));
136    }
137}
stupid_2fa/lib.rs

stupid_2fa/
lib.rs
