Fix 406 (#454)

* Added header option for encode

* added tests and None check

* updated quotes for consistency

* format fixes after tox

2 files changed, 56 insertions, 7 deletions

diff --git a/jwt/__main__.py b/jwt/__main__.py
index 1493e50..a6b9f39 100644
--- a/jwt/__main__.py
+++ b/jwt/__main__.py
@@ -43,7 +43,19 @@ def encode_payload(args):
 
         payload[k] = v
 
-    token = encode(payload, key=args.key, algorithm=args.algorithm)
+    # Build header object to encode
+    header = {}
+    if args.header:
+        try:
+            header = json.loads(args.header)
+        except Exception as e:
+            raise ValueError(
+                "Error loading header: %s. See --help for usage." % e
+            )
+
+    token = encode(
+        payload, key=args.key, algorithm=args.algorithm, headers=header
+    )
 
     return token.decode("utf-8")
 
@@ -86,6 +98,12 @@ def build_argparser():
     %(prog)s --key=secret encode foo=bar exp=+10
 
     The exp key is special and can take an offset to current Unix time.
+
+    %(prog)s --key=secret --header='{"typ":"jwt", "alg":"RS256"}' encode is=me
+
+    The header option can be provided for input to encode in the jwt. The format
+    requires the header be enclosed in single quote and key/value pairs with double
+    quotes.
     """
 
     arg_parser = argparse.ArgumentParser(prog="pyjwt", usage=usage)
@@ -110,6 +128,14 @@ def build_argparser():
         help="set crypto algorithm to sign with. default=HS256",
     )
 
+    arg_parser.add_argument(
+        "--header",
+        dest="header",
+        metavar="HEADER",
+        default=None,
+        help="set jwt header",
+    )
+
     subparsers = arg_parser.add_subparsers(
         title="PyJWT subcommands",
         description="valid subcommands",
diff --git a/tests/test_cli.py b/tests/test_cli.py
index 6246c9a..ba65280 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -27,6 +27,23 @@ class TestCli:
 
         assert "Key is required when encoding" in str(excinfo.value)
 
+    def test_encode_header_raises_value_error_bad_dict(self):
+        encode_args = [
+            "--key=secret",
+            "--header=dfsfd",
+            "encode",
+            "name=Vader",
+            "job=Sith",
+        ]
+        parser = build_argparser()
+
+        args = parser.parse_args(encode_args)
+
+        with pytest.raises(ValueError) as excinfo:
+            encode_payload(args)
+
+        assert "Error loading header:" in str(excinfo.value)
+
     def test_decode_payload_raises_decoded_error(self):
         decode_args = ["--key", "1234", "decode", "wrong-token"]
         parser = build_argparser()
@@ -56,7 +73,12 @@ class TestCli:
         assert "There was an error decoding the token" in str(excinfo.value)
 
     def test_decode_payload_terminal_tty(self, monkeypatch):
-        encode_args = ["--key=secret-key", "encode", "name=hello-world"]
+        encode_args = [
+            "--key=secret-key",
+            '--header={"alg":"HS256"}',
+            "encode",
+            "name=hello-world",
+        ]
         parser = build_argparser()
         parsed_encode_args = parser.parse_args(encode_args)
         token = encode_payload(parsed_encode_args)
@@ -84,16 +106,17 @@ class TestCli:
             )
 
     @pytest.mark.parametrize(
-        "key,name,job,exp,verify",
+        "key,header,name,job,exp,verify",
         [
-            ("1234", "Vader", "Sith", None, None),
-            ("4567", "Anakin", "Jedi", "+1", None),
-            ("4321", "Padme", "Queen", "4070926800", "true"),
+            ("1234", "{}", "Vader", "Sith", None, None),
+            ("4567", '{"typ":"test"}', "Anakin", "Jedi", "+1", None),
+            ("4321", "", "Padme", "Queen", "4070926800", "true"),
         ],
     )
-    def test_encode_decode(self, key, name, job, exp, verify):
+    def test_encode_decode(self, key, header, name, job, exp, verify):
         encode_args = [
             "--key={0}".format(key),
+            "--header={0}".format(header),
             "encode",
             "name={0}".format(name),
             "job={0}".format(job),