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.openfedlib.saml2.state.Saml2DataMapProvider;
import com.airlock.iam.opensso.jaxb.metadata.IDPSSODescriptorType;
import com.airlock.iam.opensso.jaxb.metadata.SPSSODescriptorType;
import com.airlock.iam.opensso.jaxb.metadata.SingleSignOnServiceElement;
import com.sun.identity.common.SystemConfigurationUtil;
import com.sun.identity.plugin.session.SessionException;
import com.sun.identity.plugin.session.SessionManager;
import com.sun.identity.plugin.session.SessionProvider;
import com.sun.identity.saml.common.SAMLUtils;
import com.sun.identity.saml2.assertion.AuthnContext;
import com.sun.identity.saml2.common.QuerySignatureUtil;
import com.sun.identity.saml2.common.SAML2Exception;
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.plugins.IDPAuthnContextInfo;
import com.sun.identity.saml2.plugins.IDPAuthnContextMapper;
import com.sun.identity.saml2.plugins.IDPECPSessionMapper;
import com.sun.identity.saml2.protocol.AuthnRequest;
import com.sun.identity.saml2.protocol.NameIDPolicy;
import com.sun.identity.saml2.protocol.ProtocolFactory;
import com.sun.identity.saml2.protocol.Response;
import com.sun.identity.shared.encode.Base64;
import com.sun.identity.shared.encode.URLEncDec;
import com.sun.identity.shared.xml.XMLUtils;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import jakarta.xml.soap.MimeHeaders;
import jakarta.xml.soap.SOAPException;
import jakarta.xml.soap.SOAPMessage;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.security.cert.X509Certificate;
import java.util.List;
import java.util.Set;
import java.util.logging.Level;
import org.w3c.dom.Document;
import org.w3c.dom.Element;

public class IDPSSOFederate {
   private static final String REQ_ID = "ReqID";

   private IDPSSOFederate() {
   }

   public static void doSSOFederate(HttpServletRequest request, HttpServletResponse response) {
      doSSOFederate(request, response, false);
   }

   private static void doSSOFederate(HttpServletRequest request, HttpServletResponse response, boolean isFromECP) {
      String classMethod = "IDPSSOFederate.doSSOFederate: ";
      SPSSODescriptorType spSSODescriptor = null;

      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.");
            }

