package com.sun.identity.saml2.assertion.impl;

import com.sun.identity.saml2.assertion.Advice;
import com.sun.identity.saml2.assertion.Assertion;
import com.sun.identity.saml2.assertion.AssertionFactory;
import com.sun.identity.saml2.assertion.AttributeStatement;
import com.sun.identity.saml2.assertion.AuthnStatement;
import com.sun.identity.saml2.assertion.AuthzDecisionStatement;
import com.sun.identity.saml2.assertion.Conditions;
import com.sun.identity.saml2.assertion.EncryptedAssertion;
import com.sun.identity.saml2.assertion.Issuer;
import com.sun.identity.saml2.assertion.Subject;
import com.sun.identity.saml2.common.SAML2Exception;
import com.sun.identity.saml2.common.SAML2SDKUtils;
import com.sun.identity.saml2.xmlenc.EncManager;
import com.sun.identity.saml2.xmlsig.SigManager;
import com.sun.identity.shared.DateUtils;
import com.sun.identity.shared.xml.XMLUtils;
import java.security.Key;
import java.security.PrivateKey;
import java.security.cert.X509Certificate;
import java.text.ParseException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Date;
import java.util.List;
import org.w3c.dom.Document;
import org.w3c.dom.Element;
import org.w3c.dom.Node;
import org.w3c.dom.NodeList;

public class AssertionImpl implements Assertion {
   private static final String ASSERTION_ELEMENT = "Assertion";
   private static final String ASSERTION_VERSION_ATTR = "Version";
   private static final String ASSERTION_ID_ATTR = "ID";
   private static final String ASSERTION_ISSUEINSTANT_ATTR = "IssueInstant";
   private static final String XSI_TYPE_ATTR = "xsi:type";
   private static final String ASSERTION_ISSUER = "Issuer";
   private static final String ASSERTION_SIGNATURE = "Signature";
   private static final String ASSERTION_SUBJECT = "Subject";
   private static final String ASSERTION_CONDITIONS = "Conditions";
   private static final String ASSERTION_ADVICE = "Advice";
   private static final String ASSERTION_STATEMENT = "Statement";
   private static final String ASSERTION_AUTHNSTATEMENT = "AuthnStatement";
   private static final String ASSERTION_AUTHZDECISIONSTATEMENT = "AuthzDecisionStatement";
   private static final String ASSERTION_ATTRIBUTESTATEMENT = "AttributeStatement";
   private String version;
   private Date issueInstant;
   private Subject subject;
   private Advice advice;
   private String signature;
   private Conditions conditions;
   private String id;
   private List<String> statements = new ArrayList<>();
   private List<AuthnStatement> authnStatements = new ArrayList<>();
   private List<AuthzDecisionStatement> authzDecisionStatements = new ArrayList<>();
   private List<AttributeStatement> attributeStatements = new ArrayList<>();
   private Issuer issuer;
   private boolean isMutable = true;
   private String signedXMLString;
   private Boolean isSignatureValid;

   public AssertionImpl() {
   }

