package com.airlock.iam.saml2.infrastructure.plugin;

import com.airlock.iam.base.misc.log.LogUtil;
import com.airlock.iam.elog.api.ELogger;
import com.airlock.iam.elog.api.ELoggerFactory;
import com.airlock.iam.plugin.api.config.LogTypes;
import com.sun.identity.saml2.common.SAML2Exception;
import com.sun.identity.saml2.common.SAML2SDKUtils;
import com.sun.identity.saml2.common.SignatureAlgorithm;
import com.sun.identity.saml2.xmlsig.SigProvider;
import com.sun.identity.shared.configuration.ISystemProperties;
import com.sun.identity.shared.configuration.SystemPropertiesManager;
import com.sun.identity.shared.xml.XMLUtils;
import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import jakarta.annotation.Nonnull;
import jakarta.annotation.Nullable;
import java.security.InvalidAlgorithmParameterException;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.cert.X509Certificate;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Optional;
import javax.xml.crypto.AlgorithmMethod;
import javax.xml.crypto.KeySelector;
import javax.xml.crypto.KeySelectorException;
import javax.xml.crypto.KeySelectorResult;
import javax.xml.crypto.MarshalException;
import javax.xml.crypto.XMLCryptoContext;
import javax.xml.crypto.XMLStructure;
import javax.xml.crypto.KeySelector.Purpose;
import javax.xml.crypto.dsig.Reference;
import javax.xml.crypto.dsig.SignedInfo;
import javax.xml.crypto.dsig.Transform;
import javax.xml.crypto.dsig.XMLSignature;
import javax.xml.crypto.dsig.XMLSignatureException;
import javax.xml.crypto.dsig.XMLSignatureFactory;
import javax.xml.crypto.dsig.dom.DOMSignContext;
import javax.xml.crypto.dsig.dom.DOMValidateContext;
import javax.xml.crypto.dsig.keyinfo.KeyInfo;
import javax.xml.crypto.dsig.keyinfo.KeyInfoFactory;
import javax.xml.crypto.dsig.keyinfo.X509Data;
import javax.xml.crypto.dsig.spec.C14NMethodParameterSpec;
import javax.xml.crypto.dsig.spec.TransformParameterSpec;
import javax.xml.namespace.NamespaceContext;
import javax.xml.xpath.XPath;
import javax.xml.xpath.XPathConstants;
import javax.xml.xpath.XPathExpressionException;
import javax.xml.xpath.XPathFactory;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.Strings;
import org.apache.xml.security.Init;
import org.w3c.dom.Document;
import org.w3c.dom.Element;
import org.w3c.dom.Node;

public final class SamlSignatureProvider implements SigProvider {
   private static final ELogger LOG = ELoggerFactory.getLogger(SamlSignatureProvider.class);
   public static final String DIGEST_METHOD_KEY = "digestMethod";
   private final Optional<ISystemProperties> systemProperties = Optional.ofNullable(SystemPropertiesManager.getSystemProperties());

   private List<SignatureAlgorithm> getAllowedSignatureAlgorithms() {
      List<SignatureAlgorithm> algorithms = this.systemProperties
         .<List<SignatureAlgorithm>>map(it -> SignatureAlgorithm.convertAllowedAlgosStringToList(it.get("allowedXmlAlgorithms")))
         .orElse(Collections.emptyList());
      return algorithms.isEmpty() ? List.of(SignatureAlgorithm.RSA_SHA1, SignatureAlgorithm.DSA_SHA1) : algorithms;
   }

   private String getSystemProperty(String key) {
      return this.systemProperties.<String>map(it -> it.get(key)).orElse(null);
   }

   private String getSystemProperty(String key, String defaultValue) {
      String value = this.getSystemProperty(key);
      return StringUtils.isNotBlank(value) ? value : defaultValue;
   }

