package com.sun.identity.saml2.plugins;

import com.sun.identity.plugin.datastore.DataStoreProviderException;
import com.sun.identity.saml2.assertion.Assertion;
import com.sun.identity.saml2.assertion.Attribute;
import com.sun.identity.saml2.assertion.AttributeStatement;
import com.sun.identity.saml2.assertion.EncryptedAttribute;
import com.sun.identity.saml2.assertion.EncryptedID;
import com.sun.identity.saml2.assertion.NameID;
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 java.security.PrivateKey;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;

public class DefaultLibrarySPAccountMapper extends DefaultAccountMapper implements SPAccountMapper {
   private PrivateKey decryptionKey;

   public DefaultLibrarySPAccountMapper() {
      SAML2SDKUtils.DEBUG.message("DefaultLibrarySPAccountMapper.constructor: ");
      this.role = "SPRole";
   }

   @Override
   public String getIdentity(Assertion assertion, String hostEntityID, String realm) throws SAML2Exception {
      if (assertion == null) {
         throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("nullAssertion"));
      }

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

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

      NameID nameID = null;
      EncryptedID encryptedID = assertion.getSubject().getEncryptedID();
      if (encryptedID != null) {
         this.decryptionKey = KeyUtil.getDecryptionKey(SAML2Utils.getSAML2MetaManager().getSPSSOConfig(realm, hostEntityID));
         nameID = encryptedID.decrypt(this.decryptionKey);
      } else {
         nameID = assertion.getSubject().getNameID();
      }

      String userID = null;
      String format = nameID.getFormat();
      boolean transientFormat = false;
      if (format != null && format.equals("urn:oasis:names:tc:SAML:2.0:nameid-format:transient")) {
         transientFormat = true;
         userID = this.getTransientUser(realm, hostEntityID);
      }

      if (userID != null && !userID.isEmpty()) {
         return userID;
      }

      if (!transientFormat) {
         String remoteEntityID = assertion.getIssuer().getValue();
         if (SAML2SDKUtils.DEBUG.messageEnabled()) {
            SAML2SDKUtils.DEBUG.message("DefaultLibrarySPAccountMapper.getIdentity(Assertion): realm = " + realm + " hostEntityID = " + hostEntityID);
         }

         try {
            userID = dsProvider.getUserID(realm, SAML2Utils.getNameIDKeyMap(nameID, hostEntityID, remoteEntityID, realm, this.role));
         } catch (DataStoreProviderException dse) {
            SAML2SDKUtils.DEBUG.error("DefaultLibrarySPAccountMapper.getIdentity(Assertion): DataStoreProviderException", dse);
            throw new SAML2Exception(dse.getMessage());
         }

         if (userID != null) {
            return userID;
         }
      }

      userID = this.getAutoFedUser(realm, hostEntityID, assertion);
      if (userID != null && !userID.isEmpty()) {
         return userID;
      }