   public AssertionImpl(String xml) throws SAML2Exception {
      Document document = XMLUtils.toDOMDocument(xml, SAML2SDKUtils.DEBUG);
      if (document != null) {
         Element rootElement = document.getDocumentElement();
         this.processElement(rootElement);
         this.makeImmutable();
         if (this.signature != null) {
            this.signedXMLString = xml;
         }
      } else {
         SAML2SDKUtils.DEBUG.error("AssertionImpl.processElement(): invalid XML input");
         throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("errorObtainingElement"));
      }
   }

   public AssertionImpl(Element element) throws SAML2Exception {
      this.processElement(element);
      this.makeImmutable();
      if (this.signature != null) {
         this.signedXMLString = XMLUtils.print(element, "UTF-8");
      }
   }

   private void processElement(Element element) throws SAML2Exception {
      if (element == null) {
         SAML2SDKUtils.DEBUG.error("AssertionImpl.processElement(): invalid root element");
         throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("invalid_element"));
      }

      String elemName = element.getLocalName();
      if (elemName == null) {
         SAML2SDKUtils.DEBUG.error("AssertionImpl.processElement(): local name missing");
         throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("missing_local_name"));
      }

      if (!elemName.equals("Assertion")) {
         SAML2SDKUtils.DEBUG.error("AssertionImpl.processElement(): invalid local name " + elemName);
         throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("invalid_local_name"));
      }

      String attrValue = element.getAttribute("Version");
      if (attrValue != null && !attrValue.isEmpty()) {
         this.version = attrValue;
         attrValue = element.getAttribute("ID");
         if (attrValue != null && !attrValue.isEmpty()) {
            this.id = attrValue;
            attrValue = element.getAttribute("IssueInstant");
            if (attrValue != null && !attrValue.isEmpty()) {
               try {
                  this.issueInstant = DateUtils.stringToDate(attrValue);
               } catch (ParseException pe) {
                  SAML2SDKUtils.DEBUG.error("AssertionImpl.processElement(): invalid issue instant");
                  throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("invalid_date_format"));
               }

               NodeList nodes = element.getChildNodes();
               int numOfNodes = nodes.getLength();
               if (numOfNodes < 1) {
                  SAML2SDKUtils.DEBUG.error("AssertionImpl.processElement(): assertion has no subelements");
                  throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("missing_subelements"));
               }

               AssertionFactory factory = AssertionFactory.getInstance();
               int nextElem = 0;

               Node child;
               for (child = nodes.item(nextElem); child.getNodeType() != 1; child = nodes.item(nextElem)) {
                  if (++nextElem >= numOfNodes) {
                     SAML2SDKUtils.DEBUG.error("AssertionImpl.processElement(): assertion has no subelements");
                     throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("missing_subelements"));
                  }
               }

               String childName = child.getLocalName();
               if (childName != null && childName.equals("Issuer")) {
                  this.issuer = AssertionFactory.getInstance().createIssuer((Element)child);
                  if (++nextElem < numOfNodes) {
                     for (child = nodes.item(nextElem); child.getNodeType() != 1; child = nodes.item(nextElem)) {
                        if (++nextElem >= numOfNodes) {
                           return;
                        }
                     }

                     childName = child.getLocalName();
                     if (childName != null && childName.equals("Signature")) {
                        this.signature = XMLUtils.print(child);
                        if (++nextElem >= numOfNodes) {
                           return;
                        }

                        for (child = nodes.item(nextElem); child.getNodeType() != 1; child = nodes.item(nextElem)) {
                           if (++nextElem >= numOfNodes) {
                              return;
                           }
                        }

                        childName = child.getLocalName();
                     } else {
                        this.signature = null;
                     }

                     if (childName != null && childName.equals("Subject")) {
                        this.subject = factory.createSubject((Element)child);
                        if (++nextElem >= numOfNodes) {
                           return;
                        }

                        for (child = nodes.item(nextElem); child.getNodeType() != 1; child = nodes.item(nextElem)) {
                           if (++nextElem >= numOfNodes) {
                              return;
                           }
                        }

                        childName = child.getLocalName();
                     } else {
                        this.subject = null;
                     }

                     if (childName != null && childName.equals("Conditions")) {
                        this.conditions = factory.createConditions((Element)child);
                        if (++nextElem >= numOfNodes) {
                           return;
                        }

                        for (child = nodes.item(nextElem); child.getNodeType() != 1; child = nodes.item(nextElem)) {
                           if (++nextElem >= numOfNodes) {
                              return;
                           }
                        }

                        childName = child.getLocalName();
                     } else {
                        this.conditions = null;
                     }

                     if (childName != null && childName.equals("Advice")) {
                        this.advice = factory.createAdvice((Element)child);
                        nextElem++;
                     } else {
                        this.advice = null;
                     }

                     for (; nextElem < numOfNodes; nextElem++) {
                        child = nodes.item(nextElem);
                        if (child.getNodeType() == 1) {
                           childName = child.getLocalName();
                           if (childName != null) {
                              if (childName.equals("AuthnStatement")) {
                                 this.authnStatements.add(factory.createAuthnStatement((Element)child));
                              } else if (childName.equals("AuthzDecisionStatement")) {
                                 this.authzDecisionStatements.add(factory.createAuthzDecisionStatement((Element)child));
                              } else if (childName.equals("AttributeStatement")) {
                                 this.attributeStatements.add(factory.createAttributeStatement((Element)child));
                              } else if (childName != null && childName.equals("Signature")) {
                                 this.signature = XMLUtils.print(child);
                              } else {
                                 String type = ((Element)child).getAttribute("xsi:type");
                                 if (!childName.equals("Statement") || type == null || type.isEmpty()) {
                                    SAML2SDKUtils.DEBUG.error("AssertionImpl.processElement(): unexpected subelement " + childName);
                                    throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("unexpected_subelement"));
                                 }

                                 this.statements.add(XMLUtils.print(child));
                              }
                           }
                        }
                     }
                  }
               } else {
                  SAML2SDKUtils.DEBUG.error("AssertionImpl.processElement(): the first element is not <Issuer>");
                  throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("missing_subelement_issuer"));
               }
            } else {
               SAML2SDKUtils.DEBUG.error("AssertionImpl.processElement(): issue instant missing");
               throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("missing_issue_instant"));
            }
         } else {
            SAML2SDKUtils.DEBUG.error("AssertionImpl.processElement(): assertion id missing");
            throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("missing_assertion_id"));
         }
      } else {
         SAML2SDKUtils.DEBUG.error("AssertionImpl.processElement(): version missing");
         throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("missing_assertion_version"));
      }
   }

   @Override
   public String getVersion() {
      return this.version;
   }

   @Override
   public void setVersion(String version) throws SAML2Exception {
      if (!this.isMutable) {
         throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("objectImmutable"));
      }

      this.version = version;
   }

   @Override
   public Date getIssueInstant() {
      return this.issueInstant;
   }

   @Override
   public void setIssueInstant(Date issueInstant) throws SAML2Exception {
      if (!this.isMutable) {
         throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("objectImmutable"));
      }

      this.issueInstant = issueInstant;
   }

   @Override
   public Subject getSubject() {
      return this.subject;
   }

   @Override
   public void setSubject(Subject subject) throws SAML2Exception {
      if (!this.isMutable) {
         throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("objectImmutable"));
      }

      this.subject = subject;
   }

   @Override
   public Advice getAdvice() {
      return this.advice;
   }

   @Override
   public void setAdvice(Advice advice) throws SAML2Exception {
      if (!this.isMutable) {
         throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("objectImmutable"));
      }

      this.advice = advice;
   }

   @Override
   public String getSignature() {
      return this.signature;
   }

   @Override
   public Conditions getConditions() {
      return this.conditions;
   }

   @Override
   public void setConditions(Conditions conditions) throws SAML2Exception {
      if (!this.isMutable) {
         throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("objectImmutable"));
      }

      this.conditions = conditions;
   }

   @Override
   public String getID() {
      return this.id;
   }

   @Override
   public void setID(String id) throws SAML2Exception {
      if (!this.isMutable) {
         throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("objectImmutable"));
      }

      this.id = id;
   }

   @Override
   public List<String> getStatements() {
      return this.statements;
   }

   @Override
   public List<AuthnStatement> getAuthnStatements() {
      return this.authnStatements;
   }

   @Override
   public List<AuthzDecisionStatement> getAuthzDecisionStatements() {
      return this.authzDecisionStatements;
   }

   @Override
   public List<AttributeStatement> getAttributeStatements() {
      return this.attributeStatements;
   }

   @Override
   public void setStatements(List<String> statements) throws SAML2Exception {
      if (!this.isMutable) {
         throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("objectImmutable"));
      }

      this.statements = statements;
   }

   @Override
   public void setAuthnStatements(List<AuthnStatement> statements) throws SAML2Exception {
      if (!this.isMutable) {
         throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("objectImmutable"));
      }

      this.authnStatements = statements;
   }

   @Override
   public void setAuthzDecisionStatements(List<AuthzDecisionStatement> statements) throws SAML2Exception {
      if (!this.isMutable) {
         throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("objectImmutable"));
      }

      this.authzDecisionStatements = statements;
   }

   @Override
   public void setAttributeStatements(List<AttributeStatement> statements) throws SAML2Exception {
      if (!this.isMutable) {
         throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("objectImmutable"));
      }

      this.attributeStatements = statements;
   }

   @Override
   public Issuer getIssuer() {
      return this.issuer;
   }

   @Override
   public void setIssuer(Issuer issuer) throws SAML2Exception {
      if (!this.isMutable) {
         throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("objectImmutable"));
      }

      this.issuer = issuer;
   }

   @Override
   public boolean isSigned() {
      return this.signature != null;
   }

   @Override
   public boolean isSignatureValid(List<X509Certificate> senderCerts) throws SAML2Exception {
      if (this.isSignatureValid == null) {
         if (this.signedXMLString == null) {
            this.signedXMLString = this.toXMLString(true, true);
         }

         this.isSignatureValid = SigManager.getSigInstance().verify(this.signedXMLString, this.getID(), senderCerts);
      }

      return this.isSignatureValid;
   }

   @Override
   public void sign(PrivateKey privateKey, X509Certificate cert) throws SAML2Exception {
      Element signatureElement = SigManager.getSigInstance().sign(this.toXMLString(true, true), this.getID(), privateKey, cert);
      this.signature = XMLUtils.print(signatureElement);
      this.signedXMLString = XMLUtils.print(signatureElement.getOwnerDocument().getDocumentElement(), "UTF-8");
      this.makeImmutable();
   }

   @Override
   public EncryptedAssertion encrypt(Key recipientPublicKey, String dataEncAlgorithm, int dataEncStrength, String recipientEntityID) throws SAML2Exception {
      Element el = EncManager.getEncInstance()
         .encrypt(this.toXMLString(true, true), recipientPublicKey, dataEncAlgorithm, dataEncStrength, recipientEntityID, "EncryptedAssertion");
      return AssertionFactory.getInstance().createEncryptedAssertion(el);
   }

   @Override
   public boolean isTimeValid() {
      return this.conditions == null ? true : this.conditions.checkDateValidity(System.currentTimeMillis());
   }

   @Override
   public String toXMLString(boolean includeNSPrefix, boolean declareNS) throws SAML2Exception {
      if (this.signature != null && this.signedXMLString != null) {
         return this.signedXMLString;
      }

      StringBuilder sb = new StringBuilder(2000);
      String ns = "";
      String appendNS = "";
      if (declareNS) {
         ns = " xmlns:saml=\"urn:oasis:names:tc:SAML:2.0:assertion\"";
      }

      if (includeNSPrefix) {
         appendNS = "saml:";
      }

      sb.append("<").append(appendNS).append("Assertion").append(ns);
      if (this.version != null && !this.version.isEmpty()) {
         sb.append(" ").append("Version").append("=\"").append(this.version).append("\"");
         if (this.id != null && !this.id.isEmpty()) {
            sb.append(" ").append("ID").append("=\"").append(this.id).append("\"");
            if (this.issueInstant == null) {
               SAML2SDKUtils.DEBUG.error("AssertionImpl.toXMLString(): issue instant missing");
               throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("missing_issue_instant"));
            }

            String instantStr = DateUtils.toUTCDateFormat(this.issueInstant);
            sb.append(" ").append("IssueInstant").append("=\"").append(instantStr).append("\"").append(">\n");
            if (this.issuer == null) {
               SAML2SDKUtils.DEBUG.error("AssertionImpl.toXMLString(): issuer missing");
               throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("missing_subelement_issuer"));
            }

            sb.append(this.issuer.toXMLString(includeNSPrefix, false));
            if (this.signature != null) {
               sb.append(this.signature);
            }

            if (this.subject != null) {
               sb.append(this.subject.toXMLString(includeNSPrefix, false));
            }

            if (this.conditions != null) {
               sb.append(this.conditions.toXMLString(includeNSPrefix, false));
            }

            if (this.advice != null) {
               sb.append(this.advice.toXMLString(includeNSPrefix, false));
            }

            if (this.statements != null) {
               for (String str : this.statements) {
                  sb.append(str);
               }
            }

            if (this.authnStatements != null) {
               for (AuthnStatement st : this.authnStatements) {
                  sb.append(st.toXMLString(includeNSPrefix, false));
               }
            }

            if (this.authzDecisionStatements != null) {
               for (AuthzDecisionStatement st : this.authzDecisionStatements) {
                  sb.append(st.toXMLString(includeNSPrefix, false));
               }
            }

            if (this.attributeStatements != null) {
               for (AttributeStatement st : this.attributeStatements) {
                  sb.append(st.toXMLString(includeNSPrefix, false));
               }
            }

            sb.append("</").append(appendNS).append("Assertion").append(">\n");
            return sb.toString();
         } else {
            SAML2SDKUtils.DEBUG.error("AssertionImpl.toXMLString(): assertion id missing");
            throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("missing_assertion_id"));
         }
      } else {
         SAML2SDKUtils.DEBUG.error("AssertionImpl.toXMLString(): version missing");
         throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("missing_assertion_version"));
      }
   }

   @Override
   public String toXMLString() throws SAML2Exception {
      return this.toXMLString(true, false);
   }

   @Override
   public final void makeImmutable() {
      if (this.isMutable) {
         if (this.authnStatements != null) {
            for (AuthnStatement authn : this.authnStatements) {
               authn.makeImmutable();
            }

            this.authnStatements = Collections.unmodifiableList(this.authnStatements);
         }

         if (this.authzDecisionStatements != null) {
            for (AuthzDecisionStatement authz : this.authzDecisionStatements) {
               authz.makeImmutable();
            }

            this.authzDecisionStatements = Collections.unmodifiableList(this.authzDecisionStatements);
         }

         if (this.attributeStatements != null) {
            for (AttributeStatement attr : this.attributeStatements) {
               attr.makeImmutable();
            }

            this.attributeStatements = Collections.unmodifiableList(this.attributeStatements);
         }

         if (this.statements != null) {
            this.statements = Collections.unmodifiableList(this.statements);
         }

         if (this.conditions != null) {
            this.conditions.makeImmutable();
         }

         if (this.issuer != null) {
            this.issuer.makeImmutable();
         }

         if (this.subject != null) {
            this.subject.makeImmutable();
         }

         if (this.advice != null) {
            this.advice.makeImmutable();
         }

         this.isMutable = false;
      }
   }

   @Override
   public boolean isMutable() {
      return this.isMutable;
   }
}