   @Override
   public Element sign(String xmlString, String idValue, PrivateKey privateKey, X509Certificate cert) throws SAML2Exception {
      String digestMethod = this.getSystemProperty("digestMethod", "http://www.w3.org/2001/04/xmlenc#sha256");
      String sigAlg = this.getSystemProperty("com.sun.identity.saml.xmlsig.xmlSigAlgorithm");
      return this.sign(xmlString, idValue, privateKey, cert, sigAlg, digestMethod);
   }

   public Element sign(String xmlString, String idValue, PrivateKey privateKey, X509Certificate cert, @Nullable String signatureAlgorithm, String digestMethod) throws SAML2Exception {
      String c14nMethod = this.getSystemProperty("com.sun.identity.saml.xmlsig.c14nMethod", "http://www.w3.org/2001/10/xml-exc-c14n#");
      String transformAlg = this.getSystemProperty("com.sun.identity.saml.xmlsig.transformAlg", "http://www.w3.org/2001/10/xml-exc-c14n#");
      String classMethod = "SamlSignatureProvider.sign: ";
      if (xmlString == null || xmlString.isEmpty()) {
         SAML2SDKUtils.DEBUG.error(classMethod + "Input xml string is empty.");
         throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("nullInput"));
      }

      if (idValue == null || idValue.isEmpty()) {
         SAML2SDKUtils.DEBUG.error(classMethod + "Id value is empty.");
         throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("nullInput"));
      }

      if (privateKey == null) {
         SAML2SDKUtils.DEBUG.error(classMethod + "Private key is null (keystore not configured correctly).");
         throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("nullInput"));
      }

      Document doc = XMLUtils.toDOMDocument(xmlString, SAML2SDKUtils.DEBUG);
      if (doc == null) {
         throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("errorObtainingElement"));
      }

      Element root = doc.getDocumentElement();
      XMLSignatureFactory signFactory = XMLSignatureFactory.getInstance("DOM");

      Reference ref;
      try {
         List<Transform> transforms = new ArrayList<>(2);
         transforms.add(signFactory.newTransform("http://www.w3.org/2000/09/xmldsig#enveloped-signature", (TransformParameterSpec)null));
         transforms.add(signFactory.newTransform(transformAlg, (TransformParameterSpec)null));
         ref = signFactory.newReference("#" + idValue, signFactory.newDigestMethod(digestMethod, null), transforms, null, null);
      } catch (NoSuchAlgorithmException | InvalidAlgorithmParameterException e) {
         throw new SAML2Exception(e);
      }

      if (StringUtils.isBlank(signatureAlgorithm)) {
         if (privateKey.getAlgorithm().equalsIgnoreCase("DSA")) {
            signatureAlgorithm = "http://www.w3.org/2000/09/xmldsig#dsa-sha1";
         } else if (privateKey.getAlgorithm().equalsIgnoreCase("RSA")) {
            signatureAlgorithm = "http://www.w3.org/2000/09/xmldsig#rsa-sha1";
         }
      }

      SignedInfo signedInfo;
      try {
         signedInfo = signFactory.newSignedInfo(
            signFactory.newCanonicalizationMethod(c14nMethod, (C14NMethodParameterSpec)null),
            signFactory.newSignatureMethod(signatureAlgorithm, null),
            Collections.singletonList(ref)
         );
      } catch (NoSuchAlgorithmException | InvalidAlgorithmParameterException e) {
         throw new SAML2Exception(e);
      }

      DOMSignContext domSignContext = new DOMSignContext(privateKey, doc.getDocumentElement());
      Node firstChild = root.getFirstChild();

      while (firstChild != null && (firstChild.getLocalName() == null || !firstChild.getLocalName().equals("Issuer"))) {
         firstChild = firstChild.getNextSibling();
      }

      Node nextSibling = null;
      if (firstChild != null) {
         nextSibling = firstChild.getNextSibling();
      }

      if (nextSibling != null) {
         domSignContext.setNextSibling(nextSibling);
      }

      domSignContext.setIdAttributeNS(root, null, "ID");
      domSignContext.setDefaultNamespacePrefix("ds");
      KeyInfo keyInfo = null;
      if (cert != null) {
         KeyInfoFactory keyInfoFactory = signFactory.getKeyInfoFactory();
         List x509Content = Collections.singletonList(cert);
         X509Data x509Data = keyInfoFactory.newX509Data(x509Content);
         keyInfo = keyInfoFactory.newKeyInfo(Collections.singletonList(x509Data));
      }

      XMLSignature signature = signFactory.newXMLSignature(signedInfo, keyInfo);

      try {
         signature.sign(domSignContext);
      } catch (MarshalException | XMLSignatureException e) {
         throw new SAML2Exception(e);
      }

      Node signatureElement = root.getFirstChild();

      while (signatureElement != null && (signatureElement.getLocalName() == null || !signatureElement.getLocalName().equals("Signature"))) {
         signatureElement = signatureElement.getNextSibling();
      }

      if (signatureElement == null) {
         throw new SAML2Exception("no element with name \"Signature\" in result.");
      } else {
         return (Element)signatureElement;
      }
   }

   @Override
   public boolean verify(String xmlString, String idValue, List<X509Certificate> senderCerts) throws SAML2Exception {
      String classMethod = "SamlSignatureProvider.verify: ";
      if (xmlString != null && !xmlString.isEmpty() && idValue != null && !idValue.isEmpty()) {
         XMLSignatureFactory signatureFactory = XMLSignatureFactory.getInstance("DOM");
         Document doc = XMLUtils.toDOMDocument(xmlString, SAML2SDKUtils.DEBUG);
         if (doc == null) {
            throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("errorObtainingElement"));
         }

         XPath xpath = this.getXPath();

         Node sigElement;
         try {
            sigElement = (Node)xpath.evaluate("//ds:Signature[1]", doc, XPathConstants.NODE);
         } catch (XPathExpressionException ex) {
            throw new SAML2Exception(ex);
         }

         Element refElement;
         try {
            refElement = (Element)xpath.evaluate("//ds:Reference[1]", doc, XPathConstants.NODE);
         } catch (XPathExpressionException ex) {
            throw new SAML2Exception(ex);
         }

         String refUri = refElement.getAttribute("URI");
         String signedId = ((Element)sigElement.getParentNode()).getAttribute("ID");
         if (!refUri.substring(1).equals(signedId)) {
            SAML2SDKUtils.DEBUG.error(classMethod + "Signature reference ID does not match with element ID");
            throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("signVerifyError"));
         }

         if (!refUri.substring(1).equals(idValue)) {
            SAML2SDKUtils.DEBUG.error(classMethod + "Signature reference ID: " + refUri + " does not match with given ID value: " + idValue);
            throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("signVerifyError"));
         }

         SAML2Exception latestException = null;

         for (X509Certificate cert : senderCerts) {
            try {
               boolean success = this.verifyWithOneCertificate(cert, sigElement, doc, signatureFactory, classMethod);
               if (success) {
                  return true;
               }
            } catch (SAML2Exception e) {
               latestException = e;
            }
         }

         if (latestException != null) {
            throw latestException;
         } else {
            return false;
         }
      } else {
         SAML2SDKUtils.DEBUG.error(classMethod + "Either input xmlString or idValue is null.");
         throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("nullInput"));
      }
   }

   private boolean verifyWithOneCertificate(X509Certificate senderCert, Node sigElement, Document doc, XMLSignatureFactory signatureFactory, String classMethod) throws SAML2Exception {
      DOMValidateContext valContext = createSecureDomValidateContext(this.getAllowedSignatureAlgorithms(), senderCert, sigElement);
      valContext.putNamespacePrefix("http://www.w3.org/2000/09/xmldsig#", "ds");
      valContext.setIdAttributeNS(doc.getDocumentElement(), null, "ID");

      XMLSignature signature;
      try {
         signature = signatureFactory.unmarshalXMLSignature(valContext);
      } catch (MarshalException e) {
         throw new SAML2Exception(e);
      }

      boolean result;
      try {
         result = signature.validate(valContext);
      } catch (XMLSignatureException e) {
         SAML2SDKUtils.DEBUG.error(classMethod + "Signature verification failed: " + e.getMessage());
         return false;
      }

      if (!result) {
         SAML2SDKUtils.DEBUG.error(classMethod + "Signature verification failed.");
         return false;
      }

      if (SAML2SDKUtils.DEBUG.messageEnabled()) {
         SAML2SDKUtils.DEBUG.message(classMethod + "Signature verification successful.");
      }

      return true;
   }

   private XPath getXPath() {
      NamespaceContext namespaceContext = new SimpleNamespaceContext(Collections.singletonMap("ds", "http://www.w3.org/2000/09/xmldsig#"));
      XPathFactory factory = XPathFactory.newInstance();
      XPath xpath = factory.newXPath();
      xpath.setNamespaceContext(namespaceContext);
      return xpath;
   }

   @Nonnull
   private static DOMValidateContext createSecureDomValidateContext(
      List<SignatureAlgorithm> signatureAlgorithmList, X509Certificate senderCert, Node sigElement
   ) {
      DOMValidateContext context = new DOMValidateContext(new SamlSignatureProvider.SamlKeySelector(senderCert, signatureAlgorithmList), sigElement);
      context.setProperty("org.jcp.xml.dsig.secureValidation", Boolean.TRUE);
      return context;
   }

   static {
      Init.init();
   }

   private static class SamlKeySelector extends KeySelector {
      private final X509Certificate cert;
      private final List<SignatureAlgorithm> allowedAlgorithms;

      @Override
      public KeySelectorResult select(KeyInfo keyInfo, Purpose purpose, AlgorithmMethod method, XMLCryptoContext context) throws KeySelectorException {
         certificateToUseForValidation = null;
         if (keyInfo != null) {
            label38:
            for (XMLStructure info : keyInfo.getContent()) {
               if (info instanceof X509Data x509Data) {
                  for (Object o : x509Data.getContent()) {
                     if (o instanceof X509Certificate certificateToUseForValidation) {
                        if (!this.cert.equals(certificateToUseForValidation)) {
                           throw new KeySelectorException("certificate in assertion differs from given certificate.");
                        }
                        break label38;
                     }
                  }
               }
            }
         }

         if (certificateToUseForValidation == null) {
            certificateToUseForValidation = this.cert;
         }

         PublicKey key = certificateToUseForValidation.getPublicKey();
         if (this.isSignatureAlgorithmValid(method.getAlgorithm(), key.getAlgorithm())) {
            return () -> key;
         } else {
            throw new KeySelectorException("certificate public key algorithm does not match signing algorithm!");
         }
      }

      private boolean isSignatureAlgorithmValid(String algURI, String keyAlgName) {
         if (this.allowedAlgorithms.stream().noneMatch(alg -> alg.getUri().equals(algURI))) {
            SamlSignatureProvider.LOG
               .warn(
                  LogTypes.CONFIGURATION_ERROR,
                  "Algorithm \""
                     + LogUtil.quote(algURI)
                     + " is not allowed by configuration. Allowed algorithms are: "
                     + this.allowedAlgorithms.stream().map(SignatureAlgorithm::getUri).toList()
               );
            return false;
         }

         if (keyAlgName.equalsIgnoreCase("DSA") && Strings.CI.contains(algURI, "#dsa-")) {
            return true;
         }

         if (keyAlgName.equalsIgnoreCase("RSA") && Strings.CI.contains(algURI, "#rsa-")) {
            return true;
         }

         SamlSignatureProvider.LOG
            .warn(LogTypes.CONFIGURATION_ERROR, "Unsupported combination of algorithm name \"" + keyAlgName + "\" and algorithm URI \"" + algURI + "\"");
         return false;
      }

      @SuppressFBWarnings(justification = "generated code")
      public SamlKeySelector(X509Certificate cert, List<SignatureAlgorithm> allowedAlgorithms) {
         this.cert = cert;
         this.allowedAlgorithms = allowedAlgorithms;
      }
   }
}