package com.sun.identity.saml2.profile;

import com.airlock.iam.openfedlib.saml2.XjcHelper;
import com.airlock.iam.openfedlib.saml2.state.Saml2DataMapProvider;
import com.airlock.iam.opensso.jaxb.entityconfig.BaseConfigType;
import com.airlock.iam.opensso.jaxb.metadata.AffiliationDescriptorType;
import com.airlock.iam.opensso.jaxb.metadata.EndpointType;
import com.airlock.iam.opensso.jaxb.metadata.EntityDescriptorElement;
import com.airlock.iam.opensso.jaxb.metadata.EntityDescriptorType;
import com.airlock.iam.opensso.jaxb.metadata.ExtensionsType;
import com.airlock.iam.opensso.jaxb.metadata.IDPSSODescriptorType;
import com.airlock.iam.opensso.jaxb.metadata.IndexedEndpointType;
import com.airlock.iam.opensso.jaxb.metadata.SPSSODescriptorType;
import com.airlock.iam.opensso.jaxb.metadata.SingleSignOnServiceElement;
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.QuerySignatureUtil;
import com.sun.identity.saml2.common.SAML2Exception;
import com.sun.identity.saml2.common.SAML2SDKUtils;
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.SAML2MetaManager;
import com.sun.identity.saml2.meta.SAML2MetaUtils;
import com.sun.identity.saml2.plugins.SAML2ServiceProviderAdapter;
import com.sun.identity.saml2.plugins.SPAuthnContextMapper;
import com.sun.identity.saml2.protocol.AuthnRequest;
import com.sun.identity.saml2.protocol.Extensions;
import com.sun.identity.saml2.protocol.NameIDPolicy;
import com.sun.identity.saml2.protocol.ProtocolFactory;
import com.sun.identity.saml2.protocol.RequestedAuthnContext;
import com.sun.identity.shared.datastruct.OrderedSet;
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 java.io.IOException;
import java.security.PrivateKey;
import java.util.Date;
import java.util.List;
import java.util.Map;
import java.util.logging.Level;

public class SPSSOFederate {
   private static SAML2MetaManager sm;

