package com.sun.identity.saml2.plugins;

import com.sun.identity.saml2.assertion.AssertionFactory;
import com.sun.identity.saml2.assertion.AuthnContext;
import com.sun.identity.saml2.common.SAML2Exception;
import com.sun.identity.saml2.common.SAML2Utils;
import com.sun.identity.saml2.profile.IDPCache;
import com.sun.identity.saml2.protocol.AuthnRequest;
import com.sun.identity.saml2.protocol.RequestedAuthnContext;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.StringTokenizer;
import java.util.Map.Entry;

public class DefaultIDPAuthnContextMapper implements IDPAuthnContextMapper {
   private static final String DEFAULT = "default";

   @Override
   public IDPAuthnContextInfo getIDPAuthnContextInfo(AuthnRequest authnRequest, String idpEntityID, String realm) throws SAML2Exception {
      String classMethod = "DefaultIDPAuthnContextMapper.getIDPAuthnContextInfo: ";
      Map<String, Set<String>> classRefSchemesMap = null;
      if (IDPCache.classRefSchemesHash != null) {
         classRefSchemesMap = IDPCache.classRefSchemesHash.get(idpEntityID + "|" + realm);
      }

      if (classRefSchemesMap == null || classRefSchemesMap.isEmpty()) {
         this.updateAuthnContextMapping(realm, idpEntityID);
         classRefSchemesMap = IDPCache.classRefSchemesHash.get(idpEntityID + "|" + realm);
         if (classRefSchemesMap == null) {
            classRefSchemesMap = new LinkedHashMap<>();
         }
      }

      String classRef = null;
      Set<String> authTypeAndValues = null;
      RequestedAuthnContext requestedAuthnContext = null;
      if (authnRequest != null) {
         requestedAuthnContext = authnRequest.getRequestedAuthnContext();
      }

      if (requestedAuthnContext != null) {
         List<String> requestedClassRefs = requestedAuthnContext.getAuthnContextClassRef();
         String comparison = requestedAuthnContext.getComparison();

         for (String requestedClassRef : requestedClassRefs) {
            List<String> singleClassRef = new ArrayList<>();
            singleClassRef.add(requestedClassRef);

            for (Entry<String, Set<String>> entry : classRefSchemesMap.entrySet()) {
               String tmpClassRef = entry.getKey();
               if (this.isAuthnContextMatching(singleClassRef, tmpClassRef, comparison, realm, idpEntityID)) {
                  authTypeAndValues = entry.getValue();
                  classRef = tmpClassRef;
                  break;
               }
            }

            if (classRef != null) {
               break;
            }
         }

         if (classRef == null) {
            return null;
         }
      } else {
         authTypeAndValues = classRefSchemesMap.get("default");
         classRef = IDPCache.defaultClassRefHash.get(idpEntityID + "|" + realm);
         if (classRef == null) {
            classRef = "urn:oasis:names:tc:SAML:2.0:ac:classes:PasswordProtectedTransport";
         }
      }

      AuthnContext authnContext = AssertionFactory.getInstance().createAuthnContext();
      authnContext.setAuthnContextClassRef(classRef);
      IDPAuthnContextInfo info = new IDPAuthnContextInfo(authnContext, authTypeAndValues);
      if (SAML2Utils.DEBUG.messageEnabled()) {
         SAML2Utils.DEBUG.message(classMethod + "\nreturned AuthnContextClassRef=" + classRef + "\nauthTypeAndValues=" + authTypeAndValues);
      }

      return info;
   }

   @Override
   public AuthnContext getAuthnContextFromAuthLevel(String authLevel, String realm, String idpEntityID) throws SAML2Exception {
      String classRef = null;
      Map<String, Object> classRefLevelMap = null;
      if (IDPCache.classRefLevelHash != null) {
         classRefLevelMap = IDPCache.classRefLevelHash.get(idpEntityID + "|" + realm);
      }

      if (classRefLevelMap == null || classRefLevelMap.isEmpty()) {
         this.updateAuthnContextMapping(realm, idpEntityID);
         classRefLevelMap = IDPCache.classRefLevelHash.get(idpEntityID + "|" + realm);
         if (classRefLevelMap == null) {
            classRefLevelMap = new HashMap<>();
         }
      }

      if (authLevel != null && !authLevel.isEmpty()) {
         try {
            int level = Integer.parseInt(authLevel);

            for (Entry<String, Object> entry : classRefLevelMap.entrySet()) {
               Object value = entry.getValue();
               if (value instanceof Integer && level == (Integer)value) {
                  classRef = entry.getKey();
                  break;
               }
            }
         } catch (NumberFormatException ne) {
            if (SAML2Utils.DEBUG.messageEnabled()) {
               SAML2Utils.DEBUG.message("DefaultIDPAuthnContextMapper.getAuthnContextFromLevel: input authLevel is not valid.", ne);
            }
         }
      }

      if (classRef == null) {
         classRef = IDPCache.defaultClassRefHash.get(idpEntityID + "|" + realm);
         if (classRef == null) {
            classRef = "urn:oasis:names:tc:SAML:2.0:ac:classes:PasswordProtectedTransport";
         }
      }

      if (SAML2Utils.DEBUG.messageEnabled()) {
         SAML2Utils.DEBUG
            .message(
               "DefaultIDPAuthnContext.getClassRefFromLevel: authLevel=" + authLevel + ", classRef=" + classRef + ", classRefLevelMap=" + classRefLevelMap
            );
      }

      AuthnContext result = AssertionFactory.getInstance().createAuthnContext();
      result.setAuthnContextClassRef(classRef);
      return result;
   }

