Change optparse for argparse. (jpadilla#238)

Author: Froilan Irizarry

.coveragerc

@@ -4,4 +4,4 @@ omit =
     .tox/*
     setup.py
     *.egg/*
-    */__main__.py
+

.travis.yml

@@ -23,6 +23,7 @@ env:
   - TOXENV=py35-contrib_crypto
   - TOXENV=py36-contrib_crypto
   - TOXENV=py27-contrib_crypto
+  
 install:
   - pip install -U pip
   - pip install -U tox coveralls

jwt/__main__.py

@@ -2,135 +2,161 @@
 
 from __future__ import absolute_import, print_function
 
+import argparse
 import json
-import optparse
 import sys
 import time
 
-from . import DecodeError, __package__, __version__, decode, encode
+from . import DecodeError, __version__, decode, encode
 
 
-def main():
+def encode_payload(args):
+    # Try to encode
+    if args.key is None:
+        raise ValueError('Key is required when encoding. See --help for usage.')
+
+    # Build payload object to encode
+    payload = {}
+
+    for arg in args.payload:
+        k, v = arg.split('=', 1)
+
+        # exp +offset special case?
+        if k == 'exp' and v[0] == '+' and len(v) > 1:
+            v = str(int(time.time()+int(v[1:])))
+
+        # Cast to integer?
+        if v.isdigit():
+            v = int(v)
+        else:
+            # Cast to float?
+            try:
+                v = float(v)
+            except ValueError:
+                pass
+
+        # Cast to true, false, or null?
+        constants = {'true': True, 'false': False, 'null': None}
+
+        if v in constants:
+            v = constants[v]
+
+        payload[k] = v
+
+    token = encode(
+        payload,
+        key=args.key,
+        algorithm=args.algorithm
+    )
+
+    return token.decode('utf-8')
+
+
+def decode_payload(args):
+    try:
+        if sys.stdin.isatty():
+            token = sys.stdin.read()
+        else:
+            token = args.token
+
+        token = token.encode('utf-8')
+        data = decode(token, key=args.key, verify=args.verify)
+
+        return json.dumps(data)
+
+    except DecodeError as e:
+        raise DecodeError('There was an error decoding the token: %s' % e)
+
 
-    usage = '''Encodes or decodes JSON Web Tokens based on input.
+def build_argparser():
 
-  %prog [options] input
+    usage = '''
+    Encodes or decodes JSON Web Tokens based on input.
 
-Decoding examples:
+    %(prog)s [options] <command> [options] input
 
-  %prog --key=secret json.web.token
-  %prog --no-verify json.web.token
+    Decoding examples:
 
-Encoding requires the key option and takes space separated key/value pairs
-separated by equals (=) as input. Examples:
+    %(prog)s --key=secret decode json.web.token
+    %(prog)s decode --no-verify json.web.token
 
-  %prog --key=secret iss=me exp=1302049071
-  %prog --key=secret foo=bar exp=+10
+    Encoding requires the key option and takes space separated key/value pairs
+    separated by equals (=) as input. Examples:
 
-The exp key is special and can take an offset to current Unix time.\
-'''
-    p = optparse.OptionParser(
-        usage=usage,
+    %(prog)s --key=secret encode iss=me exp=1302049071
+    %(prog)s --key=secret encode foo=bar exp=+10
+
+    The exp key is special and can take an offset to current Unix time.
+    '''
+
+    arg_parser = argparse.ArgumentParser(
         prog='pyjwt',
-        version='%s %s' % (__package__, __version__),
+        usage=usage
     )
 
-    p.add_option(
-        '-n', '--no-verify',
-        action='store_false',
-        dest='verify',
-        default=True,
-        help='ignore signature and claims verification on decode'
+    arg_parser.add_argument(
+        '-v', '--version',
+        action='version',
+        version='%(prog)s ' + __version__
     )
 
-    p.add_option(
+    arg_parser.add_argument(
         '--key',
         dest='key',
         metavar='KEY',
         default=None,
         help='set the secret key to sign with'
     )
 
-    p.add_option(
+    arg_parser.add_argument(
         '--alg',
         dest='algorithm',
         metavar='ALG',
         default='HS256',
         help='set crypto algorithm to sign with. default=HS256'
     )
 
-    options, arguments = p.parse_args()
+    subparsers = arg_parser.add_subparsers(
+        title='PyJWT subcommands',
+        description='valid subcommands',
+        help='additional help'
+    )
 
-    if len(arguments) > 0 or not sys.stdin.isatty():
-        if len(arguments) == 1 and (not options.verify or options.key):
-            # Try to decode
-            try:
-                if not sys.stdin.isatty():
-                    token = sys.stdin.read()
-                else:
-                    token = arguments[0]
+    # Encode subcommand
+    encode_parser = subparsers.add_parser('encode', help='use to encode a supplied payload')
 
-                token = token.encode('utf-8')
-                data = decode(token, key=options.key, verify=options.verify)
+    payload_help = """Payload to encode. Must be a space separated list of key/value
+    pairs separated by equals (=) sign."""
 
-                print(json.dumps(data))
-                sys.exit(0)
-            except DecodeError as e:
-                print(e)
-                sys.exit(1)
+    encode_parser.add_argument('payload', nargs='+', help=payload_help)
+    encode_parser.set_defaults(func=encode_payload)
 
-        # Try to encode
-        if options.key is None:
-            print('Key is required when encoding. See --help for usage.')
-            sys.exit(1)
+    # Decode subcommand
+    decode_parser = subparsers.add_parser('decode', help='use to decode a supplied JSON web token')
+    decode_parser.add_argument('token', help='JSON web token to decode.')
 
-        # Build payload object to encode
-        payload = {}
+    decode_parser.add_argument(
+        '-n', '--no-verify',
+        action='store_false',
+        dest='verify',
+        default=True,
+        help='ignore signature and claims verification on decode'
+    )
 
-        for arg in arguments:
-            try:
-                k, v = arg.split('=', 1)
+    decode_parser.set_defaults(func=decode_payload)
+
+    return arg_parser
 
-                # exp +offset special case?
-                if k == 'exp' and v[0] == '+' and len(v) > 1:
-                    v = str(int(time.time()+int(v[1:])))
 
-                # Cast to integer?
-                if v.isdigit():
-                    v = int(v)
-                else:
-                    # Cast to float?
-                    try:
-                        v = float(v)
-                    except ValueError:
-                        pass
+def main():
+    arg_parser = build_argparser()
 
-                # Cast to true, false, or null?
-                constants = {'true': True, 'false': False, 'null': None}
+    try:
+        arguments = arg_parser.parse_args(sys.argv[1:])
 
-                if v in constants:
-                    v = constants[v]
+        output = arguments.func(arguments)
 
-                payload[k] = v
-            except ValueError:
-                print('Invalid encoding input at {}'.format(arg))
-                sys.exit(1)
-
-        try:
-            token = encode(
-                payload,
-                key=options.key,
-                algorithm=options.algorithm
-            )
-
-            print(token)
-            sys.exit(0)
-        except Exception as e:
-            print(e)
-            sys.exit(1)
-    else:
-        p.print_help()
-
-
-if __name__ == '__main__':
-    main()
+        print(output)
+    except Exception as e:
+        print('There was an unforseen error: ', e)
+        arg_parser.print_help()

tests/test_cli.py

@@ -0,0 +1,127 @@
+
+import argparse
+import json
+import sys
+
+import jwt
+from jwt.__main__ import build_argparser, decode_payload, encode_payload, main
+
+import pytest
+
+
+class TestCli:
+
+    def test_build_argparse(self):
+        args = ['--key', '1234', 'encode', 'name=Vader']
+        parser = build_argparser()
+        parsed_args = parser.parse_args(args)
+
+        assert parsed_args.key == '1234'
+
+    def test_encode_payload_raises_value_error_key_is_required(self):
+        encode_args = ['encode', 'name=Vader', 'job=Sith']
+        parser = build_argparser()
+
+        args = parser.parse_args(encode_args)
+
+        with pytest.raises(ValueError) as excinfo:
+            encode_payload(args)
+
+        assert 'Key is required when encoding' in str(excinfo.value)
+
+    def test_decode_payload_raises_decoded_error(self):
+        decode_args = ['--key', '1234', 'decode', 'wrong-token']
+        parser = build_argparser()
+
+        args = parser.parse_args(decode_args)
+
+        with pytest.raises(jwt.DecodeError) as excinfo:
+            decode_payload(args)
+
+        assert 'There was an error decoding the token' in str(excinfo.value)
+
+    def test_decode_payload_raises_decoded_error_isatty(self, monkeypatch):
+        def patched_sys_stdin_read():
+            raise jwt.DecodeError()
+
+        decode_args = ['--key', '1234', 'decode', 'wrong-token']
+        parser = build_argparser()
+
+        args = parser.parse_args(decode_args)
+
+        monkeypatch.setattr(sys.stdin, 'isatty', lambda: True)
+        monkeypatch.setattr(sys.stdin, 'read', patched_sys_stdin_read)
+
+        with pytest.raises(jwt.DecodeError) as excinfo:
+            decode_payload(args)
+
+        assert 'There was an error decoding the token' in str(excinfo.value)
+
+    @pytest.mark.parametrize('key,name,job,exp,verify', [
+        ('1234', 'Vader', 'Sith', None, None),
+        ('4567', 'Anakin', 'Jedi', '+1', None),
+        ('4321', 'Padme', 'Queen', '4070926800', 'true'),
+    ])
+    def test_encode_decode(self, key, name, job, exp, verify):
+        encode_args = [
+            '--key={0}'.format(key),
+            'encode',
+            'name={0}'.format(name),
+            'job={0}'.format(job),
+        ]
+        if exp:
+            encode_args.append('exp={0}'.format(exp))
+        if verify:
+            encode_args.append('verify={0}'.format(verify))
+
+        parser = build_argparser()
+        parsed_encode_args = parser.parse_args(encode_args)
+        token = encode_payload(parsed_encode_args)
+        assert token is not None
+        assert token is not ''
+
+        decode_args = [
+            '--key={0}'.format(key),
+            'decode',
+            token
+        ]
+        parser = build_argparser()
+        parsed_decode_args = parser.parse_args(decode_args)
+
+        actual = json.loads(decode_payload(parsed_decode_args))
+        expected = {
+            'job': job,
+            'name': name,
+        }
+        assert actual['name'] == expected['name']
+        assert actual['job'] == expected['job']
+
+    @pytest.mark.parametrize('key,name,job,exp,verify', [
+        ('1234', 'Vader', 'Sith', None, None),
+        ('4567', 'Anakin', 'Jedi', '+1', None),
+        ('4321', 'Padme', 'Queen', '4070926800', 'true'),
+    ])
+    def test_main(self, monkeypatch, key, name, job, exp, verify):
+        args = [
+            'test_cli.py',
+            '--key={0}'.format(key),
+            'encode',
+            'name={0}'.format(name),
+            'job={0}'.format(job),
+        ]
+        if exp:
+            args.append('exp={0}'.format(exp))
+        if verify:
+            args.append('verify={0}'.format(verify))
+        monkeypatch.setattr(sys, 'argv', args)
+        main()
+
+    def test_main_throw_exception(self, monkeypatch, capsys):
+        def patched_argparser_parse_args(self, args):
+            raise Exception('NOOOOOOOOOOO!')
+
+        monkeypatch.setattr(argparse.ArgumentParser, 'parse_args', patched_argparser_parse_args)
+        main()
+        out, _ = capsys.readouterr()
+
+        assert 'NOOOOOOOOOOO!' in out