            sendError(response, "Client", "IDPMetaAliasNotFound", null, isFromECP);
            return;
         }

         String idpEntityID = null;
         String realm = null;

         try {
            if (IDPSSOUtil.META_MANAGER == null) {
               SAML2Utils.DEBUG.error(classMethod + "Unable to get meta manager.");
               sendError(response, "Server", "errorMetaManager", null, isFromECP);
               return;
            }

            idpEntityID = IDPSSOUtil.META_MANAGER.getEntityByMetaAlias(idpMetaAlias);
            if (idpEntityID == null || idpEntityID.trim().isEmpty()) {
               SAML2Utils.DEBUG.error(classMethod + "Unable to get IDP Entity ID from meta.");
               String[] data = new String[]{idpEntityID};
               LogUtil.error(Level.INFO, "INVALID_IDP", data, null);
               sendError(response, "Client", "nullIDPEntityID", null, isFromECP);
               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);
            sendError(response, "Server", "nullIDPEntityID", sme.getMessage(), isFromECP);
            return;
         }

         String reqID = request.getParameter("ReqID");
         if (reqID != null && reqID.trim().isEmpty()) {
            reqID = null;
         }

         Object session = null;
         AuthnRequest authnReq = null;
         if (reqID == null) {
            String binding = "urn:oasis:names:tc:SAML:2.0:bindings:HTTP-Redirect";
            if (request.getMethod().equals("POST")) {
               binding = "urn:oasis:names:tc:SAML:2.0:bindings:HTTP-POST";
            }

            authnReq = getAuthnRequest(request, isFromECP, binding);
            if (authnReq == null) {
               sendError(response, "Client", "InvalidSAMLRequest", null, isFromECP);
               return;
            }

            String spEntityID = authnReq.getIssuer().getValue();

            try {
               String authnRequestStr = authnReq.toXMLString();
               String[] logdata = new String[]{spEntityID, idpMetaAlias, authnRequestStr};
               String logId = isFromECP ? "RECEIVED_AUTHN_REQUEST_ECP" : "RECEIVED_AUTHN_REQUEST";
               LogUtil.access(Level.INFO, logId, logdata, null);
            } catch (SAML2Exception saml2ex) {
               SAML2Utils.DEBUG.error(classMethod, saml2ex);
               sendError(response, "Client", "InvalidSAMLRequest", saml2ex.getMessage(), isFromECP);
               return;
            }

            if (!SAML2Utils.isSourceSiteValid(authnReq.getIssuer(), realm, idpEntityID)) {
               if (SAML2Utils.DEBUG.warningEnabled()) {
                  SAML2Utils.DEBUG.warning(classMethod + "Issuer in Request is not valid.");
               }

               sendError(response, "Client", "InvalidSAMLRequest", null, isFromECP);
               return;
            }

            IDPSSODescriptorType idpSSODescriptor = null;

            try {
               idpSSODescriptor = IDPSSOUtil.META_MANAGER.getIDPSSODescriptor(realm, idpEntityID);
            } catch (SAML2MetaException sme) {
               SAML2Utils.DEBUG.error(classMethod, sme);
            }

            if (idpSSODescriptor == null) {
               SAML2Utils.DEBUG.error(classMethod + "Unable to get IDP SSO Descriptor from meta.");
               sendError(response, "Server", "metaDataError", null, isFromECP);
               return;
            }

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

            if (isFromECP || XjcHelper.mapBooleanWrapperToPrimitiveWithFalseAsDefault(idpSSODescriptor.isWantAuthnRequestsSigned())) {
               if (spEntityID == null || spEntityID.trim().isEmpty()) {
                  sendError(response, "Client", "InvalidSAMLRequest", null, isFromECP);
                  return;
               }

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

               List<X509Certificate> spCert = KeyUtil.getVerificationCerts(spSSODescriptor, spEntityID, "SPRole");

               try {
                  boolean isSignatureOK;
                  if (isFromECP) {
                     isSignatureOK = authnReq.isSignatureValid(spCert);
                  } else {
                     String method = request.getMethod();
                     if (method.equals("POST")) {
                        isSignatureOK = authnReq.isSignatureValid(spCert);
                     } else {
                        String queryString = request.getQueryString();
                        isSignatureOK = QuerySignatureUtil.verify(queryString, spCert);
                     }
                  }

                  if (!isSignatureOK) {
                     SAML2Utils.DEBUG.error(classMethod + "authn request verification failed.");
                     sendError(response, "Client", "invalidSignInRequest", null, isFromECP);
                     return;
                  }

                  if (!isFromECP) {
                     List<SingleSignOnServiceElement> ssoServiceList = idpSSODescriptor.getSingleSignOnService();
                     String ssoURL = SPSSOFederate.getSSOURL(ssoServiceList, binding);
                     if (!SAML2Utils.verifyDestination(authnReq.getDestination(), ssoURL)) {
                        SAML2Utils.DEBUG.error(classMethod + "authn request destination verification failed.");
                        sendError(response, "Client", "invalidDestination", null, isFromECP);
                        return;
                     }
                  }
               } catch (SAML2Exception se) {
                  SAML2Utils.DEBUG.error(classMethod + "authn request verification failed.", se);
                  sendError(response, "Client", "invalidSignInRequest", null, isFromECP);
                  return;
               }

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

            reqID = authnReq.getID();
            if (SAML2Utils.DEBUG.messageEnabled()) {
               SAML2Utils.DEBUG.message(classMethod + "request id=" + reqID);
            }

            if (reqID == null) {
               SAML2Utils.DEBUG.error(classMethod + "Request id is null");
               sendError(response, "Client", "InvalidSAMLRequestID", null, isFromECP);
               return;
            }

            if (isFromECP) {
               try {
                  IDPECPSessionMapper idpECPSessonMapper = IDPSSOUtil.getIDPECPSessionMapper(realm, idpEntityID);
                  session = idpECPSessonMapper.getSession(request, response);
               } catch (SAML2Exception se) {
                  if (SAML2Utils.DEBUG.messageEnabled()) {
                     SAML2Utils.DEBUG.message(classMethod + "Unable to retrieve user session.");
                  }
               }
            } else {
               try {
                  session = SessionManager.getProvider().getSession(request);
               } catch (SessionException se) {
                  if (SAML2Utils.DEBUG.messageEnabled()) {
                     SAML2Utils.DEBUG.message(classMethod + "Unable to retrieve user session.");
                  }
               }
            }

            IDPAuthnContextMapper idpAuthnContextMapper = null;

            try {
               idpAuthnContextMapper = IDPSSOUtil.getIDPAuthnContextMapper(realm, idpEntityID);
            } catch (SAML2Exception sme) {
               SAML2Utils.DEBUG.error(classMethod, sme);
            }

            if (idpAuthnContextMapper == null) {
               SAML2Utils.DEBUG.error(classMethod + "Unable to get IDPAuthnContextMapper from meta.");
               sendError(response, "Server", "metaDataError", null, isFromECP);
               return;
            }

            IDPAuthnContextInfo idpAuthnContextInfo = null;

            try {
               idpAuthnContextInfo = idpAuthnContextMapper.getIDPAuthnContextInfo(authnReq, idpEntityID, realm);
            } catch (SAML2Exception sme) {
               SAML2Utils.DEBUG.error(classMethod, sme);
            }

            if (idpAuthnContextInfo == null) {
               if (SAML2Utils.DEBUG.messageEnabled()) {
                  SAML2Utils.DEBUG.message(classMethod + "Unable to find valid AuthnContext. Sending error Response.");
               }

               try {
                  Response res = SAML2Utils.getErrorResponse(
                     authnReq, "urn:oasis:names:tc:SAML:2.0:status:Requester", "urn:oasis:names:tc:SAML:2.0:status:NoAuthnContext", null, idpEntityID
                  );
                  StringBuffer returnedBinding = new StringBuffer();
                  String acsURL = IDPSSOUtil.getACSurl(spEntityID, idpEntityID, realm, authnReq, request, returnedBinding);
                  String acsBinding = returnedBinding.toString();
                  IDPSSOUtil.sendResponse(request, response, acsBinding, spEntityID, idpEntityID, idpMetaAlias, realm, null, acsURL, res, session);
               } catch (SAML2Exception sme) {
                  SAML2Utils.DEBUG.error(classMethod, sme);
                  sendError(response, "Server", "metaDataError", null, isFromECP);
               }

               return;
            }

            if (Boolean.TRUE.equals(authnReq.isForceAuthn()) && !Boolean.TRUE.equals(authnReq.isPassive()) && session != null) {
               try {
                  SessionManager.getProvider().invalidateSession(session, request, response);
               } catch (SessionException ssoe) {
                  SAML2Utils.DEBUG.error(classMethod + "Unable to invalidate the sso session.");
               }

               session = null;
            }

            String relayState = request.getParameter("RelayState");
            AuthnContext matchingAuthnContext = idpAuthnContextInfo.getAuthnContext();
            if (session == null) {
               synchronized (Saml2DataMapProvider.DataMapId.IDP_AUTHN_REQUEST_CACHE) {
                  IDPCache.authnRequestCache().put(reqID, authnReq);
               }

               synchronized (Saml2DataMapProvider.DataMapId.IDP_IDP_AUTHN_CONTEXT_CACHE) {
                  IDPCache.idpAuthnContextCache().put(reqID, matchingAuthnContext);
               }

               if (relayState != null && !relayState.trim().isEmpty()) {
                  IDPCache.relayStateCache().put(reqID, relayState);
               }

               try {
                  redirectToFlowAuthUri(response, reqID);
               } catch (IOException | SAML2Exception ioe) {
                  SAML2Utils.DEBUG.error(classMethod + "Unable to redirect to authentication.", ioe);
                  sendError(response, "Server", "UnableToRedirectToAuth", ioe.getMessage(), isFromECP);
               }
            } else {
               if (SAML2Utils.DEBUG.messageEnabled()) {
                  SAML2Utils.DEBUG.message(classMethod + "Session is valid");
               }

               String sessionIndex = IDPSSOUtil.getSessionIndex(session);
               boolean sessionUpgrade;
               if (isForceSessionUpgrade(session)) {
                  sessionUpgrade = true;
               } else {
                  sessionUpgrade = isSessionUpgrade(idpAuthnContextInfo.getAuthnContext(), sessionIndex);
               }

               if (SAML2Utils.DEBUG.messageEnabled()) {
                  SAML2Utils.DEBUG.message(classMethod + "Session Upgrade is :" + sessionUpgrade);
               }

               if (sessionUpgrade) {
                  if (sessionIndex != null) {
                     IDPSession oldIDPSession = IDPCache.idpSessionsByIndices().get(sessionIndex);
                     if (oldIDPSession != null) {
                        IDPCache.oldIDPSessionCache().put(reqID, oldIDPSession);
                     }
                  }

                  IDPCache.authnRequestCache().put(reqID, authnReq);
                  IDPCache.idpAuthnContextCache().put(reqID, matchingAuthnContext);
                  IDPCache.isSessionUpgradeCache().put(reqID, "true");
                  if (relayState != null && !relayState.trim().isEmpty()) {
                     IDPCache.relayStateCache().put(reqID, relayState);
                  }

                  try {
                     redirectToFlowAuthUri(response, reqID);
                     return;
                  } catch (IOException | SAML2Exception ioe) {
                     SAML2Utils.DEBUG.error(classMethod + "Unable to redirect to authentication.", ioe);
                     sessionUpgrade = false;
                     sendError(response, "Server", "UnableToRedirectToAuth", ioe.getMessage(), isFromECP);
                  }
               }

               if (!sessionUpgrade) {
                  NameIDPolicy policy = authnReq.getNameIDPolicy();
                  String nameIDFormat = policy == null ? null : policy.getFormat();

                  try {
                     IDPSSOUtil.sendResponseToACS(
                        request, response, session, authnReq, spEntityID, idpEntityID, idpMetaAlias, realm, nameIDFormat, relayState, matchingAuthnContext
                     );
                  } catch (SAML2Exception se) {
                     SAML2Utils.DEBUG.error(classMethod + "Unable to do sso or federation.", se);
                     sendError(response, "Server", "UnableToDOSSOOrFederation", se.getMessage(), isFromECP);
                  }
               }
            }
         } else {
            synchronized (Saml2DataMapProvider.DataMapId.IDP_AUTHN_REQUEST_CACHE) {
               authnReq = IDPCache.authnRequestCache().remove(reqID);
            }

            AuthnContext matchingAuthnContext;
            synchronized (Saml2DataMapProvider.DataMapId.IDP_IDP_AUTHN_CONTEXT_CACHE) {
               matchingAuthnContext = IDPCache.idpAuthnContextCache().remove(reqID);
            }

            String relayState = IDPCache.relayStateCache().remove(reqID);
            if (authnReq == null) {
               SAML2Utils.DEBUG.error(classMethod + "Unable to get AuthnRequest from cache.");
               sendError(response, "Server", "UnableToGetAuthnReq", null, isFromECP);
               return;
            }

            if (SAML2Utils.DEBUG.messageEnabled()) {
               SAML2Utils.DEBUG.message(classMethod + "RequestID=" + reqID);
            }

            boolean isSessionUpgrade = false;
            if (IDPCache.isSessionUpgradeCache().contains(reqID)) {
               isSessionUpgrade = true;
            }

            SessionProvider sessionProvider = SessionManager.getProvider();
            session = sessionProvider.getSession(request);
            if (isSessionUpgrade) {
               IDPSession oldSess = IDPCache.oldIDPSessionCache().remove(reqID);
               String sessionIndex = IDPSSOUtil.getSessionIndex(session);
               if (sessionIndex != null && !sessionIndex.isEmpty() && oldSess != null) {
                  IDPCache.idpSessionsByIndices().put(sessionIndex, oldSess);
               }
            }

            String spEntityID = authnReq.getIssuer().getValue();
            NameIDPolicy policy = authnReq.getNameIDPolicy();
            String nameIDFormat = policy == null ? null : policy.getFormat();

            try {
               IDPSSOUtil.sendResponseToACS(
                  request, response, session, authnReq, spEntityID, idpEntityID, idpMetaAlias, realm, nameIDFormat, relayState, matchingAuthnContext
               );
            } catch (SAML2Exception se) {
               SAML2Utils.DEBUG.error(classMethod + "Unable to do sso or federation.", se);
               sendError(response, "Server", "UnableToDOSSOOrFederation", se.getMessage(), isFromECP);
            }
         }
      } catch (IOException ioe) {
         SAML2Utils.DEBUG.error(classMethod + "I/O rrror", ioe);
      } catch (SessionException sso) {
         SAML2Utils.DEBUG.error("SSOException : ", sso);
      } catch (SOAPException soapex) {
         SAML2Utils.DEBUG.error("IDPSSOFederate.doSSOFederate:", soapex);
      }
   }

   private static boolean isForceSessionUpgrade(Object session) throws SessionException {
      String[] forceSessionUpgrade = SessionManager.getProvider().getProperty(session, "FORCE_SESSION_UPGRADE");
      return forceSessionUpgrade != null && forceSessionUpgrade.length == 1 && "true".equalsIgnoreCase(forceSessionUpgrade[0]);
   }

   private static void sendError(HttpServletResponse response, String faultCode, String rbKey, String detail, boolean isFromECP) throws IOException, SOAPException {
      if (isFromECP) {
         SOAPMessage soapFault = SAML2Utils.createSOAPFault(faultCode, rbKey, detail);
         if (soapFault != null) {
            if (soapFault.saveRequired()) {
               soapFault.saveChanges();
            }

            response.setStatus(200);
            SAML2Utils.putHeaders(soapFault.getMimeHeaders(), response);
            ByteArrayOutputStream buffer = new ByteArrayOutputStream(2048);
            soapFault.writeTo(buffer);
            Saml2ResponseBodyWriter.writeBodyWithoutCommittingResponse(buffer, response);
         } else {
            response.setStatus(204);
         }
      } else {
         SAMLUtils.sendError(response, rbKey, SAML2Utils.BUNDLE.getString(rbKey));
      }
   }

   private static AuthnRequest getAuthnRequest(String compressedReq) {
      AuthnRequest authnReq = null;
      String outputString = SAML2Utils.decodeFromRedirect(compressedReq);
      if (outputString != null) {
         try {
            authnReq = ProtocolFactory.getInstance().createAuthnRequest(outputString);
         } catch (SAML2Exception se) {
            SAML2Utils.DEBUG.error("IDPSSOFederate.getAuthnRequest(): cannot construct a AuthnRequest object from the SAMLRequest value:", se);
         }
      }

      return authnReq;
   }

   private static AuthnRequest getAuthnRequest(HttpServletRequest request, boolean isFromECP, String binding) {
      if (isFromECP) {
         MimeHeaders headers = SAML2Utils.getHeaders(request);

         try {
            InputStream is = request.getInputStream();
            SOAPMessage msg = SAML2Utils.MESSAGE_FACTORY.createMessage(headers, is);
            Element elem = SAML2Utils.getSamlpElement(msg, "AuthnRequest");
            return ProtocolFactory.getInstance().createAuthnRequest(elem);
         } catch (Exception ex) {
            SAML2Utils.DEBUG.error("IDPSSOFederate.getAuthnRequest:", ex);
            return null;
         }
      } else if (binding.equals("urn:oasis:names:tc:SAML:2.0:bindings:HTTP-Redirect")) {
         String samlRequest = request.getParameter("SAMLRequest");
         if (SAML2Utils.DEBUG.messageEnabled()) {
            SAML2Utils.DEBUG.message("IDPSSOFederate.getAuthnRequest: saml request = " + samlRequest);
         }

         return samlRequest == null ? null : getAuthnRequest(samlRequest);
      } else if (binding.equals("urn:oasis:names:tc:SAML:2.0:bindings:HTTP-POST")) {
         String samlRequest = request.getParameter("SAMLRequest");
         ByteArrayInputStream bis = null;
         AuthnRequest authnRequest = null;

         try {
            byte[] raw = Base64.decode(samlRequest);
            if (raw != null) {
               bis = new ByteArrayInputStream(raw);
               Document doc = XMLUtils.toDOMDocument(bis, SAML2Utils.DEBUG);
               if (doc != null) {
                  authnRequest = ProtocolFactory.getInstance().createAuthnRequest(doc.getDocumentElement());
               }
            }
         } catch (Exception ex) {
            SAML2Utils.DEBUG.error("IDPSSOFederate.getAuthnRequest:", ex);
            return null;
         } finally {
            if (bis != null) {
               try {
                  bis.close();
               } catch (Exception ie) {
                  if (SAML2Utils.DEBUG.messageEnabled()) {
                     SAML2Utils.DEBUG.message("IDPSSOFederate.getAuthnRequest:", ie);
                  }
               }
            }
         }

         return authnRequest;
      } else {
         return null;
      }
   }

   public static void redirectToFlowAuthUri(HttpServletResponse response, String reqID) throws SAML2Exception, IOException {
      String authUri = SystemConfigurationUtil.getProperty("AUTH_URI");
      if (authUri == null) {
         throw new SAML2Exception("No Auth URI defined for flows");
      }

      if (authUri.contains("?")) {
         authUri = authUri + "&";
      } else {
         authUri = authUri + "?";
      }

      authUri = authUri + "ReqID=" + URLEncDec.encode(reqID);
      if (SAML2Utils.DEBUG.messageEnabled()) {
         SAML2Utils.DEBUG.message("Redirecting to auth uri: " + authUri);
      }

      response.sendRedirect(authUri);
   }

   private static boolean isSessionUpgrade(AuthnContext authnContext, String sessionIndex) {
      if (sessionIndex == null) {
         return false;
      }

      Set<AuthnContext> oldAuthnContexts = IDPCache.authnContextCache().get(sessionIndex);
      if (oldAuthnContexts != null && !oldAuthnContexts.isEmpty()) {
         String newAuthnContextClassRef = authnContext.getAuthnContextClassRef();

         for (AuthnContext oldAuthnContext : oldAuthnContexts) {
            if (newAuthnContextClassRef.equals(oldAuthnContext.getAuthnContextClassRef())) {
               return false;
            }
         }

         return true;
      } else {
         return false;
      }
   }
}