Mercurial > public > m209

changeset 37:efdb3b57fb46

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
Rework command-line for keygen. Created IndicatorIter & tests.
author Brian Neal <bgneal@gmail.com>
date Thu, 27 Jun 2013 21:54:55 -0500
parents 7a1b5740236f
children c10322a646d9
files m209/keylist/key_list.py m209/keylist/tests/test_key_list.py m209/main.py
diffstat 3 files changed, 143 insertions(+), 17 deletions(-) [+]
line wrap: on
line diff
--- a/m209/keylist/key_list.py	Wed Jun 26 21:37:07 2013 -0500
+++ b/m209/keylist/key_list.py	Thu Jun 27 21:54:55 2013 -0500
@@ -18,3 +18,30 @@
 def valid_indicator(indicator):
     """Returns True if the given indicator is valid and False otherwise."""
     return True if VALID_IND_RE.match(indicator) else False
+
+
+class IndicatorIter:
+    """Iterator class for key list indicators AA-ZZ"""
+
+    MAX_N = 26 ** 2
+
+    def __init__(self, start='AA'):
+        if not valid_indicator(start):
+            raise ValueError('invalid key list indicator')
+        self.n = (ord(start[0]) - ord('A')) * 26 + ord(start[1]) - ord('A')
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        if self.n < self.MAX_N:
+            x, y = self.n // 26, self.n % 26
+            s = chr(x + ord('A')) + chr(y + ord('A'))
+            self.n += 1
+            return s
+        raise StopIteration
+
+    def __len__(self):
+        """Returns how many indicators are available"""
+        return self.MAX_N - self.n
+
--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/m209/keylist/tests/test_key_list.py	Thu Jun 27 21:54:55 2013 -0500
@@ -0,0 +1,63 @@
+# Copyright (C) 2013 by Brian Neal.
+# This file is part of m209, the M-209 simulation.
+# m209 is released under the MIT License (see LICENSE.txt).
+
+"""Unit tests for the functions in the key_list module."""
+
+import unittest
+
+from ..key_list import valid_indicator, IndicatorIter
+
+
+INVALID_LIST = ['', 'a', 'aaa', '12', '1', '@', 'A8', 'Az', 'A)']
+
+
+class IndicatorIterTestCase(unittest.TestCase):
+
+    def test_bad_args(self):
+
+        for ind in INVALID_LIST:
+            self.assertRaises(ValueError, IndicatorIter, ind)
+
+    def test_next_initial(self):
+
+        self.assertEqual(next(IndicatorIter()), 'AA')
+        self.assertEqual(next(IndicatorIter('AA')), 'AA')
+        self.assertEqual(next(IndicatorIter('AB')), 'AB')
+        self.assertEqual(next(IndicatorIter('JA')), 'JA')
+        self.assertEqual(next(IndicatorIter('ZZ')), 'ZZ')
+
+    def test_range(self):
+
+        result = [v for v in IndicatorIter()]
+        self.assertEqual(len(result), 26 ** 2)
+
+        for n in range(26 ** 2):
+            x = (ord(result[n][0]) - ord('A')) * 26 + ord(result[n][1]) - ord('A')
+            self.assertEqual(x, n)
+
+    def test_len(self):
+
+        n = 26 ** 2
+        it = IndicatorIter()
+        while n >= 0:
+            self.assertEqual(n, len(it))
+            try:
+                next(it)
+            except StopIteration:
+                pass
+            n -= 1
+
+
+class ValidIndicatorTestCase(unittest.TestCase):
+
+    def test_invalid(self):
+
+        for ind in INVALID_LIST:
+            self.assertFalse(valid_indicator(ind))
+
+    def test_valid(self):
+
+        for i in IndicatorIter():
+            self.assertTrue(valid_indicator(i))
+
--- a/m209/main.py	Wed Jun 26 21:37:07 2013 -0500
+++ b/m209/main.py	Thu Jun 27 21:54:55 2013 -0500
@@ -8,10 +8,12 @@
 """
 import argparse
 import logging
+import os.path
+import random
 import sys
 
 from .keylist.generate import generate_key_list
-from .keylist.key_list import valid_indicator
+from .keylist.key_list import valid_indicator, IndicatorIter
 
 
 DESC = "M-209 simulator and utility program"
@@ -19,22 +21,36 @@
 LOG_CHOICES = ['debug', 'info', 'warning', 'error', 'critical']
 
 
-def keylist_range(start, end):
-    """A generator function to generate key list indicators.
+def validate_key_list_indicator(s):
+    """Validation/conversion function for validating the supplied starting key
+    list indicator.
 
-    Generates a range of indicators from start to end, inclusive.
+    Returns the string valud if valid, otherwise raises an ArgumentTypeError.
 
     """
