package com.sun.identity.saml2.profile;

import com.airlock.iam.openfedlib.saml2.XjcHelper;
import com.airlock.iam.openfedlib.saml2.servlet.Saml2ResponseBodyWriter;
import com.airlock.iam.opensso.jaxb.metadata.SPSSODescriptorType;
import com.sun.identity.saml.common.SAMLUtils;
import com.sun.identity.saml.xmlsig.KeyProvider;
import com.sun.identity.saml2.assertion.AssertionFactory;
import com.sun.identity.saml2.assertion.Issuer;
import com.sun.identity.saml2.common.IOExceptionThrowingSupplier;
import com.sun.identity.saml2.common.SAML2Exception;
import com.sun.identity.saml2.common.SAML2Repository;
import com.sun.identity.saml2.common.SAML2Utils;
import com.sun.identity.saml2.key.KeyUtil;
import com.sun.identity.saml2.logging.LogUtil;
import com.sun.identity.saml2.meta.SAML2MetaException;
import com.sun.identity.saml2.meta.SAML2MetaUtils;
import com.sun.identity.saml2.protocol.Artifact;
import com.sun.identity.saml2.protocol.ArtifactResolve;
import com.sun.identity.saml2.protocol.ArtifactResponse;
import com.sun.identity.saml2.protocol.ProtocolFactory;
import com.sun.identity.saml2.protocol.Response;
import com.sun.identity.saml2.protocol.Status;
import com.sun.identity.saml2.protocol.StatusCode;
import com.sun.identity.shared.datastruct.OrderedSet;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import jakarta.xml.soap.MessageFactory;
import jakarta.xml.soap.MimeHeader;
import jakarta.xml.soap.MimeHeaders;
import jakarta.xml.soap.SOAPConnection;
import jakarta.xml.soap.SOAPException;
import jakarta.xml.soap.SOAPMessage;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.security.cert.X509Certificate;
import java.util.Date;
import java.util.Enumeration;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.StringTokenizer;
import java.util.logging.Level;
import org.w3c.dom.Element;

public class IDPArtifactResolution {
   private static MessageFactory messageFactory;

   private IDPArtifactResolution() {
   }