   public static void initiateAuthnRequest(
      HttpServletRequest request,
      HttpServletResponse response,
      String spEntityID,
      String idpEntityID,
      String realmName,
      Map<String, List<String>> paramsMap,
      String requestID
   ) throws SAML2Exception {
      if (spEntityID == null) {
         SAML2SDKUtils.DEBUG.error("SPSSOFederate:Service Provider ID is missing.");
         String[] data = new String[]{spEntityID};
         LogUtil.error(Level.INFO, "INVALID_SP", data, null);
         throw new SAML2Exception(SAML2Utils.BUNDLE.getString("nullSPEntityID"));
      }

      if (idpEntityID == null) {
         SAML2SDKUtils.DEBUG.error("SPSSOFederate: Identity Provider ID is missing .");
         String[] data = new String[]{idpEntityID};
         LogUtil.error(Level.INFO, "INVALID_IDP", data, null);
         throw new SAML2Exception(SAML2Utils.BUNDLE.getString("nullIDPEntityID"));
      }

      String binding = getParameter(paramsMap, "reqBinding");
      if (binding == null) {
         binding = "urn:oasis:names:tc:SAML:2.0:bindings:HTTP-Redirect";
      }

      if (SAML2SDKUtils.DEBUG.messageEnabled()) {
         SAML2SDKUtils.DEBUG.message("SPSSOFederate: in initiateSSOFed");
         SAML2SDKUtils.DEBUG.message("SPSSOFederate: spEntityID is : " + spEntityID);
         SAML2SDKUtils.DEBUG.message("SPSSOFederate: idpEntityID : " + idpEntityID);
      }

      String realm = getRealm(realmName);

      try {
         if (sm == null) {
            throw new SAML2Exception(SAML2Utils.BUNDLE.getString("errorMetaManager"));
         }

         BaseConfigType spEntityCfg = sm.getSPSSOConfig(realm, spEntityID);
         Map<String, List<String>> spConfigAttrsMap = null;
         if (spEntityCfg != null) {
            spConfigAttrsMap = SAML2MetaUtils.getAttributes(spEntityCfg);
         }

         SPSSODescriptorType spsso = sm.getSPSSODescriptor(realm, spEntityID);
         if (spsso == null) {
            String[] data = new String[]{spEntityID};
            LogUtil.error(Level.INFO, "SP_METADATA_ERROR", data, null);
            throw new SAML2Exception(SAML2Utils.BUNDLE.getString("metaDataError"));
         }

         List<String> extensionsList = getExtensionsList(sm, spEntityID, realm);
         IDPSSODescriptorType idpsso = sm.getIDPSSODescriptor(realm, idpEntityID);
         if (idpsso == null) {
            String[] data = new String[]{idpEntityID};
            LogUtil.error(Level.INFO, "IDP_METADATA_ERROR", data, null);
            throw new SAML2Exception(SAML2Utils.BUNDLE.getString("metaDataError"));
         }

         List<SingleSignOnServiceElement> ssoServiceList = idpsso.getSingleSignOnService();
         String ssoURL = getSSOURL(ssoServiceList, binding);
         if (ssoURL != null && !ssoURL.isEmpty()) {
            AuthnRequest authnRequest = createAuthnRequest(
               realm, spEntityID, paramsMap, spConfigAttrsMap, extensionsList, spsso, idpsso, ssoURL, false, requestID
            );
            SAML2ServiceProviderAdapter spAdapter = SAML2Utils.getSPAdapterClass(spEntityID, realmName);
            if (spAdapter != null) {
               spAdapter.preSingleSignOnRequest(spEntityID, idpEntityID, realmName, request, response, authnRequest);
            }

            String authReqXMLString = authnRequest.toXMLString(true, true);
            if (SAML2SDKUtils.DEBUG.messageEnabled()) {
               SAML2SDKUtils.DEBUG.message("SPSSOFederate: AuthnRequest:" + authReqXMLString);
            }

            String relayState = getParameter(paramsMap, "RelayState");
            String relayStateID = null;
            if (relayState != null && !relayState.isEmpty()) {
               relayStateID = getRelayStateID(relayState, authnRequest.getID());
            }

            if (binding.equals("urn:oasis:names:tc:SAML:2.0:bindings:HTTP-POST")) {
               if (idpsso != null && XjcHelper.mapBooleanWrapperToPrimitiveWithFalseAsDefault(idpsso.isWantAuthnRequestsSigned())
                  || spsso != null && XjcHelper.mapBooleanWrapperToPrimitiveWithFalseAsDefault(spsso.isAuthnRequestsSigned())) {
                  String certAlias = getParameter(spConfigAttrsMap, "signingCertAlias");
                  signAuthnRequest(certAlias, authnRequest);
               }

               String authXMLString = authnRequest.toXMLString(true, true);
               if (SAML2SDKUtils.DEBUG.messageEnabled()) {
                  SAML2SDKUtils.DEBUG.message("SPSSOFederate.initiateAuthnRequest: SAML Response content :\n" + authXMLString);
               }

               String encodedReqMsg = SAML2Utils.encodeForPOST(authXMLString);
               SAML2Utils.postToTarget(request, response, "SAMLRequest", encodedReqMsg, "RelayState", relayStateID, ssoURL);
            } else {
               String encodedXML = SAML2Utils.encodeForRedirect(authReqXMLString);
               StringBuilder queryString = new StringBuilder();
               queryString.append("SAMLRequest").append("=").append(encodedXML);
               if (relayStateID != null && !relayStateID.isEmpty()) {
                  queryString.append("&").append("RelayState").append("=").append(URLEncDec.encode(relayStateID));
               }

               StringBuilder redirectURL = new StringBuilder().append(ssoURL).append(ssoURL.contains("?") ? "&" : "?");
               if ((idpsso == null || !XjcHelper.mapBooleanWrapperToPrimitiveWithFalseAsDefault(idpsso.isWantAuthnRequestsSigned()))
                  && (spsso == null || !XjcHelper.mapBooleanWrapperToPrimitiveWithFalseAsDefault(spsso.isAuthnRequestsSigned()))) {
                  redirectURL.append(queryString);
               } else {
                  String certAlias = getParameter(spConfigAttrsMap, "signingCertAlias");
                  String signedQueryStr = signQueryString(queryString.toString(), certAlias);
                  redirectURL.append(signedQueryStr);
               }

               response.sendRedirect(redirectURL.toString());
            }

            String[] data = new String[]{ssoURL};
            LogUtil.access(Level.INFO, "REDIRECT_TO_IDP", data, null);
            synchronized (Saml2DataMapProvider.DataMapId.SP_REQUEST_HASH) {
               SPCache.requestHash().put(authnRequest.getID(), authnRequest);
            }
         } else {
            String[] data = new String[]{idpEntityID};
            LogUtil.error(Level.INFO, "SSO_NOT_FOUND", data, null);
            throw new SAML2Exception(SAML2Utils.BUNDLE.getString("ssoServiceNotfound"));
         }
      } catch (IOException ioe) {
         SAML2SDKUtils.DEBUG.error("SPSSOFederate: Exception :", ioe);
         throw new SAML2Exception(SAML2Utils.BUNDLE.getString("errorCreatingAuthnRequest"));
      } catch (SAML2MetaException sme) {
         SAML2SDKUtils.DEBUG.error("SPSSOFederate:Error retreiving metadata", sme);
         throw new SAML2Exception(SAML2Utils.BUNDLE.getString("metaDataError"));
      }
   }