-    def to_int(s):
-        return (ord(s[0]) - ord('A')) * 26 + ord(s[1]) - ord('A')
+    if s == '*' or valid_indicator(s):
+        return s
 
-    x = to_int(start)
-    y = to_int(end)
+    raise argparse.ArgumentTypeError('must be * or in the range AA-ZZ')
 
-    for n in range(x, y + 1):
-        c = n // 26
-        d = n % 26
-        yield chr(c + ord('A')) + chr(d + ord('A'))
+
+def validate_num_key_lists(s):
+    """Validation/conversion function for validating the number of key lists to
+    generate.
+
+    Returns the integer value if valid, otherwise raises an ArgumentTypeError
+
+    """
+    bounds = (1, 26 ** 2)
+    msg = "value must be in range {}-{}".format(*bounds)
+    try:
+        val = int(s)
+    except ValueError:
+        raise argparse.ArgumentTypeError(msg)
+
+    if not (bounds[0] <= val <= bounds[1]):
+        raise argparse.ArgumentTypeError(msg)
+    return val
 
 
 def encrypt(args):
@@ -51,6 +67,25 @@
     """Key list generation subcommand processor"""
     print('Creating key list!', args)
 
+    if not args.overwrite and os.path.exists(args.file):
+        sys.exit("File '{}' exists. Use -o to overwrite\n".format(args.file))
+
+    if args.start == '*':   # random indicators
+        indicators = random.sample([i for i in IndicatorIter()], args.number)
+    else:
+        it = IndicatorIter(args.start)
+        n = len(it)
+        if n < args.number:
+            sys.exit("Error: can only produce {} key lists when starting at {}\n".format(
+                n, args.start))
+
+        indicators = (next(it) for n in range(args.number))
+
+    key_lists = (generate_key_list(indicator) for indicator in indicators)
+
+    for key_list in key_lists:
+        print(key_list)
+
 
 def main(argv=None):
     """Entry point for the m209 command-line utility."""
@@ -89,16 +124,17 @@
     # create the parser for generating key lists
 
     kg_parser = subparsers.add_parser('keygen', aliases=['kg'],
-        description='Generate key list files',
+        description='Generate key list file',
         help='generate key list')
     kg_parser.add_argument('-f', '--file', default=DEFAULT_KEY_LIST,
         help='path to key list file [default: %(default)s]')
     kg_parser.add_argument('-o', '--overwrite', action='store_true',
         help='overwrite key list file if it exists')
-    kg_parser.add_argument('-i', '--indicators', nargs='+', metavar='XX',
-        help='key list indicators [e.g. AA BB XA-XZ]')
-    kg_parser.add_argument('-r', '--random', type=int, metavar='N',
-        help='generate N random key lists')
+    kg_parser.add_argument('-s', '--start', metavar='XX', default='AA',
+        type=validate_key_list_indicator,
+        help='starting indicator [default: %(default)s, * for random]')
+    kg_parser.add_argument('-n', '--number', type=validate_num_key_lists, default=1,
+        help='number of key lists to generate [default: %(default)s]')
     kg_parser.set_defaults(subcommand=keygen)
 
     args = parser.parse_args(args=argv)