1 files changed, 89 insertions, 64 deletions

diff --git a/src/otp2289/__main__.py b/src/otp2289/__main__.py
index 0fdcc52..c5d28b6 100644
--- a/src/otp2289/__main__.py
+++ b/src/otp2289/__main__.py
@@ -51,6 +51,56 @@ def eprint(
     print(*value, sep=sep, end=end, file=sys.stderr, flush=True)
 
 
+def get_rnd_seed() -> str:
+    """
+    Returns a random seed in the format:
+
+    2 random letters (capitalize()) + 5 random digits
+    """
+    rnd = secrets.SystemRandom()
+    return ''.join(
+        rnd.choices(string.ascii_lowercase, k=2)
+    ).capitalize() + ''.join(rnd.choices(string.digits, k=5))
+
+
+def initiate_new_sequence(args: argparse.Namespace) -> str:
+    """
+    Generates a new sequence based on the parameters sent from the parser.
+
+    :param args: The arguments assigned from argparse
+    :type args: argparse.Namespace
+
+    :raises otp2289.OTPChallengeError: If the challenge is invalid
+
+    :raises otp2289.OTPGeneratorError: If generator parameters are wrong
+
+    :raises otp2289.OTPResponseError: If tokens or hex can not be generated
+
+    :return: The response string
+    :rtype: str
+    """
+    if not args.seed:
+        args.seed = get_rnd_seed()
+
+    header = ''
+    if not args.quiet:
+        header = (
+            f'Seed: {args.seed}, Step: {args.step}, '
+            f'Hash: {args.hash_algo}{os.linesep}'
+        )
+    generator = otp2289.OTPGenerator(
+        args.password.encode(), args.seed, args.hash_algo
+    )
+    if args.challenge_string:
+        return (
+            header
+            + generator.generate_otp_response_from_challenge(
+                args.challenge_string
+            ).hexdigest
+        )
+    return header + generator.generate_otp_response(args.step).hexdigest
+
+
 def generate_otp_response(args: argparse.Namespace) -> str:
     """
     Generates a response based on the parameters sent from the parser
@@ -62,20 +112,22 @@ def generate_otp_response(args: argparse.Namespace) -> str:
 
     :raises otp2289.OTPGeneratorError: If generator parameters are wrong
 
+    :raises otp2289.OTPResponseError: If tokens or hex can not be generated
+
     :return: The response string
     :rtype: str
     """
-    generator = otp2289.generator.OTPGenerator(
+    generator = otp2289.OTPGenerator(
         args.password.encode(), args.seed, args.hash_algo
     )
     if args.challenge_string:
-        if args.output_format == 'token':
-            return generator.generate_otp_words_from_challenge(
-                args.challenge_string
-            )
-        return generator.generate_otp_hexdigest_from_challenge(
-            args.challenge_string
+        response = generator.generate_otp_response_from_challenge(
+            args.challange_string
         )
+        if args.output_format == 'token':
+            return response.words
+        return response.hexdigest
+
     # regular parameters
     header = ''
     if not args.quiet:
@@ -83,9 +135,12 @@ def generate_otp_response(args: argparse.Namespace) -> str:
             f'Seed: {args.seed}, Step: {args.step}, '
             f'Hash: {args.hash_algo}{os.linesep}'
         )
+
+    response = generator.generate_otp_response(args.step)
+
     if args.output_format == 'token':
-        return header + generator.generate_otp_words(args.step)
-    return header + generator.generate_otp_hexdigest(args.step)
+        return header + response.words
+    return header + response.hexdigest
 
 
 def generate_otp_range(args: argparse.Namespace) -> str:
@@ -99,20 +154,24 @@ def generate_otp_range(args: argparse.Namespace) -> str:
 
     :raises otp2289.OTPGeneratorError: If generator parameters are wrong
 
+    :raises otp2289.OTPResponseError: If tokens or hex can not be generated
+
     :return: The responses string
     :rtype: str
     """
-    generator = otp2289.generator.OTPGenerator(
+    generator = otp2289.OTPGenerator(
         args.password.encode(), args.seed, args.hash_algo
     )
-    if args.output_format == 'token':
-        method = generator.generate_otp_words
-    else:
-        method = generator.generate_otp_hexdigest
 
     # handle most cases explicitly
     if args.range == 1:
-        return f'{args.step}: ' + method(args.step)
+        response = generator.generate_otp_response(args.step)
+        if args.output_format == 'token':
+            response_str = response.words
+        else:
+            response_str = response.hexdigest
+        return f'{args.step}: ' + response_str
+
     args.range = min(args.range, args.step + 1)
 
     # any need for quiet?
@@ -122,9 +181,19 @@ def generate_otp_range(args: argparse.Namespace) -> str:
             f'Seed: {args.seed}, Step: {args.step}, '
             f'Hash: {args.hash_algo}, Range: {args.range}{os.linesep}'
         )
+
+    if args.output_format == 'token':
+        return header + os.linesep.join(
+            [
+                f'{step}: ' + generator.generate_otp_response(step).words
+                for step in range(args.step, args.step - args.range, -1)
+            ]
+        )
+
+    # assume hex
     return header + os.linesep.join(
         [
-            f'{step}: ' + method(step)
+            f'{step}: ' + generator.generate_otp_response(step).hexdigest
             for step in range(args.step, args.step - args.range, -1)
         ]
     )
@@ -172,50 +241,6 @@ def get_password(args: argparse.Namespace) -> str:
     return args.password
 
 
-def get_rnd_seed() -> str:
-    """
-    Returns a random seed in the format:
-
-    2 random letters (capitalize()) + 5 random digits
-    """
-    rnd = secrets.SystemRandom()
-    return ''.join(
-        rnd.choices(string.ascii_lowercase, k=2)
-    ).capitalize() + ''.join(rnd.choices(string.digits, k=5))
-
-
-def initiate_new_sequence(args: argparse.Namespace) -> str:
-    """
-    Generates a new sequence based on the parameters sent from the parser.
-
-    :param args: The arguments assigned from argparse
-    :type args: argparse.Namespace
-
-    :raises otp2289.OTPChallengeError: If the challenge is invalid
-
-    :raises otp2289.OTPGeneratorError: If generator parameters are wrong
-
-    :return: The response string
-    :rtype: str
-    """
-    if not args.seed:
-        args.seed = get_rnd_seed()
-    header = ''
-    if not args.quiet:
-        header = (
-            f'Seed: {args.seed}, Step: {args.step}, '
-            f'Hash: {args.hash_algo}{os.linesep}'
-        )
-    generator = otp2289.generator.OTPGenerator(
-        args.password.encode(), args.seed, args.hash_algo
-    )
-    if args.challenge_string:
-        return header + generator.generate_otp_hexdigest_from_challenge(
-            args.challenge_string
-        )
-    return header + generator.generate_otp_hexdigest(args.step)
-
-
 def main(inargs: list[str] | None = None) -> None:
     """the main entry point"""
     parser = argparse.ArgumentParser(
@@ -278,8 +303,8 @@ def main(inargs: list[str] | None = None) -> None:
     parser.add_argument(
         '-a',
         '--hash-algorithm',
-        metavar='<md5 | sha1>',
         type=str,
+        choices=['md5', 'sha1'],
         dest='hash_algo',
         default='md5',
         help='The hash algorithm to use. Possible values: md5 (default), sha1',
@@ -296,8 +321,8 @@ def main(inargs: list[str] | None = None) -> None:
     parser.add_argument(
         '-f',
         '--output-format',
-        metavar='<hex | token>',
         type=str,
+        choices=['hex', 'token'],
         dest='output_format',
         default='hex',
         help='The output format to use. Possible values: hex (default), token',
@@ -368,9 +393,9 @@ def main(inargs: list[str] | None = None) -> None:
         if args.generate_otp_response:
             print(generate_otp_response(args))
         sys.exit(0)
-    except otp2289.generator.OTPGeneratorError as exp:
+    except otp2289.OTPGeneratorError as exp:
         eprint(f'GeneratorException: {exp}')
-    except otp2289.generator.OTPChallengeError as exp:
+    except otp2289.OTPChallengeError as exp:
         eprint(f'ChallengeException: {exp}')
     except Exception as exp:
         eprint(f'Unknown error: {exp}')