   public static void doArtifactResolution(
      HttpServletRequest request, HttpServletResponse response, IOExceptionThrowingSupplier<InputStream> inputStreamSupplier
   ) {
      String classMethod = "IDPArtifactResolution.doArtifactResolution: ";

      try {
         String idpMetaAlias = request.getParameter("metaAlias");
         if (idpMetaAlias == null || idpMetaAlias.trim().isEmpty()) {
            idpMetaAlias = SAML2MetaUtils.getMetaAliasByUri(request.getRequestURI());
         }

         if (idpMetaAlias == null || idpMetaAlias.trim().isEmpty()) {
            if (SAML2Utils.DEBUG.messageEnabled()) {
               SAML2Utils.DEBUG.message(classMethod + "unable to get IDP meta alias from request.");
            }

            String[] data = new String[]{idpMetaAlias};
            LogUtil.error(Level.INFO, "IDP_METADATA_ERROR", data, null);
            SAMLUtils.sendError(response, "nullIDPMetaAlias", SAML2Utils.BUNDLE.getString("nullIDPMetaAlias"));
            return;
         }

         String idpEntityID = null;
         String realm = null;

         try {
            idpEntityID = IDPSSOUtil.META_MANAGER.getEntityByMetaAlias(idpMetaAlias);
            if (idpEntityID == null || idpEntityID.trim().isEmpty()) {
               SAMLUtils.DEBUG.error(classMethod + "Unable to get IDP Entity ID from meta.");
               String[] data = new String[]{idpEntityID};
               LogUtil.error(Level.INFO, "INVALID_IDP", data, null);
               SAMLUtils.sendError(response, "nullIDPEntityID", SAML2Utils.BUNDLE.getString("nullIDPEntityID"));
               return;
            }

            realm = SAML2MetaUtils.getRealmByMetaAlias(idpMetaAlias);
         } catch (SAML2MetaException sme) {
            SAML2Utils.DEBUG.error(classMethod + "Unable to get IDP Entity ID from meta.");
            String[] data = new String[]{idpMetaAlias};
            LogUtil.error(Level.INFO, "IDP_METADATA_ERROR", data, null);
            SAMLUtils.sendError(response, "metaDataError", SAML2Utils.BUNDLE.getString("metaDataError"));
            return;
         }

         try {
            MimeHeaders headers = getHeaders(request);
            InputStream is = inputStreamSupplier.get();
            SOAPMessage msg = messageFactory.createMessage(headers, is);
            SOAPMessage reply = null;
            reply = onMessage(msg, request, realm, idpEntityID);
            if (reply != null) {
               if (reply.saveRequired()) {
                  reply.saveChanges();
               }

               response.setStatus(200);
               putHeaders(reply.getMimeHeaders(), response);
               ByteArrayOutputStream buffer = new ByteArrayOutputStream(8192);
               reply.writeTo(buffer);
               Saml2ResponseBodyWriter.writeBodyWithoutCommittingResponse(buffer, response);
            } else {
               response.setStatus(204);
            }
         } catch (SOAPException ex) {
            SAML2Utils.DEBUG.error(classMethod + "SOAP error", ex);
            String[] data = new String[]{idpEntityID};
            LogUtil.error(Level.INFO, "INVALID_SOAP_MESSAGE", data, null);
            SAMLUtils.sendError(response, "invalidSOAPMessage", SAML2Utils.BUNDLE.getString("invalidSOAPMessage") + " " + ex.getMessage());
         } catch (SAML2Exception se) {
            SAML2Utils.DEBUG.error(classMethod + "SAML2 error", se);
            SAMLUtils.sendError(
               response, "unableToCreateArtifactResponse", SAML2Utils.BUNDLE.getString("unableToCreateArtifactResponse") + " " + se.getMessage()
            );
         }
      } catch (IOException ioe) {
         SAML2Utils.DEBUG.error(classMethod + "I/O rrror", ioe);
      }
   }

