package com.sun.identity.saml2.xmlenc;

import com.sun.identity.common.SystemConfigurationUtil;
import com.sun.identity.saml2.common.SAML2Exception;
import com.sun.identity.saml2.common.SAML2SDKUtils;
import com.sun.identity.shared.xml.XMLUtils;
import java.security.Key;
import java.security.NoSuchAlgorithmException;
import java.util.Hashtable;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import org.apache.xml.security.Init;
import org.apache.xml.security.encryption.EncryptedData;
import org.apache.xml.security.encryption.EncryptedKey;
import org.apache.xml.security.encryption.XMLCipher;
import org.apache.xml.security.encryption.XMLEncryptionException;
import org.w3c.dom.Document;
import org.w3c.dom.Element;
import org.w3c.dom.Node;
import org.w3c.dom.NodeList;

public final class FMEncProvider implements EncProvider {
   private static final Hashtable<String, SecretKey> CACHED_KEYS = new Hashtable<>();
   private static boolean encryptedKeyInKeyInfo = true;

   @Override
   public Element encrypt(
      String xmlString, Key recipientPublicKey, String dataEncAlgorithm, int dataEncStrength, String recipientEntityID, String outerElementName
   ) throws SAML2Exception {
      return this.encrypt(xmlString, recipientPublicKey, null, dataEncAlgorithm, dataEncStrength, recipientEntityID, outerElementName);
   }

   @Override
   public Element encrypt(
      String xmlString,
      Key recipientPublicKey,
      SecretKey secretKey,
      String dataEncAlgorithm,
      int dataEncStrength,
      String recipientEntityID,
      String outerElementName
   ) throws SAML2Exception {
      String classMethod = "FMEncProvider.encrypt: ";
      if (xmlString == null
         || xmlString.isEmpty()
         || recipientPublicKey == null
         || dataEncAlgorithm == null
         || dataEncAlgorithm.isEmpty()
         || outerElementName == null
         || outerElementName.isEmpty()) {
         SAML2SDKUtils.DEBUG.error(classMethod + "Null input parameter(s).");
         throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("nullInput"));
      }

