<?php

namespace iDServer\Internal\Crypto;

use DOMElement;
use DOMNode;
use DOMNodeList;
use DOMXPath;
use Exception;
use iDServer\Internal\XMLSecurityDSig as EGXMLSecurityDSig;
use RobRichards\XMLSecLibs\XMLSecurityDSig;
use RobRichards\XMLSecLibs\XMLSecEnc;
use iDServer\Constants;
use iDServer\Exception\iDServerException;
use iDServer\Log;
use DOMDocument;
use RobRichards\XMLSecLibs\XMLSecurityKey;
use iDServer\Internal\XMLSecurityKey as EGXMLSecurityKey;
use function iDServer\Internal\Utilities\validateOpenSSLKeyType;

/**
 * Sign the given XML with a key and cert
 *
 * @param DOMDocument $doc
 * @param mixed $signKey
 *   - PEM formatted decrypted private key string or
 *   - the return value of {@link https://www.php.net/manual/en/function.openssl-sign.php}
 *     which is resource in PHP 7.4 and below and OpenSSLAsymmetricKey in PHP 8.0 and above.
 *
 * @param string $signCert
 * @param string|null $algorithm Automatically detected if not set or set to null
 * @param string $passphrase It can be used if $signKey is an encrypted PEM formatted private key string
 *                           and not and not an already decrypted key.
 * @return DOMDocument
 * @throws Exception
 */
function sign(DOMDocument $doc, $signKey, string $signCert, string $algorithm = null, string $passphrase = ''): DOMDocument
{
    Log::debug('Signing XML document.');

    validateOpenSSLKeyType($signKey);

    $objDSig = new EGXMLSecurityDSig();
    $objDSig->setCanonicalMethod(XMLSecurityDSig::EXC_C14N);
    $objDSig->addReference(
        $doc->firstChild,
        XMLSecurityDSig::SHA256,
        array(
            'http://www.w3.org/2000/09/xmldsig#enveloped-signature',
            'http://www.w3.org/2001/10/xml-exc-c14n#'
        ),
        array(
            'id_name' => 'ID',
            'overwrite' => false
        )
    );

    if (!$algorithm) {
        $pubKeyInfo = openssl_pkey_get_details(openssl_pkey_get_public($signCert));
        $bits = $pubKeyInfo['bits'];
        if ($pubKeyInfo['type'] === OPENSSL_KEYTYPE_EC) {
            switch ($bits) {
                case 256: $algorithm = EGXMLSecurityKey::ECC_SECP256R1; break;
                case 384: $algorithm = EGXMLSecurityKey::ECC_SECP384R1; break;
                case 521: $algorithm = EGXMLSecurityKey::ECC_SECP521R1; break;
            }
        } else {
            // TODO: implement http://www.w3.org/2007/05/xmldsig-more%23sha256-rsa-MGF1
            $algorithm = XMLSecurityKey::RSA_SHA256;
        }
    }

    $objKey = new EGXMLSecurityKey($algorithm, array('type' => 'private'));
    $objKey->passphrase = $passphrase;
    $objKey->loadKey($signKey);

    $objDSig->sign($objKey);
    $objDSig->add509Cert($signCert);
    $objDSig->insertSignature($doc->documentElement, $doc->documentElement->firstChild->nextSibling);

    return $doc;
}

/**
 * Verify if the given SAML XML is valid
 *
 * @param DOMDocument $doc
 * @param string $cert
 * @return bool
 * @throws iDServerException In case of invalid SAML
 * @throws Exception Thrown by XMLSecurityDSig
 */
function verifySaml(DOMDocument $doc, string $cert, $passphrase = ''): bool
{
    $objXMLSecDSig = new EGXMLSecurityDSig();
    $objDSig = $objXMLSecDSig->locateSignature($doc);
    if (!$objDSig) {
        Log::error("SAML verification failed: missing SAML signature.");
        throw new iDServerException("Missing SAML signature");
    }

    $objXMLSecDSig->canonicalizeSignedInfo();
    $objXMLSecDSig->idKeys = array('ID');
    $result = $objXMLSecDSig->validateReference();
    if (!$result) {
        Log::error("SAML verification failed: reference validation failed.");
        throw new iDServerException("SAML reference validation failed");
    }

    $objKey = $objXMLSecDSig->locateKey();
    if (!$objKey) {
        Log::error("SAML verification failed: the message did not contain the key.");
        throw new iDServerException("SAML message did not contain the key");
    }
    XMLSecEnc::staticLocateKeyInfo($objKey, $objDSig);
    $objKey->passphrase = $passphrase;
    $objKey->loadKey($cert);

    $result = $objXMLSecDSig->verify($objKey);
    if (!$result) {
        Log::error("SAML verification failed: signature validation failed.");
        throw new iDServerException("SAML signature validation failed");
    }

    return true;
}

/**
 * Verify a SAML signature against an array of certificates.
 *
 * If any certificate matches, the message is considered verified.
 *
 * @param DomDocument $doc The SAML message in XML form. This function modifies it.
 * @param array $certs An array of certificates to validate the message against.
 *
 * @throws iDServerException if the verification fails.
 * @throws Exception Thrown by XMLSecurityDSig
 */
function verifySamlWithFallbacks(DOMDocument $doc, array $certs)
{
    $failed = false;
    foreach ($certs as $k => $cert) {
        $vdoc = clone $doc;
        Log::debug("Verifying message with certificate #" . $k);
        try {
            verifySaml($vdoc, $cert);
            $failed = false;
            $doc = $vdoc;
            break;
        } catch (iDServerException $e) {
            Log::debug("Verification failed: " . $e->getMessage());
            $failed = true;
        }
    }
    if ($failed) {
        throw new iDServerException("SAML signature validation failed");
    }
}