   private static SOAPMessage onMessage(SOAPMessage message, HttpServletRequest request, String realm, String idpEntityID) throws SAML2Exception {
      String classMethod = "IDPArtifactResolution.onMessage: ";
      if (SAML2Utils.DEBUG.messageEnabled()) {
         SAML2Utils.DEBUG.message(classMethod + "Entering onMessage().");
      }

      Element reqElem = SAML2Utils.getSamlpElement(message, "ArtifactResolve");
      ArtifactResolve artResolve = ProtocolFactory.getInstance().createArtifactResolve(reqElem);
      if (artResolve == null) {
         if (SAML2Utils.DEBUG.messageEnabled()) {
            SAML2Utils.DEBUG.message(classMethod + "no valid ArtifactResolve node found in SOAP body.");
         }

         return SAML2Utils.createSOAPFault("Client", "noArtifactResolve", null);
      } else {
         String spEntityID = artResolve.getIssuer().getValue();
         SPSSODescriptorType spSSODescriptor = null;

         try {
            spSSODescriptor = IDPSSOUtil.META_MANAGER.getSPSSODescriptor(realm, spEntityID);
         } catch (SAML2MetaException sme) {
            SAML2Utils.DEBUG.error(classMethod, sme);
            spSSODescriptor = null;
         }

         if (spSSODescriptor == null) {
            SAML2Utils.DEBUG.error(classMethod + "Unable to get SP SSO Descriptor from meta.");
            return SAML2Utils.createSOAPFault("Server", "metaDataError", null);
         }

         OrderedSet<String> acsSet = SPSSOFederate.getACSUrl(spSSODescriptor, "urn:oasis:names:tc:SAML:2.0:bindings:HTTP-Artifact");
         String acsURL = acsSet.get(0);
         String isArtifactResolveSigned = SAML2Utils.getAttributeValueFromSSOConfig(realm, idpEntityID, "IDPRole", "wantArtifactResolveSigned");
         if (isArtifactResolveSigned != null && isArtifactResolveSigned.equals("true")) {
            if (!artResolve.isSigned()) {
               SAML2Utils.DEBUG.error(classMethod + "The artifact resolve is not signed when it is expected to be signed.");
               return SAML2Utils.createSOAPFault("Client", "ArtifactResolveNotSigned", null);
            }

            List<X509Certificate> spCert = KeyUtil.getVerificationCerts(spSSODescriptor, spEntityID, "SPRole");
            if (!artResolve.isSignatureValid(spCert)) {
               SAML2Utils.DEBUG.error(classMethod + "artifact resolve verification failed.");
               return SAML2Utils.createSOAPFault("Client", "invalidArtifact", null);
            }

            if (SAML2Utils.DEBUG.messageEnabled()) {
               SAML2Utils.DEBUG.message(classMethod + "artifact resolve signature verification is successful.");
            }
         }

         Artifact art = artResolve.getArtifact();
         if (art == null) {
            SAML2Utils.DEBUG.error(classMethod + "Unable to get an artifact from ArtifactResolve.");
            return SAML2Utils.createSOAPFault("Client", "invalidArtifactSignature", null);
         }

         String artStr = art.getArtifactValue();
         Response res = IDPCache.responsesByArtifacts().remove(artStr);
         String remoteArtURL = null;
         if (res == null) {
            String remoteServiceURL = SAML2Utils.getRemoteServiceURL(art.getMessageHandle());
            if (remoteServiceURL != null) {
               remoteArtURL = remoteServiceURL + SAML2Utils.removeDeployUri(request.getRequestURI());

               try {
                  SOAPConnection con = SAML2Utils.CONNECTION_FACTORY.createConnection();
                  return con.call(message, remoteArtURL);
               } catch (Exception ex) {
                  if (SAML2Utils.DEBUG.messageEnabled()) {
                     SAML2Utils.DEBUG.message(classMethod + "unable to forward request to remote server. remote url = " + remoteArtURL, ex);
                  }
               }
            }

            try {
               if (SAML2Utils.isSAML2FailOverEnabled()) {
                  if (SAML2Utils.DEBUG.messageEnabled()) {
                     SAML2Utils.DEBUG.message("Artifact=" + artStr);
                  }

                  String tmp = (String)SAML2Repository.getInstance().retrieve(artStr);
                  res = ProtocolFactory.getInstance().createResponse(tmp);
               }
            } catch (SAML2Exception e) {
               SAML2Utils.DEBUG.error(classMethod + "DB ERROR!!!");
            }
         }

         if (res == null) {
            return SAML2Utils.createSOAPFault("Server", "UnableToFindResponse", null);
         }

         try {
            if (SAML2Utils.isSAML2FailOverEnabled()) {
               SAML2Repository.getInstance().delete(artStr);
            }
         } catch (SAML2Exception e) {
            SAML2Utils.DEBUG.error(classMethod + "DB ERROR!!!");
         }

         Map<String, String> props = new HashMap<>();
         String nameIDString = SAML2Utils.getNameIDStringFromResponse(res);
         if (nameIDString != null) {
            props.put("NameID", nameIDString);
         }

         boolean signAssertion = XjcHelper.mapBooleanWrapperToPrimitiveWithFalseAsDefault(spSSODescriptor.isWantAssertionsSigned());
         if (signAssertion && SAML2Utils.DEBUG.messageEnabled()) {
            SAML2Utils.DEBUG.message(classMethod + "signing the assertion.");
         }

         IDPSSOUtil.signAndEncryptResponseComponents(realm, spEntityID, idpEntityID, res, signAssertion);
         ArtifactResponse artResponse = ProtocolFactory.getInstance().createArtifactResponse();
         Status status = ProtocolFactory.getInstance().createStatus();
         StatusCode statusCode = ProtocolFactory.getInstance().createStatusCode();
         statusCode.setValue("urn:oasis:names:tc:SAML:2.0:status:Success");
         status.setStatusCode(statusCode);
         Issuer issuer = AssertionFactory.getInstance().createIssuer();
         issuer.setValue(idpEntityID);
         artResponse.setStatus(status);
         artResponse.setID(SAML2Utils.generateID());
         artResponse.setInResponseTo(artResolve.getID());
         artResponse.setVersion("2.0");
         artResponse.setIssueInstant(new Date());
         artResponse.setAny(res.toXMLString(true, true));
         artResponse.setIssuer(issuer);
         artResponse.setDestination(acsURL);
         String wantArtifactResponseSigned = SAML2Utils.getAttributeValueFromSSOConfig(realm, spEntityID, "SPRole", "wantArtifactResponseSigned");
         if (wantArtifactResponseSigned != null && wantArtifactResponseSigned.equals("true")) {
            KeyProvider kp = KeyUtil.getKeyProviderInstance();
            if (kp == null) {
               SAML2Utils.DEBUG.error(classMethod + "Unable to get a key provider instance.");
               return SAML2Utils.createSOAPFault("Server", "nullKeyProvider", null);
            }

            String idpSignCertAlias = SAML2Utils.getSigningCertAlias(realm, idpEntityID, "IDPRole");
            if (idpSignCertAlias == null) {
               SAML2Utils.DEBUG.error(classMethod + "Unable to get the hosted IDP signing certificate alias.");
               return SAML2Utils.createSOAPFault("Server", "missingSigningCertAlias", null);
            }

            artResponse.sign(kp.getPrivateKey(idpSignCertAlias), kp.getX509Certificate(idpSignCertAlias));
         }

         String str = artResponse.toXMLString(true, true);
         String[] logdata = new String[]{idpEntityID, artStr, str};
         LogUtil.access(Level.INFO, "ARTIFACT_RESPONSE", logdata, null, props);
         if (str != null) {
            if (SAML2Utils.DEBUG.messageEnabled()) {
               SAML2Utils.DEBUG.message(classMethod + "ArtifactResponse message:\n" + str);
            }
         } else if (SAML2Utils.DEBUG.messageEnabled()) {
            SAML2Utils.DEBUG.message(classMethod + "Unable to print ArtifactResponse message.");
         }

         SOAPMessage msg = null;

         try {
            return SAML2Utils.createSOAPMessage(str, false);
         } catch (SOAPException se) {
            SAML2Utils.DEBUG.error(classMethod + "Unable to create a SOAPMessage and add a document ", se);
            return SAML2Utils.createSOAPFault("Server", "unableToCreateSOAPMessage", null);
         }
      }
   }