      if (!dataEncAlgorithm.equals("http://www.w3.org/2001/04/xmlenc#aes128-cbc")
         && !dataEncAlgorithm.equals("http://www.w3.org/2001/04/xmlenc#aes192-cbc")
         && !dataEncAlgorithm.equals("http://www.w3.org/2001/04/xmlenc#aes256-cbc")
         && !dataEncAlgorithm.equals("http://www.w3.org/2001/04/xmlenc#tripledes-cbc")) {
         throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("unsupportedKeyAlg"));
      }

      if ((!dataEncAlgorithm.equals("http://www.w3.org/2001/04/xmlenc#aes128-cbc") || dataEncStrength == 128)
         && (!dataEncAlgorithm.equals("http://www.w3.org/2001/04/xmlenc#aes192-cbc") || dataEncStrength == 192)
         && (!dataEncAlgorithm.equals("http://www.w3.org/2001/04/xmlenc#aes256-cbc") || dataEncStrength == 256)) {
         Document doc = XMLUtils.toDOMDocument(xmlString, SAML2SDKUtils.DEBUG);
         if (doc == null) {
            throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("errorObtainingElement"));
         }

         if (dataEncStrength <= 0) {
            dataEncStrength = 128;
         }

         Element rootElement = doc.getDocumentElement();
         if (rootElement == null) {
            SAML2SDKUtils.DEBUG.error(classMethod + "Empty document.");
            throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("emptyDoc"));
         }

         if (secretKey == null) {
            if (recipientEntityID != null) {
               if (CACHED_KEYS.containsKey(recipientEntityID)) {
                  secretKey = CACHED_KEYS.get(recipientEntityID);
               } else {
                  secretKey = this.generateSecretKey(dataEncAlgorithm, dataEncStrength);
                  CACHED_KEYS.put(recipientEntityID, secretKey);
               }
            } else {
               secretKey = this.generateSecretKey(dataEncAlgorithm, dataEncStrength);
            }

            if (secretKey == null) {
               throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("errorGenerateKey"));
            }
         }

         XMLCipher cipher = null;
         String publicKeyEncAlg = recipientPublicKey.getAlgorithm();

         try {
            cipher = switch (publicKeyEncAlg) {
               case "RSA" -> XMLCipher.getInstance("http://www.w3.org/2001/04/xmlenc#rsa-1_5");
               case "DESede" -> XMLCipher.getInstance("http://www.w3.org/2001/04/xmlenc#kw-tripledes");
               case "AES" -> XMLCipher.getInstance("http://www.w3.org/2001/04/xmlenc#kw-aes128");
               default -> throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("unsupportedKeyAlg"));
            };
         } catch (XMLEncryptionException xe1) {
            SAML2SDKUtils.DEBUG.error(classMethod + "Unable to obtain cipher with public key algorithm.", xe1);
            throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("noCipherForPublicKeyAlg"));
         }

         try {
            cipher.init(3, recipientPublicKey);
         } catch (XMLEncryptionException xe2) {
            SAML2SDKUtils.DEBUG.error(classMethod + "Failed to initialize cipher with public key", xe2);
            throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("failedInitCipherWithPublicKey"));
         }

         EncryptedKey encryptedKey = null;

         try {
            encryptedKey = cipher.encryptKey(doc, secretKey);
         } catch (XMLEncryptionException xe3) {
            SAML2SDKUtils.DEBUG.error(classMethod + "Failed to encrypt secret key with public key", xe3);
            throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("failedEncryptingSecretKeyWithPublicKey"));
         }

         try {
            cipher = XMLCipher.getInstance(dataEncAlgorithm);
         } catch (XMLEncryptionException xe4) {
            SAML2SDKUtils.DEBUG.error(classMethod + "Failed to obtain a cipher for data encryption algorithm" + dataEncAlgorithm, xe4);
            throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("cipherNotAvailableForDataEncAlg"));
         }

         try {
            cipher.init(1, secretKey);
         } catch (XMLEncryptionException xe5) {
            SAML2SDKUtils.DEBUG.error(classMethod + "Failed to initialize cipher with secret key.", xe5);
            throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("failedInitCipherWithSecretKey"));
         }

         Document resultDoc = null;

         try {
            resultDoc = cipher.doFinal(doc, rootElement);
         } catch (Exception e) {
            SAML2SDKUtils.DEBUG.error(classMethod + "Failed to do the final data encryption.", e);
            throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("failedEncryptingData"));
         }

         Element ek = null;

         try {
            ek = cipher.martial(doc, encryptedKey);
         } catch (Exception xe6) {
            SAML2SDKUtils.DEBUG.error(classMethod + "Failed to martial the encrypted key", xe6);
            throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("failedMartialingEncryptedKey"));
         }

         String outerElemNS = "urn:oasis:names:tc:SAML:2.0:assertion";
         String outerElemPrefix = "saml";
         if (outerElementName.equals("NewEncryptedID")) {
            outerElemNS = "urn:oasis:names:tc:SAML:2.0:protocol";
            outerElemPrefix = "samlp";
         }

         Element outerElement = resultDoc.createElementNS(outerElemNS, outerElemPrefix + ":" + outerElementName);
         outerElement.setAttributeNS("http://www.w3.org/2000/xmlns/", "xmlns:" + outerElemPrefix, outerElemNS);
         Element ed = resultDoc.getDocumentElement();
         resultDoc.replaceChild(outerElement, ed);
         outerElement.appendChild(ed);
         if (encryptedKeyInKeyInfo) {
            Element dsElement = resultDoc.createElementNS("http://www.w3.org/2000/09/xmldsig#", "ds:KeyInfo");
            dsElement.setAttributeNS("http://www.w3.org/2000/xmlns/", "xmlns:ds", "http://www.w3.org/2000/09/xmldsig#");
            dsElement.appendChild(ek);
            NodeList nl = ed.getElementsByTagNameNS("http://www.w3.org/2001/04/xmlenc#", "CipherData");
            if (nl == null || nl.getLength() == 0) {
               SAML2SDKUtils.DEBUG.error(classMethod + "Unable to find required xenc:CipherData Element.");
               throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("failedEncryptingData"));
            }

            Element cipherDataElement = (Element)nl.item(0);
            ed.insertBefore(dsElement, cipherDataElement);
         } else {
            outerElement.appendChild(ek);
         }

         return resultDoc.getDocumentElement();
      } else {
         SAML2SDKUtils.DEBUG.error(classMethod + "Data encryption algorithm " + dataEncAlgorithm + "and strength " + dataEncStrength + " mismatch.");
         throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("algSizeMismatch"));
      }
   }

   @Override
   public SecretKey getSecretKey(String xmlString, Key recipientPrivateKey) throws SAML2Exception {
      String classMethod = "FMEncProvider.getSecretKey: ";
      if (SAML2SDKUtils.DEBUG.messageEnabled()) {
         SAML2SDKUtils.DEBUG.message(classMethod + "Entering ...");
      }

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

         Element rootElement = doc.getDocumentElement();
         if (rootElement == null) {
            SAML2SDKUtils.DEBUG.error(classMethod + "Empty document.");
            throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("emptyDoc"));
         }

         Element firstChild = this.getNextElementNode(rootElement.getFirstChild());
         if (firstChild == null) {
            SAML2SDKUtils.DEBUG.error(classMethod + "Missing the EncryptedData element.");
            throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("missingElementEncryptedData"));
         }

         Element secondChild = this.getNextElementNode(firstChild.getNextSibling());
         if (secondChild == null) {
            if (SAML2SDKUtils.DEBUG.messageEnabled()) {
               SAML2SDKUtils.DEBUG.message(classMethod + "looking for encrytion key inside first child.");
            }

            NodeList nl = firstChild.getElementsByTagNameNS("http://www.w3.org/2001/04/xmlenc#", "EncryptedKey");
            if (nl == null || nl.getLength() == 0) {
               SAML2SDKUtils.DEBUG.error(classMethod + "Missing the EncryptedKey element.");
               throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("missingElementEncryptedKey"));
            }

            secondChild = (Element)nl.item(0);
         }

         XMLCipher cipher = null;

         try {
            cipher = XMLCipher.getInstance();
         } catch (XMLEncryptionException xe1) {
            SAML2SDKUtils.DEBUG.error(classMethod + "Unable to get a cipher instance.", xe1);
            throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("noCipher"));
         }

         try {
            cipher.init(2, null);
         } catch (XMLEncryptionException xe2) {
            SAML2SDKUtils.DEBUG.error(classMethod + "Failed to initialize cipher for decryption mode", xe2);
            throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("failedInitCipherForDecrypt"));
         }

         EncryptedData encryptedData = null;

         try {
            encryptedData = cipher.loadEncryptedData(doc, firstChild);
         } catch (XMLEncryptionException xe3) {
            SAML2SDKUtils.DEBUG.error(classMethod + "Failed to load encrypted data", xe3);
            throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("failedLoadingEncryptedData"));
         }

         EncryptedKey encryptedKey = null;

         try {
            encryptedKey = cipher.loadEncryptedKey(doc, secondChild);
         } catch (XMLEncryptionException xe4) {
            SAML2SDKUtils.DEBUG.error(classMethod + "Failed to load encrypted key", xe4);
            throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("failedLoadingEncryptedKey"));
         }

         if (encryptedKey != null && encryptedData != null) {
            XMLCipher keyCipher = null;

            try {
               keyCipher = XMLCipher.getInstance();
            } catch (XMLEncryptionException xe5) {
               SAML2SDKUtils.DEBUG.error(classMethod + "Failed to get a cipher instance for decrypting secret key.", xe5);
               throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("noCipher"));
            }

            try {
               keyCipher.init(4, recipientPrivateKey);
            } catch (XMLEncryptionException xe6) {
               SAML2SDKUtils.DEBUG.error(classMethod + "Failed to initialize cipher in unwrap mode with private key", xe6);
               throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("noCipherForUnwrap"));
            }

            try {
               return (SecretKey)keyCipher.decryptKey(encryptedKey, encryptedData.getEncryptionMethod().getAlgorithm());
            } catch (XMLEncryptionException xe7) {
               SAML2SDKUtils.DEBUG.error(classMethod + "Failed to decrypt the secret key", xe7);
               throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("failedDecryptingSecretKey"));
            }
         } else {
            return null;
         }
      } else {
         throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("nullInput"));
      }
   }

   @Override
   public Element decrypt(String xmlString, Key recipientPrivateKey) throws SAML2Exception {
      String classMethod = "FMEncProvider.decrypt: ";
      if (SAML2SDKUtils.DEBUG.messageEnabled()) {
         SAML2SDKUtils.DEBUG.message(classMethod + "Entering ...");
      }

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

         Element rootElement = doc.getDocumentElement();
         if (rootElement == null) {
            SAML2SDKUtils.DEBUG.error(classMethod + "Empty document.");
            throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("emptyDoc"));
         }

         Element firstChild = this.getNextElementNode(rootElement.getFirstChild());
         if (firstChild == null) {
            SAML2SDKUtils.DEBUG.error(classMethod + "Missing the EncryptedData element.");
            throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("missingElementEncryptedData"));
         }

         Element secondChild = this.getNextElementNode(firstChild.getNextSibling());
         if (secondChild == null) {
            if (SAML2SDKUtils.DEBUG.messageEnabled()) {
               SAML2SDKUtils.DEBUG.message(classMethod + "looking for encrytion key inside first child.");
            }

            NodeList nl = firstChild.getElementsByTagNameNS("http://www.w3.org/2001/04/xmlenc#", "EncryptedKey");
            if (nl == null || nl.getLength() == 0) {
               SAML2SDKUtils.DEBUG.error(classMethod + "Missing the EncryptedKey element.");
               throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("missingElementEncryptedKey"));
            }

            secondChild = (Element)nl.item(0);
         }

         XMLCipher cipher = null;

         try {
            cipher = XMLCipher.getInstance();
         } catch (XMLEncryptionException xe1) {
            SAML2SDKUtils.DEBUG.error(classMethod + "Unable to get a cipher instance.", xe1);
            throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("noCipher"));
         }

         try {
            cipher.init(2, null);
         } catch (XMLEncryptionException xe2) {
            SAML2SDKUtils.DEBUG.error(classMethod + "Failed to initialize cipher for decryption mode", xe2);
            throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("failedInitCipherForDecrypt"));
         }

         EncryptedData encryptedData = null;

         try {
            encryptedData = cipher.loadEncryptedData(doc, firstChild);
         } catch (XMLEncryptionException xe3) {
            SAML2SDKUtils.DEBUG.error(classMethod + "Failed to load encrypted data", xe3);
            throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("failedLoadingEncryptedData"));
         }

         EncryptedKey encryptedKey = null;

         try {
            encryptedKey = cipher.loadEncryptedKey(doc, secondChild);
         } catch (XMLEncryptionException xe4) {
            SAML2SDKUtils.DEBUG.error(classMethod + "Failed to load encrypted key", xe4);
            throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("failedLoadingEncryptedKey"));
         }

         Document decryptedDoc = null;
         if (encryptedKey != null && encryptedData != null) {
            XMLCipher keyCipher = null;

            try {
               keyCipher = XMLCipher.getInstance();
            } catch (XMLEncryptionException xe5) {
               SAML2SDKUtils.DEBUG.error(classMethod + "Failed to get a cipher instance for decrypting secret key.", xe5);
               throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("noCipher"));
            }

            try {
               keyCipher.init(4, recipientPrivateKey);
            } catch (XMLEncryptionException xe6) {
               SAML2SDKUtils.DEBUG.error(classMethod + "Failed to initialize cipher in unwrap mode with private key", xe6);
               throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("noCipherForUnwrap"));
            }

            Key encryptionKey = null;

            try {
               encryptionKey = keyCipher.decryptKey(encryptedKey, encryptedData.getEncryptionMethod().getAlgorithm());
            } catch (XMLEncryptionException xe7) {
               SAML2SDKUtils.DEBUG.error(classMethod + "Failed to decrypt the secret key", xe7);
               throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("failedDecryptingSecretKey"));
            }

            cipher = null;

            try {
               cipher = XMLCipher.getInstance();
            } catch (XMLEncryptionException xe8) {
               SAML2SDKUtils.DEBUG.error(classMethod + "Failed to get cipher instance for final data decryption.", xe8);
               throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("noCipher"));
            }

            try {
               cipher.init(2, encryptionKey);
            } catch (XMLEncryptionException xe9) {
               SAML2SDKUtils.DEBUG.error(classMethod + "Failed to initialize cipher with secret key.", xe9);
               throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("failedInitCipherForDecrypt"));
            }

            try {
               decryptedDoc = cipher.doFinal(doc, firstChild);
            } catch (Exception e) {
               SAML2SDKUtils.DEBUG.error(classMethod + "Failed to decrypt data.", e);
               throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("failedDecryptingData"));
            }
         }

         Element root = decryptedDoc.getDocumentElement();
         Element child = this.getNextElementNode(root.getFirstChild());
         if (child == null) {
            SAML2SDKUtils.DEBUG.error(classMethod + "decrypted document contains empty element.");
            throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("failedDecryptingData"));
         } else {
            root.removeChild(child);
            decryptedDoc.replaceChild(child, root);
            return decryptedDoc.getDocumentElement();
         }
      } else {
         throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("nullInput"));
      }
   }

   private Element getNextElementNode(Node node) {
      while (node != null) {
         if (node.getNodeType() == 1) {
            return (Element)node;
         }

         node = node.getNextSibling();
      }

      return null;
   }

   private SecretKey generateSecretKey(String algorithm, int keyStrength) throws SAML2Exception {
      KeyGenerator keygen = null;

      try {
         if (!algorithm.equals("http://www.w3.org/2001/04/xmlenc#aes128-cbc")
            && !algorithm.equals("http://www.w3.org/2001/04/xmlenc#aes192-cbc")
            && !algorithm.equals("http://www.w3.org/2001/04/xmlenc#aes256-cbc")) {
            if (!algorithm.equals("http://www.w3.org/2001/04/xmlenc#tripledes-cbc")) {
               throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("unsupportedKeyAlg"));
            }

            keygen = KeyGenerator.getInstance("TripleDES");
         } else {
            keygen = KeyGenerator.getInstance("AES");
         }

         if (keyStrength != 0) {
            keygen.init(keyStrength);
         }
      } catch (NoSuchAlgorithmException ne) {
         throw new SAML2Exception(ne);
      }

      return keygen != null ? keygen.generateKey() : null;
   }

   static {
      Init.init();
      String tmp = SystemConfigurationUtil.getProperty("com.sun.identity.saml.xmlenc.encryptedKeyInKeyInfo");
      if (tmp != null && tmp.equalsIgnoreCase("false")) {
         encryptedKeyInKeyInfo = false;
      }
   }
}