   private static NameIDPolicy createNameIDPolicy(
      String spEntityID,
      String format,
      boolean allowCreate,
      SPSSODescriptorType spsso,
      IDPSSODescriptorType idpsso,
      String realm,
      Map<String, List<String>> paramsMap
   ) throws SAML2Exception {
      format = SAML2Utils.verifyNameIDFormat(format, spsso, idpsso);
      NameIDPolicy nameIDPolicy = ProtocolFactory.getInstance().createNameIDPolicy();
      String affiliationID = getParameter(paramsMap, "affiliationID");
      if (affiliationID != null) {
         AffiliationDescriptorType affiDesc = sm.getAffiliationDescriptor(realm, affiliationID);
         if (affiDesc == null) {
            throw new SAML2Exception(SAML2Utils.BUNDLE.getString("affiliationNotFound"));
         }

         if (!affiDesc.getAffiliateMember().contains(spEntityID)) {
            throw new SAML2Exception(SAML2Utils.BUNDLE.getString("spNotAffiliationMember"));
         }

         nameIDPolicy.setSPNameQualifier(affiliationID);
      } else {
         nameIDPolicy.setSPNameQualifier(spEntityID);
      }

      nameIDPolicy.setAllowCreate(allowCreate);
      nameIDPolicy.setFormat(format);
      return nameIDPolicy;
   }

   private static Issuer createIssuer(String spEntityID) throws SAML2Exception {
      Issuer issuer = AssertionFactory.getInstance().createIssuer();
      issuer.setValue(spEntityID);
      return issuer;
   }