/**
 * @param DOMElement|DOMDocument|DOMNode $node
 * @param mixed $cryptKey
 *   - PEM formatted decrypted private key string or
 *   - the return value of {@link https://www.php.net/manual/en/function.openssl-sign.php}
 *     which is resource in PHP 7.4 and below and OpenSSLAsymmetricKey in PHP 8.0 and above.
 * @throws iDServerException If the decryption fails
 * @throws Exception Thrown by XMLSecEnc
 */
function decryptNode($node, $cryptKey)
{
    validateOpenSSLKeyType($cryptKey);

    $objenc = new XMLSecEnc();
    /* @var $encData DOMElement */
    $xpath = new DOMXPath($node->ownerDocument);
    // Do not use $objenc->locateEncryptedData()!
    // It always locates the root node's first encrypted data.
    $query = "*[local-name()='EncryptedData' and namespace-uri()='" . XMLSecEnc::XMLENCNS . "']";
    $nodeset = $xpath->query($query, $node);
    $encData = $nodeset->item(0);
    $objenc->setNode($encData);
    $objenc->type = $encData->getAttribute("Type");

    $xpath = new DOMXPath($node->ownerDocument);
    $xpath->registerNamespace('xmlsecenc', XMLSecEnc::XMLENCNS);
    /* @var $encmeth DOMElement */
    $encmeth = $xpath->query(".//xmlsecenc:EncryptionMethod", $encData)->item(0);
    $attrAlgorithm = $encmeth->getAttribute("Algorithm");

    if ($attrAlgorithm === Constants\AES256_CBC) {
        $crypto = [
            'algo' => 'AES-256-CBC',
            'flags' => OPENSSL_RAW_DATA | OPENSSL_ZERO_PADDING,
            'postProcess' => 'iDServer\Internal\Crypto\stripPkcs7Padding', /* @see stripPkcs7Padding() */
            'is_authenticated' => false
        ];
    } elseif ($attrAlgorithm === Constants\AES128_CBC) {
        $crypto = [
            'algo' => 'AES-128-CBC',
            'flags' => OPENSSL_RAW_DATA | OPENSSL_ZERO_PADDING,
            'postProcess' => 'iDServer\Internal\Crypto\stripPkcs7Padding', /* @see stripPkcs7Padding() */
            'is_authenticated' => false
        ];
    } elseif ($attrAlgorithm === Constants\AES256_GCM) {
        $crypto = [
            'algo' => 'aes-256-gcm',
            'flags' => OPENSSL_RAW_DATA,
            'is_authenticated' => true,
            'tag_length' => 16
        ];
    } else {
        Log::error("Node decryption failed: encountered unknown encryption algorithm: " . $attrAlgorithm);
        throw new iDServerException("Unknown encryption algorithm: " . $attrAlgorithm);
    }

    $objKeyInfo = $objenc->locateKeyInfo();
    // passing null in place of the objBaseKey is cute,
    // but it's basically just a big hack to avoid re-implementing half of XMLSecLib
    // see https://github.com/robrichards/xmlseclibs/blob/master/src/XMLSecEnc.php#L430 for failure conditions
    if (empty($objKeyInfo)) {
        Log::error("Node decryption failed: Unexpected key storage method.");
        throw new iDServerException("Unexpected key storage method.");
    }

    $objencKey = $objKeyInfo->encryptedCtx;
    $k = $objencKey->locateKey();
    $k->loadKey($cryptKey);
    $key = $objencKey->decryptNode($k);

    // and now we exit XMLSecLib completely because we'd like to support GCM
    $ivlen = openssl_cipher_iv_length($crypto['algo']);
    $cipherValue = $objenc->getCipherValue();

    $iv = substr($cipherValue, 0, $ivlen);
    $ciph = substr($cipherValue, $ivlen);

    if ($crypto['is_authenticated']) {
        $tag = substr($ciph, -1 * $crypto['tag_length']);
        $ciph = substr($ciph, 0, -1 * $crypto['tag_length']);
        $plaintext = openssl_decrypt($ciph, $crypto['algo'], $key, $crypto['flags'], $iv, $tag);
    } else {
        $plaintext = openssl_decrypt($ciph, $crypto['algo'], $key, $crypto['flags'], $iv);
    }

    if ($plaintext === false) {
        Log::error("OpenSSL decryption failed.");
        throw new iDServerException("Decryption failed.");
    }

    if (isset($crypto['postProcess'])) {
        $plaintext = call_user_func($crypto['postProcess'], $plaintext);
    }

    $newdoc = new DOMDocument();
    $newdoc->loadXML($plaintext);
    $decNode = $node->ownerDocument->importNode($newdoc->documentElement, true);
    $node->parentNode->replaceChild($decNode, $node);
}

/**
 * @param DOMNodeList $list
 * @param mixed $cryptKey
 *   - PEM formatted decrypted private key string or
 *   - the return value of {@link https://www.php.net/manual/en/function.openssl-sign.php}
 *     which is resource in PHP 7.4 and below and OpenSSLAsymmetricKey in PHP 8.0 and above.
 * @throws Exception
 */
function decryptNodeList(DOMNodeList $list, $cryptKey)
{
    validateOpenSSLKeyType($cryptKey);

    $nodes = [];
    foreach ($list as $node) {
        $nodes[] = $node;
    }
    foreach ($nodes as $node) {
        decryptNode($node, $cryptKey);
    }
}

/**
 * @param string $plaintext
 * @return bool|string
 */
function stripPkcs7Padding(string $plaintext)
{
    return substr($plaintext, 0, -ord(substr($plaintext, -1)));
}