ulif · alanhoyle · Apr 16, 2022 · Apr 16, 2022
diff --git a/diceware/__init__.py b/diceware/__init__.py
@@ -30,7 +30,8 @@
 __version__ = pkg_resources.get_distribution('diceware').version
 
 #: Special chars inserted on demand
-SPECIAL_CHARS = r"~!#$%^&*()-=+[]\{}:;" + r'"' + r"'<>?/0123456789"
+SPECIAL_CHARS = r"~!#$%^&*()-=+[]\{}:;" + r'"' + r"'<>?/"
+DIGIT_CHARS = r"0123456789"
 
 
 GPL_TEXT = (
@@ -104,6 +105,9 @@ def handle_options(args):
     parser.add_argument(
         '-s', '--specials', default=0, type=int, metavar='NUM',
         help="Insert NUM special chars into generated word.")
+    parser.add_argument(
+        '-D', '--digits', default=0, type=int, metavar='NUM',
+        help="Insert NUM digit chars into generated word.")
     parser.add_argument(
         '-d', '--delimiter', default='',
         help="Separate words by DELIMITER. Empty string by default.")
@@ -161,6 +165,21 @@ def insert_special_char(word, specials=SPECIAL_CHARS, rnd=None):
     return ''.join(char_list)
 
 
+def insert_digit_char(word, digits=DIGIT_CHARS, rnd=None):
+    """Insert a char out of `digits` into `word`.
+
+    `rnd`, if passed in, will be used as a (pseudo) random number
+    generator. We use `.choice()` only.
+
+    Returns the modified word.
+    """
+    if rnd is None:
+        rnd = SystemRandom()
+    char_list = list(word)
+    char_list[rnd.choice(range(len(char_list)))] = rnd.choice(digits)
+    return ''.join(char_list)
+
+
 def get_passphrase(options=None):
     """Get a diceware passphrase.
 
@@ -198,6 +217,8 @@ def get_passphrase(options=None):
     result = options.delimiter.join(words)
     for _ in range(options.specials):
         result = insert_special_char(result, rnd=rnd)
+    for _ in range(options.digits):
+        result = insert_digit_char(result, rnd=rnd)
     return result
 
 

diff --git a/tests/test_diceware.py b/tests/test_diceware.py
@@ -7,7 +7,8 @@
 from io import StringIO
 from errno import EISDIR
 from diceware import (
-    get_wordlists_dir, SPECIAL_CHARS, insert_special_char, get_passphrase,
+    get_wordlists_dir, SPECIAL_CHARS, insert_special_char,
+    DIGIT_CHARS, insert_digit_char, get_passphrase,
     handle_options, main, __version__, print_version, get_random_sources,
     get_wordlist_names
     )
@@ -206,9 +207,39 @@ def test_insert_special_char_defaults(self):
         for x in range(100):
             assert insert_special_char('foo') in expected_matrix
 
+    def test_insert_digit_char(self):
+        # we can insert special chars in words.
+        fake_rnd = FakeRandom()
+        result1 = insert_digit_char('foo', digits='123', rnd=fake_rnd)
+        assert result1 == '1oo'
+        fake_rnd.nums_to_draw = [1, 1]
+        result2 = insert_digit_char('foo', digits='123', rnd=fake_rnd)
+        assert result2 == 'f2o'
+        fake_rnd.nums_to_draw = [2, 2]
+        result3 = insert_digit_char('foo', digits='123', rnd=fake_rnd)
+        assert result3 == 'fo3'
+        fake_rnd.nums_to_draw = [0, 0]
+        result4 = insert_digit_char('foo', rnd=fake_rnd)
+        assert result4 == '0oo'
+
+    def test_insert_digit_char_defaults(self):
+        # defaults are respected
+        expected_matrix = []
+        for i in range(3):
+            for j in range(len(DIGIT_CHARS)):
+                word = list('foo')
+                word[i] = DIGIT_CHARS[j]
+                expected_matrix.append(''.join(word))
+        for x in range(100):
+            assert insert_digit_char('foo') in expected_matrix
+
     def test_special_chars_do_not_quote(self):
         # backslashes in SPECIAL_CHAR do not hide away chars
-        assert len(SPECIAL_CHARS) == 36
+        assert len(SPECIAL_CHARS) == 26
+
+    def test_digit_chars_do_not_quote(self):
+        # confirm that we have all 10 digits
+        assert len(DIGIT_CHARS) == 10
 
     def test_get_passphrase(self):
         # we can get passphrases
@@ -237,6 +268,15 @@ def test_get_passphrase_specialchars(self):
         # the 2nd special char position might be equal to 1st.
         assert len(specials) > 0
 
+    def test_get_passphrase_digitchars(self):
+        # we can request special chars in passphrases
+        options = handle_options(args=[])
+        options.digits = 2
+        phrase = get_passphrase(options)
+        digits = [x for x in phrase if x in DIGIT_CHARS]
+        # the 2nd special char position might be equal to 1st.
+        assert len(digits) > 0
+
     def test_get_passphrase_delimiters(self):
         # we can set separators
         options = handle_options(args=[])
@@ -355,6 +395,15 @@ def test_main_specialchars(self, argv_handler, capsys):
         specials = [x for x in out if x in SPECIAL_CHARS]
         assert len(specials) > 0
 
+    def test_main_digitchars(self, argv_handler, capsys):
+        # number of specialchars is respected in calls to main.
+        sys.stdin = StringIO("word1\n")
+        sys.argv = ['diceware', '-n', '1', '-D', '1', '-']
+        main()
+        out, err = capsys.readouterr()
+        digits = [x for x in out if x in DIGIT_CHARS]
+        assert len(digits) > 0
+
     def test_main_wordlist(self, argv_handler, capsys, wordlists_dir):
         # we can pick the wordlist we prefer
         wordlists_dir.join('wordlist_foo.txt').write("foo\n")