   private static AuthnRequest createAuthnRequest(
      String realmName,
      String spEntityID,
      Map<String, List<String>> paramsMap,
      Map<String, List<String>> spConfigMap,
      List<String> extensionsList,
      SPSSODescriptorType spsso,
      IDPSSODescriptorType idpsso,
      String ssourl,
      boolean isForECP,
      String requestID
   ) throws SAML2Exception {
      String destinationURI = getParameter(paramsMap, "Destination");
      Boolean isPassive = doPassive(paramsMap, spConfigMap);
      Boolean isforceAuthn = isForceAuthN(paramsMap, spConfigMap);
      boolean allowCreate = isAllowCreate(paramsMap, spConfigMap);
      String consent = getParameter(paramsMap, "Consent");
      Extensions extensions = createExtensions(extensionsList);
      String nameIDPolicyFormat = getParameter(paramsMap, "NameIDFormat");
      NameIDPolicy nameIDPolicy = createNameIDPolicy(spEntityID, nameIDPolicyFormat, allowCreate, spsso, idpsso, realmName, paramsMap);
      Issuer issuer = createIssuer(spEntityID);
      Integer acsIndex = getIndex(paramsMap, "AssertionConsumerServiceIndex");
      Integer attrIndex = getIndex(paramsMap, "AttributeConsumingServiceIndex");
      String protocolBinding = isForECP ? "urn:oasis:names:tc:SAML:2.0:bindings:PAOS" : getParameter(paramsMap, "binding");
      OrderedSet<String> acsSet = getACSUrl(spsso, protocolBinding);
      String acsURL = acsSet.get(0);
      protocolBinding = acsSet.get(1);
      RequestedAuthnContext reqAuthnContext = createReqAuthnContext(realmName, spEntityID, paramsMap, spConfigMap);
      AuthnRequest authnReq = ProtocolFactory.getInstance().createAuthnRequest();
      if (!isForECP) {
         if (destinationURI != null && !destinationURI.isEmpty()) {
            authnReq.setDestination(destinationURI);
         } else {
            authnReq.setDestination(ssourl);
         }
      }

      authnReq.setConsent(consent);
      authnReq.setIsPassive(isPassive);
      authnReq.setForceAuthn(isforceAuthn);
      authnReq.setAttributeConsumingServiceIndex(attrIndex);
      authnReq.setAssertionConsumerServiceIndex(acsIndex);
      authnReq.setAssertionConsumerServiceURL(XMLUtils.escapeSpecialCharacters(acsURL));
      authnReq.setProtocolBinding(protocolBinding);
      authnReq.setIssuer(issuer);
      authnReq.setNameIDPolicy(nameIDPolicy);
      authnReq.setRequestedAuthnContext(reqAuthnContext);
      if (extensions != null) {
         authnReq.setExtensions(extensions);
      }

      authnReq.setID(requestID);
      authnReq.setVersion("2.0");
      authnReq.setIssueInstant(new Date());
      return authnReq;
   }

   private static Boolean getAttrValueFromMap(Map<String, List<String>> attrMap, String attrName) {
      Boolean boolVal = null;
      if (attrMap != null && !attrMap.isEmpty()) {
         String attrVal = getParameter(attrMap, attrName);
         if (attrVal != null && (attrVal.equals("true") || attrVal.equals("false"))) {
            boolVal = Boolean.valueOf(attrVal);
         }
      }

      return boolVal;
   }

   static String getSSOURL(List<SingleSignOnServiceElement> ssoServiceList, String binding) {
      String ssoURL = null;
      if (ssoServiceList != null && !ssoServiceList.isEmpty()) {
         for (SingleSignOnServiceElement sso : ssoServiceList) {
            if (sso != null && ((EndpointType)sso.getValue()).getBinding() != null && ((EndpointType)sso.getValue()).getBinding().equals(binding)) {
               ssoURL = ((EndpointType)sso.getValue()).getLocation();
               break;
            }
         }
      }

      if (SAML2SDKUtils.DEBUG.messageEnabled()) {
         SAML2SDKUtils.DEBUG.message("SPSSOFederate:  SingleSignOnService URL :" + ssoURL);
      }

      return ssoURL;
   }