   @Override
   public boolean isAuthnContextMatching(List<String> requestedACClassRefs, String acClassRef, String comparison, String realm, String idpEntityID) {
      Map<String, Object> acClassRefLevelMap = null;
      if (IDPCache.classRefLevelHash != null) {
         acClassRefLevelMap = IDPCache.classRefLevelHash.get(idpEntityID + "|" + realm);
      }

      if (acClassRefLevelMap == null || acClassRefLevelMap.isEmpty()) {
         this.updateAuthnContextMapping(realm, idpEntityID);
         acClassRefLevelMap = IDPCache.classRefLevelHash.get(idpEntityID + "|" + realm);
         if (acClassRefLevelMap == null) {
            acClassRefLevelMap = new HashMap<>();
         }
      }

      return SAML2Utils.isAuthnContextMatching(requestedACClassRefs, acClassRef, comparison, acClassRefLevelMap);
   }

   private void updateAuthnContextMapping(String realm, String idpEntityID) {
      List<String> values = SAML2Utils.getAllAttributeValueFromSSOConfig(realm, idpEntityID, "IDPRole", "idpAuthncontextClassrefMapping");
      Map<String, Object> classRefLevelMap = new LinkedHashMap<>();
      String defaultClassRef = null;
      Map<String, Set<String>> classRefSchemesMap = new LinkedHashMap<>();
      if (values != null && !values.isEmpty()) {
         for (String s : values) {
            boolean isDefault = false;
            String value = s.trim();
            if (value.endsWith("|default")) {
               value = value.substring(0, value.length() - "default".length());
               isDefault = true;
            }

            StringTokenizer st = new StringTokenizer(value, "|");
            if (st.hasMoreTokens()) {
               String classRef = st.nextToken().trim();
               Set<String> authTypeAndValues = new HashSet<>();
               if (st.hasMoreTokens()) {
                  String level = st.nextToken();
                  if (!level.contains("=")) {
                     try {
                        Integer authLevel = Integer.valueOf(level);
                        classRefLevelMap.put(classRef, authLevel);
                        if (isDefault && !classRefLevelMap.containsKey("default")) {
                           classRefLevelMap.put("default", authLevel);
                           defaultClassRef = classRef;
                        }
                     } catch (NumberFormatException nfe) {
                        if (SAML2Utils.DEBUG.messageEnabled()) {
                           SAML2Utils.DEBUG.message("DefaultIDPAuthnContextMapper.getACClassRefLevelMap:", nfe);
                        }
                     }
                  } else if (!level.trim().isEmpty()) {
                     authTypeAndValues.add(level);
                  }

                  while (st.hasMoreTokens()) {
                     String authTypeAndValue = st.nextToken().trim();
                     if (!authTypeAndValue.isEmpty()) {
                        authTypeAndValues.add(authTypeAndValue);
                     }
                  }
               }

               classRefSchemesMap.put(classRef, authTypeAndValues);
               if (isDefault) {
                  classRefSchemesMap.put("default", authTypeAndValues);
               }
            }
         }
      }

      String key = idpEntityID + "|" + realm;
      if (!classRefSchemesMap.isEmpty()) {
         IDPCache.classRefSchemesHash.put(key, classRefSchemesMap);
      }

      if (!classRefLevelMap.isEmpty()) {
         IDPCache.classRefLevelHash.put(key, classRefLevelMap);
      }

      if (defaultClassRef != null) {
         IDPCache.defaultClassRefHash.put(key, defaultClassRef);
      }
   }
}