1 files changed, 58 insertions, 6 deletions

diff --git a/pw_rand.c b/pw_rand.c
index 5c9cd4f..127687d 100644
--- a/pw_rand.c
+++ b/pw_rand.c
@@ -9,20 +9,72 @@
  */
 
 #include <string.h>
+#include <stdio.h>
+#include <stdlib.h>
 #include "pwgen.h"
 
-#define PW_CHARS "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
+const char *pw_digits = "0123456789";
+const char *pw_uppers = "ABCDEFGHIJKLMNOPQRSTUVWXYZ";
+const char *pw_lowers = "abcdefghijklmnopqrstuvwxyz";
+const char *pw_symbols = "!\"#$%&'()*+,-./:;<=>?@[\\]^_`{|}~";
+const char *pw_ambiguous = "B8G6I1l0OQDS5Z2";
 
 void pw_rand(char *buf, int size, int pw_flags)
 {
-	char	ch, *chars = PW_CHARS;
-	int	i = 0, len;
+	char		ch, *chars, *wchars;
+	int		i, len, feature_flags;
 
+	len = 0;
+	if (pw_flags & PW_DIGITS) {
+		len += strlen(pw_digits);
+	}
+	if (pw_flags & PW_UPPERS) {
+		len += strlen(pw_uppers);
+	}
+	len += strlen(pw_lowers);
+	if (pw_flags & PW_SYMBOLS) {
+		len += strlen(pw_symbols);
+	}
+        chars = malloc(len+1);
+        if (!chars) {
+		fprintf(stderr, "Couldn't malloc pw_rand buffer.\n");
+		exit(1);
+	}
+	wchars = chars;
+	if (pw_flags & PW_DIGITS) {
+		strcpy(wchars, pw_digits);
+		wchars += strlen(pw_digits);
+	}
+	if (pw_flags & PW_UPPERS) {
+		strcpy(wchars, pw_uppers);
+		wchars += strlen(pw_uppers);
+	}
+	strcpy(wchars, pw_lowers);
+	wchars += strlen(pw_lowers);
+	if (pw_flags & PW_SYMBOLS) {
+		strcpy(wchars, pw_symbols);
+	}
+try_again:
 	len = strlen(chars);
-	while (i < size) {
-		ch = chars[pw_random_number(len)];
-		buf[i++] = ch;
+	feature_flags = pw_flags;
+	for (i = 0; i < size; i++) {
+		ch = chars[pw_number(len)];
+		if (pw_flags & PW_AMBIGUOUS) {
+			if (strchr(pw_ambiguous,ch)) {
+				continue;
+			}
+		}
+		buf[i] = ch;
+		if (strchr(pw_digits, ch))
+			feature_flags &= ~PW_DIGITS;
+		if (strchr(pw_uppers, ch))
+			feature_flags &= ~PW_UPPERS;
+		if (strchr(pw_symbols, ch))
+			feature_flags &= ~PW_SYMBOLS;
 	}
+	if (feature_flags & (PW_UPPERS | PW_DIGITS | PW_SYMBOLS))
+		goto try_again;
 	buf[size] = 0;
+	free(chars);
 	return;
 }