   static OrderedSet<String> getACSUrl(SPSSODescriptorType spsso, String binding) {
      String responseBinding = binding;
      if (binding != null && !binding.isEmpty() && !binding.contains("urn:oasis:names:tc:SAML:2.0:bindings:")) {
         responseBinding = "urn:oasis:names:tc:SAML:2.0:bindings:" + binding;
      }

      List<IndexedEndpointType> acsList = spsso.getAssertionConsumerService();
      String acsURL = null;
      if (acsList != null && !acsList.isEmpty()) {
         for (IndexedEndpointType ace : acsList) {
            if (ace != null
               && XjcHelper.mapBooleanWrapperToPrimitiveWithFalseAsDefault(ace.isIsDefault())
               && (responseBinding == null || responseBinding.isEmpty())) {
               acsURL = ace.getLocation();
               responseBinding = ace.getBinding();
               break;
            }

            if (ace != null && ace.getBinding().equals(responseBinding)) {
               acsURL = ace.getLocation();
               break;
            }
         }
      }

      OrderedSet<String> ol = new OrderedSet<>();
      ol.add(acsURL);
      ol.add(responseBinding);
      if (SAML2SDKUtils.DEBUG.messageEnabled()) {
         SAML2SDKUtils.DEBUG.message("SPSSOFederate: AssertionConsumerService : URL :" + acsURL);
         SAML2SDKUtils.DEBUG.message("SPSSOFederate: AssertionConsumerService : Binding Passed in Query: " + binding);
         SAML2SDKUtils.DEBUG.message("SPSSOFederate: AssertionConsumerService : Binding : " + responseBinding);
      }

      return ol;
   }

   private static String getRealm(String realm) {
      return realm != null && !realm.isEmpty() ? realm : "/";
   }

   private static Boolean doPassive(Map<String, List<String>> paramsMap, Map<String, List<String>> spConfigAttrsMap) {
      String isPassiveStr = getParameter(paramsMap, "IsPassive");
      Boolean isPassive;
      if (isPassiveStr == null || !isPassiveStr.equals("true") && !isPassiveStr.equals("false")) {
         isPassive = getAttrValueFromMap(spConfigAttrsMap, "IsPassive");
      } else {
         isPassive = Boolean.valueOf(isPassiveStr);
      }

      if (SAML2SDKUtils.DEBUG.messageEnabled()) {
         SAML2SDKUtils.DEBUG.message("SPSSOFederate: isPassive : " + isPassive);
      }

      return isPassive == null ? Boolean.FALSE : isPassive;
   }

   private static Boolean isForceAuthN(Map<String, List<String>> paramsMap, Map<String, List<String>> spConfigAttrsMap) {
      String forceAuthn = getParameter(paramsMap, "ForceAuthn");
      Boolean isforceAuthn;
      if (forceAuthn == null || !forceAuthn.equals("true") && !forceAuthn.equals("false")) {
         isforceAuthn = getAttrValueFromMap(spConfigAttrsMap, "ForceAuthn");
      } else {
         isforceAuthn = Boolean.valueOf(forceAuthn);
      }

      if (SAML2SDKUtils.DEBUG.messageEnabled()) {
         SAML2SDKUtils.DEBUG.message("SPSSOFederate:ForceAuthn: " + forceAuthn);
      }

      return isforceAuthn == null ? Boolean.FALSE : isforceAuthn;
   }

   private static boolean isAllowCreate(Map<String, List<String>> paramsMap, Map<String, List<String>> spConfigAttrsMap) {
      boolean allowCreate = true;
      String allowCreateStr = getParameter(paramsMap, "AllowCreate");
      if (allowCreateStr == null || !allowCreateStr.equals("true") && !allowCreateStr.equals("false")) {
         Boolean val = getAttrValueFromMap(spConfigAttrsMap, "AllowCreate");
         if (val != null) {
            allowCreate = val;
         }
      } else {
         allowCreate = Boolean.parseBoolean(allowCreateStr);
      }

      if (SAML2SDKUtils.DEBUG.messageEnabled()) {
         SAML2SDKUtils.DEBUG.message("SPSSOFederate:AllowCreate:" + allowCreate);
      }

      return allowCreate;
   }

   private static Integer getIndex(Map<String, List<String>> paramsMap, String attrName) {
      Integer attrIndex = null;
      String index = getParameter(paramsMap, attrName);
      if (index != null && !index.isEmpty()) {
         attrIndex = Integer.valueOf(index);
      }

      return attrIndex;
   }

