From bc4a2e85a59548064953734a2e5520ae314a0270 Mon Sep 17 00:00:00 2001
From: Erik Jaegervall <erik.jaegervall@se.bosch.com>
Date: Fri, 1 Mar 2024 15:11:38 +0100
Subject: [PATCH] Update createToken.py

Origin https://github.com/eclipse/kuksa.val/pull/676
---
 jwt/README.md      |  2 --
 jwt/createToken.py | 34 +++++++++++++++++++++++++---------
 2 files changed, 25 insertions(+), 11 deletions(-)

diff --git a/jwt/README.md b/jwt/README.md
index db486c3..b0e70cc 100644
--- a/jwt/README.md
+++ b/jwt/README.md
@@ -20,8 +20,6 @@ Two helper scripts exist for generating keys and tokens
 * [recreateJWTkeyPair.sh](recreateJWTkeyPair.sh) to regenerate the JWT keys used for signing
 * [createToken.py](createToken.py) to create signed tokens, requires `*.json` files as parameters
 
-Note that token generation must take place from the directory containing `createToken.py`
-
 An example is shown below:
 
 ```
diff --git a/jwt/createToken.py b/jwt/createToken.py
index 55e26d1..e5342e7 100755
--- a/jwt/createToken.py
+++ b/jwt/createToken.py
@@ -16,22 +16,30 @@
 # SPDX-License-Identifier: Apache-2.0
 ########################################################################
 
+
 import argparse
+import sys
+from os import path
+
 import json
 import jwt
 
-from os import path
 
+def error_exit(msg):
+    print(msg, file=sys.stderr)
+    sys.exit(1)
 
-def createJWTToken(input_filename, priv_key):
+
+def createJWTToken(input_filename, priv_key, output_filename=None):
     print("Reading JWT payload from {}".format(input_filename))
     with open(input_filename, "r") as file:
         payload = json.load(file)
 
     encoded = jwt.encode(payload, priv_key, algorithm="RS256")
 
-    output_filename = input_filename[:-5] if input_filename.endswith(".json") else input_filename
-    output_filename += ".token"
+    if output_filename is None:
+        output_filename = input_filename[:-5] if input_filename.endswith(".json") else input_filename
+        output_filename += ".token"
 
     print("Writing signed access token to {}".format(output_filename))
     with open(output_filename, "w") as output:
@@ -41,17 +49,25 @@ def createJWTToken(input_filename, priv_key):
 def main():
     parser = argparse.ArgumentParser()
     parser.add_argument("files", help="Read JWT payload from these files", nargs="+")
+    script_dir = path.abspath(path.dirname(__file__))
+    default_key_filename = path.join(script_dir, "jwt.key")
+
+    parser.add_argument("--key", help="Private key location", dest="priv_key_filename", default=default_key_filename)
+    parser.add_argument("--output", help="Name of the output file to store token to", dest="output")
     args = parser.parse_args()
 
-    script_dir = path.abspath(path.dirname(__file__))
-    priv_key_filename = path.join(script_dir, "jwt.key")
+    if args.output is not None and len(args.files) > 1:
+        error_exit("""
+        Both --output option and multiple files have been specified.
+        Output filename can be specified for single input file only!
+        """)
 
     print("Reading private key from {}".format("jwt.key"))
-    with open(priv_key_filename, "r") as file:
+    with open(args.priv_key_filename, "r") as file:
         priv_key = file.read()
 
-    for input in args.files:
-        createJWTToken(input, priv_key)
+    for input_file in args.files:
+        createJWTToken(input_file, priv_key, args.output)
 
 
 if __name__ == "__main__":