# HG changeset patch # User Brian Neal # Date 1372388095 18000 # Node ID efdb3b57fb46714c9b8eec22cd65232bfce3a2b9 # Parent 7a1b5740236f1d978d2cd41a4878b13776f3bd11 Rework command-line for keygen. Created IndicatorIter & tests. diff -r 7a1b5740236f -r efdb3b57fb46 m209/keylist/key_list.py --- 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 + diff -r 7a1b5740236f -r efdb3b57fb46 m209/keylist/tests/test_key_list.py --- /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)) + diff -r 7a1b5740236f -r efdb3b57fb46 m209/main.py --- 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)