1 files changed, 27 insertions, 7 deletions

diff --git a/src/lib.rs b/src/lib.rs
index cfa078a..a4fa227 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -14,18 +14,38 @@ extern {
     fn alert(s: &str);
 }
 
+
 #[wasm_bindgen]
-pub fn generate(len: usize) -> String {
-    const CHARSET: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZ\
-                            abcdefghijklmnopqrstuvwxyz\
-                            1234567890\
-                            !@#$%^&*()+{}[]";
+pub fn generate(len: usize, lower: bool, upper: bool, number: bool, special: bool) -> String {
+    let mut upper_charset: Vec<char> = String::from("ABCDEFGHIJKLMNOPQRSTUVWXYZ").chars().collect();
+    let mut lower_charset: Vec<char> = String::from("abcdefghijklmnopqrstuvwxyz").chars().collect();
+    let mut number_charset: Vec<char> = String::from("1234567890").chars().collect();
+    let mut special_charset: Vec<char> = String::from("!@#$%^&*()+{}[]").chars().collect();
+    let mut result: Vec<char> = Vec::new();
+
     let mut rng = rand::thread_rng();
 
+    if lower {
+        result.append(&mut lower_charset);
+    }
+
+    if upper {
+        result.append(&mut upper_charset);
+    }
+
+    if number {
+        result.append(&mut number_charset);
+    }
+
+    if special {
+        result.append(&mut special_charset);
+    }
+
+
     (0..len)
         .map(|_| {
-            let idx = rng.gen_range(0..CHARSET.len());
-            CHARSET[idx] as char
+            let idx = rng.gen_range(0..result.len());
+            result[idx] as char
         })
         .collect()
 }