      String useNameID = this.getAttribute(realm, hostEntityID, "useNameIDAsSPUserID");
      if (useNameID != null && useNameID.equalsIgnoreCase("true")) {
         if (SAML2SDKUtils.DEBUG.messageEnabled()) {
            SAML2SDKUtils.DEBUG.message("DefaultLibrarySPAccountMapper.getIdentity: use NameID value as userID: " + nameID.getValue());
         }

         return nameID.getValue();
      } else {
         return null;
      }
   }

   protected String getTransientUser(String realm, String entityID) {
      return this.getAttribute(realm, entityID, "transientUser");
   }

   private String getAutoFedUser(String realm, String entityID, Assertion assertion) throws SAML2Exception {
      List<AttributeStatement> attributeStatements = assertion.getAttributeStatements();
      if (attributeStatements != null && !attributeStatements.isEmpty()) {
         String autoFedEnable = this.getAttribute(realm, entityID, "autofedEnabled");
         if (autoFedEnable != null && !autoFedEnable.equals("false")) {
            String autoFedAttribute = this.getAttribute(realm, entityID, "autofedAttribute");
            if (autoFedAttribute != null && !autoFedAttribute.isEmpty()) {
               Set<String> autoFedAttributeValue = null;

               for (AttributeStatement statement : attributeStatements) {
                  autoFedAttributeValue = this.getAttribute(statement, autoFedAttribute, realm, entityID);
                  if (autoFedAttributeValue != null && !autoFedAttributeValue.isEmpty()) {
                     break;
                  }
               }

               if (autoFedAttributeValue != null && !autoFedAttributeValue.isEmpty()) {
                  DefaultSPAttributeMapper attributeMapper = new DefaultSPAttributeMapper();
                  Map<String, String> attributeMap = attributeMapper.getConfigAttributeMap(realm, entityID, "SPRole");
                  if ((attributeMap == null || attributeMap.isEmpty()) && SAML2SDKUtils.DEBUG.messageEnabled()) {
                     SAML2SDKUtils.DEBUG.message("DefaultLibrarySPAccountMapper.getAutoFedUser: attribute map is not configured.");
                  }

                  String autoFedMapAttribute = attributeMap.get(autoFedAttribute);
                  if (autoFedMapAttribute == null) {
                     if (SAML2SDKUtils.DEBUG.messageEnabled()) {
                        SAML2SDKUtils.DEBUG.message("DefaultLibrarySPAccountMapper.getAutoFedUser: Auto federation attribute map is not specified in config.");
                     }

                     autoFedMapAttribute = autoFedAttribute;
                  }

                  try {
                     Map<String, Set<String>> map = new HashMap<>();
                     map.put(autoFedMapAttribute, autoFedAttributeValue);
                     if (SAML2SDKUtils.DEBUG.messageEnabled()) {
                        SAML2SDKUtils.DEBUG.message("DefaultLibrarySPAccountMapper.getAutoFedUser: Search map: " + map);
                     }

                     String userId = dsProvider.getUserID(realm, map);
                     if (userId != null && !userId.isEmpty()) {
                        return userId;
                     }

                     if (this.isDynamicalOrIgnoredProfile(realm)) {
                        if (SAML2SDKUtils.DEBUG.messageEnabled()) {
                           SAML2SDKUtils.DEBUG
                              .message("DefaultLibrarySPAccountMapper: dynamical user creation or ignore profile enabled : uid=" + autoFedAttributeValue);
                        }

                        return autoFedAttributeValue.iterator().next();
                     }
                  } catch (DataStoreProviderException dse) {
                     if (SAML2SDKUtils.DEBUG.warningEnabled()) {
                        SAML2SDKUtils.DEBUG.warning("DefaultLibrarySPAccountMapper.getAutoFedUser: Datastore provider exception", dse);
                     }
                  }

                  return null;
               } else {
                  if (SAML2SDKUtils.DEBUG.messageEnabled()) {
                     SAML2SDKUtils.DEBUG.message("DefaultLibrarySPAccountMapper.getAutoFedUser: Auto federation attribute is not specified in the assertion.");
                  }

                  return null;
               }
            } else {
               if (SAML2SDKUtils.DEBUG.messageEnabled()) {
                  SAML2SDKUtils.DEBUG.message("DefaultLibrarySPAccountMapper.getAutoFedUser: Auto federation attribute is not configured.");
               }

               return null;
            }
         } else {
            if (SAML2SDKUtils.DEBUG.messageEnabled()) {
               SAML2SDKUtils.DEBUG.message("DefaultLibrarySPAccountMapper.getAutoFedUser: Auto federation is disabled.");
            }

            return null;
         }
      } else {
         if (SAML2SDKUtils.DEBUG.messageEnabled()) {
            SAML2SDKUtils.DEBUG.message("DefaultLibrarySPAccountMapper.getAutoFedUser: Assertion does not have attribute statements.");
         }

         return null;
      }
   }

   private boolean isDynamicalOrIgnoredProfile(String realm) {
      return true;
   }

   private Set<String> getAttribute(AttributeStatement statement, String attributeName, String realm, String hostEntityID) {
      if (SAML2SDKUtils.DEBUG.messageEnabled()) {
         SAML2SDKUtils.DEBUG.message("DefaultLibrarySPAccountMapper.getAttribute: attributeName =" + attributeName);
      }

      List<Attribute> list = statement.getAttribute();
      List<EncryptedAttribute> encList = statement.getEncryptedAttribute();
      if (encList != null && !encList.isEmpty()) {
         List<Attribute> allList = new ArrayList<>();
         if (list != null && !list.isEmpty()) {
            allList.addAll(list);
         }

         list = allList;
         Iterator<EncryptedAttribute> encIter = encList.iterator();

         while (encIter.hasNext()) {
            try {
               if (this.decryptionKey == null) {
                  this.decryptionKey = KeyUtil.getDecryptionKey(SAML2Utils.getSAML2MetaManager().getSPSSOConfig(realm, hostEntityID));
               }

               list.add(encIter.next().decrypt(this.decryptionKey));
            } catch (SAML2Exception se) {
               SAML2SDKUtils.DEBUG.error("Decryption error:", se);
               return null;
            }
         }
      }

      for (Attribute attribute : list) {
         if (attributeName.equalsIgnoreCase(attribute.getName())) {
            List<String> values = attribute.getAttributeValueString();
            if (values != null && !values.isEmpty()) {
               return new HashSet<>(values);
            }

            return null;
         }
      }

      return null;
   }
}