   public static String getParameter(Map<String, List<String>> paramsMap, String attrName) {
      String attrVal = null;
      if (paramsMap != null && !paramsMap.isEmpty()) {
         List<String> attrValList = paramsMap.get(attrName);
         if (attrValList != null && !attrValList.isEmpty()) {
            attrVal = attrValList.getFirst();
         }
      }

      return attrVal;
   }

   private static List<String> getExtensionsList(SAML2MetaManager sm, String entityID, String realm) {
      List<String> extensionsList = null;

      try {
         EntityDescriptorElement ed = sm.getEntityDescriptor(realm, entityID);
         if (ed != null) {
            ExtensionsType ext = ((EntityDescriptorType)ed.getValue()).getExtensions();
            if (ext != null) {
               extensionsList = getAny(ext).stream().map(Object::toString).toList();
            }
         }
      } catch (SAML2Exception e) {
         SAML2SDKUtils.DEBUG.error("SPSSOFederate:Error retrieving EntityDescriptor");
      }

      return extensionsList;
   }

   private static List<Object> getAny(ExtensionsType type) {
      return type.getAny();
   }

   private static Extensions createExtensions(List<String> extensionsList) throws SAML2Exception {
      Extensions extensions = null;
      if (extensionsList != null && !extensionsList.isEmpty()) {
         extensions = ProtocolFactory.getInstance().createExtensions();
         extensions.setAny(extensionsList);
      }

      return extensions;
   }

   private static String getRelayStateID(String relayState, String requestID) {
      SPCache.relayStateHash().put(requestID, relayState);
      return requestID;
   }

   private static RequestedAuthnContext createReqAuthnContext(
      String realmName, String spEntityID, Map<String, List<String>> paramsMap, Map<String, List<String>> spConfigMap
   ) {
      RequestedAuthnContext reqCtx = null;
      String className = null;
      if (spConfigMap != null && !spConfigMap.isEmpty()) {
         List<String> listVal = spConfigMap.get("spAuthncontextMapper");
         if (listVal != null && !listVal.isEmpty()) {
            className = listVal.getFirst().trim();
         }
      }

      SPAuthnContextMapper spAuthnContextMapper = SAML2Utils.getSPAuthnContextMapper(realmName, spEntityID, className);

      try {
         reqCtx = spAuthnContextMapper.getRequestedAuthnContext(realmName, spEntityID, paramsMap);
      } catch (SAML2Exception e) {
         if (SAML2SDKUtils.DEBUG.messageEnabled()) {
            SAML2SDKUtils.DEBUG.message("SPSSOFederate:Error creating RequestedAuthnContext", e);
         }
      }

      return reqCtx;
   }

   private static String signQueryString(String queryString, String certAlias) throws SAML2Exception {
      if (SAML2SDKUtils.DEBUG.messageEnabled()) {
         SAML2SDKUtils.DEBUG.message("SPSSOFederate:queryString:" + queryString);
         SAML2SDKUtils.DEBUG.message("SPSSOFederate: certAlias :" + certAlias);
      }

      KeyProvider kp = KeyUtil.getKeyProviderInstance();
      PrivateKey privateKey = kp.getPrivateKey(certAlias);
      return QuerySignatureUtil.sign(queryString, privateKey);
   }

   private static void signAuthnRequest(String certAlias, AuthnRequest authnRequest) throws SAML2Exception {
      KeyProvider kp = KeyUtil.getKeyProviderInstance();
      if (kp == null) {
         SAML2SDKUtils.DEBUG.error("SPSSOFederate:signAuthnRequest: Unable to get a key provider instance.");
         throw new SAML2Exception(SAML2Utils.BUNDLE.getString("nullKeyProvider"));
      }

      authnRequest.sign(kp.getPrivateKey(certAlias), kp.getX509Certificate(certAlias));
   }

   static {
      try {
         sm = new SAML2MetaManager();
      } catch (SAML2MetaException sme) {
         SAML2SDKUtils.DEBUG.error("SPSSOFederate: Error retreiving metadata", sme);
      }
   }
}