   private static MimeHeaders getHeaders(HttpServletRequest req) {
      Enumeration<String> enumerator = req.getHeaderNames();
      MimeHeaders headers = new MimeHeaders();

      while (enumerator.hasMoreElements()) {
         String headerName = enumerator.nextElement();
         String headerValue = req.getHeader(headerName);
         StringTokenizer values = new StringTokenizer(headerValue, ",");

         while (values.hasMoreTokens()) {
            headers.addHeader(headerName, values.nextToken().trim());
         }
      }

      return headers;
   }

   private static void putHeaders(MimeHeaders headers, HttpServletResponse res) {
      Iterator<MimeHeader> it = headers.getAllHeaders();

      while (it.hasNext()) {
         MimeHeader header = it.next();
         if (!header.getName().equalsIgnoreCase("Content-Length")) {
            String[] values = headers.getHeader(header.getName());
            if (values.length == 1) {
               res.setHeader(header.getName(), header.getValue());
            } else {
               StringBuilder concat = new StringBuilder();

               for (int i = 0; i < values.length; concat.append(values[i++])) {
                  if (i != 0) {
                     concat.append(',');
                  }
               }

               res.setHeader(header.getName(), concat.toString());
            }
         }
      }
   }

   static {
      try {
         messageFactory = MessageFactory.newInstance();
      } catch (SOAPException se) {
         SAML2Utils.DEBUG.error("Unable to obtain SOAPFactory.", se);
      }
   }
}