package com.sun.identity.saml2.common;

import com.sun.identity.shared.configuration.SystemPropertiesManager;
import com.sun.identity.shared.encode.Base64;
import com.sun.identity.shared.encode.URLEncDec;
import java.nio.charset.StandardCharsets;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.Signature;
import java.security.SignatureException;
import java.security.cert.X509Certificate;
import java.util.Arrays;
import java.util.List;
import java.util.Optional;
import java.util.StringTokenizer;

public class QuerySignatureUtil {
   private QuerySignatureUtil() {
   }

   public static String sign(String queryString, PrivateKey privateKey) throws SAML2Exception {
      String algURI = SystemPropertiesManager.get("com.sun.identity.saml.xmlsig.xmlSigAlgorithm", null);
      return sign(queryString, privateKey, algURI != null ? SignatureAlgorithm.forUri(algURI) : null);
   }

   public static String sign(String queryString, PrivateKey privateKey, SignatureAlgorithm signatureAlgorithm) throws SAML2Exception {
      String classMethod = "QuerySignatureUtil.sign: ";
      if (queryString != null && !queryString.isEmpty() && privateKey != null) {
         if (SAML2Utils.DEBUG.messageEnabled()) {
            SAML2Utils.DEBUG.message(classMethod + "Input query string:\n" + queryString);
         }

         String alg = privateKey.getAlgorithm();
         Signature sig;
         String algURI;
         if (signatureAlgorithm != null) {
            try {
               sig = Signature.getInstance(signatureAlgorithm.getAlgorithm());
               algURI = signatureAlgorithm.getUri();
            } catch (NoSuchAlgorithmException nsae) {
               throw new SAML2Exception(nsae);
            }
         } else if (alg.equals("RSA")) {
            try {
               sig = Signature.getInstance("SHA1withRSA");
               algURI = "http://www.w3.org/2000/09/xmldsig#rsa-sha1";
            } catch (NoSuchAlgorithmException nsae) {
               throw new SAML2Exception(nsae);
            }
         } else {
            if (!alg.equals("DSA")) {
               SAML2Utils.DEBUG.error(classMethod + "Algorithm not supported: " + alg);
               throw new SAML2Exception(SAML2Utils.BUNDLE.getString("algorithmNotSupported"));
            }

            try {
               sig = Signature.getInstance("SHA1withDSA");
               algURI = "http://www.w3.org/2000/09/xmldsig#dsa-sha1";
            } catch (NoSuchAlgorithmException nsae) {
               throw new SAML2Exception(nsae);
            }
         }

         if (queryString.charAt(queryString.length() - 1) != '&') {
            queryString = queryString + "&";
         }

         queryString = queryString + "SigAlg=" + URLEncDec.encode(algURI);
         if (SAML2Utils.DEBUG.messageEnabled()) {
            SAML2Utils.DEBUG.message(classMethod + "Final string to be signed:\n" + queryString);
         }

         try {
            sig.initSign(privateKey);
         } catch (InvalidKeyException ike) {
            throw new SAML2Exception(ike);
         }

         try {
            sig.update(queryString.getBytes(StandardCharsets.UTF_8));
         } catch (SignatureException se1) {
            throw new SAML2Exception(se1);
         }

         byte[] sigBytes = null;

         try {
            sigBytes = sig.sign();
         } catch (SignatureException se2) {
            throw new SAML2Exception(se2);
         }

         if (sigBytes != null && sigBytes.length != 0) {
            String encodedSig = Base64.encode(sigBytes);
            queryString = queryString + "&Signature=" + URLEncDec.encode(encodedSig);
            if (SAML2Utils.DEBUG.messageEnabled()) {
               SAML2Utils.DEBUG.message(classMethod + "Signed query string:\n" + queryString);
            }

            return queryString;
         } else {
            SAML2Utils.DEBUG.error(classMethod + "Generated signature is null");
            throw new SAML2Exception(SAML2Utils.BUNDLE.getString("nullSigGenerated"));
         }
      } else {
         SAML2Utils.DEBUG.error(classMethod + "Either input query string or private key is null.");
         throw new SAML2Exception(SAML2Utils.BUNDLE.getString("nullInput"));
      }
   }

   public static boolean verify(String queryString, List<X509Certificate> certs) throws SAML2Exception {
      List<SignatureAlgorithm> allowedSignatureAlgorithms = SignatureAlgorithm.convertAllowedAlgosStringToList(
         SystemPropertiesManager.get("allowedXmlAlgorithms", "")
      );
      if (allowedSignatureAlgorithms == null || allowedSignatureAlgorithms.isEmpty()) {
         allowedSignatureAlgorithms = Arrays.asList(SignatureAlgorithm.RSA_SHA1, SignatureAlgorithm.DSA_SHA1);
      }

      SAML2Exception latestException = null;

      for (X509Certificate cert : certs) {
         try {
            boolean success = verify(queryString, cert, allowedSignatureAlgorithms);
            if (success) {
               return true;
            }
         } catch (SAML2Exception e) {
            latestException = e;
         }
      }

      if (latestException != null) {
         throw latestException;
      } else {
         return false;
      }
   }

   public static boolean verify(String queryString, X509Certificate cert, List<SignatureAlgorithm> allowedSignatureAlgorithms) throws SAML2Exception {
      String classMethod = "QuerySignatureUtil.verify: ";
      if (queryString != null && !queryString.isEmpty() && cert != null) {
         if (SAML2Utils.DEBUG.messageEnabled()) {
            SAML2Utils.DEBUG.message(classMethod + "Query string to be verifed:\n" + queryString);
         }

         StringTokenizer st = new StringTokenizer(queryString, "&");
         String token = null;
         String samlReq = null;
         String samlRes = null;
         String relay = null;
         String sigAlg = null;
         String encSig = null;

         while (st.hasMoreTokens()) {
            token = st.nextToken();
            if (token.startsWith("SAMLRequest")) {
               samlReq = token;
            } else if (token.startsWith("SAMLResponse")) {
               samlRes = token;
            } else if (token.startsWith("RelayState")) {
               relay = token;
            } else if (token.startsWith("SigAlg")) {
               sigAlg = token;
            } else if (token.startsWith("Signature")) {
               encSig = token;
            }
         }

         if (sigAlg == null || sigAlg.isEmpty()) {
            SAML2Utils.DEBUG.error(classMethod + "Null SigAlg query parameter.");
            throw new SAML2Exception(SAML2Utils.BUNDLE.getString("nullSigAlg"));
         }

         if (encSig != null && !encSig.isEmpty()) {
            String newQueryString;
            if (samlReq != null) {
               newQueryString = samlReq;
            } else {
               newQueryString = samlRes;
            }

            if (relay != null) {
               newQueryString = newQueryString + "&" + relay;
            }

            newQueryString = newQueryString + "&" + sigAlg;
            if (SAML2Utils.DEBUG.messageEnabled()) {
               SAML2Utils.DEBUG.message(classMethod + "Query string to be verifed (re-arranged):\n" + newQueryString);
            }

            int sigAlgValueIndex = sigAlg.indexOf(61);
            String sigAlgValue = sigAlg.substring(sigAlgValueIndex + 1);
            if (sigAlgValue != null && !sigAlgValue.isEmpty()) {
               sigAlgValue = URLEncDec.decode(sigAlgValue);
               if (SAML2Utils.DEBUG.messageEnabled()) {
                  SAML2Utils.DEBUG.message(classMethod + "SigAlg query parameter value: " + sigAlgValue);
               }

               String signatureAlgorithmFromQuery = sigAlgValue;
               Optional<SignatureAlgorithm> algorithm = allowedSignatureAlgorithms.stream()
                  .filter(algo -> algo.getUri().equals(signatureAlgorithmFromQuery))
                  .findFirst();
               if (algorithm.isEmpty()) {
                  SAML2Utils.DEBUG
                     .error(
                        classMethod
                           + "Signature algorithm not allowed: \""
                           + sigAlgValue
                           + "\". Allowed algorithms: "
                           + allowedSignatureAlgorithms.stream().map(SignatureAlgorithm::getUri).toList()
                     );
                  return false;
               }

               int encSigValueIndex = encSig.indexOf(61);
               String encSigValue = encSig.substring(encSigValueIndex + 1);
               if (encSigValue != null && !encSigValue.isEmpty()) {
                  encSigValue = URLEncDec.decode(encSigValue);
                  if (SAML2Utils.DEBUG.messageEnabled()) {
                     SAML2Utils.DEBUG.message(classMethod + "Signature query parameter value:\n" + encSigValue);
                  }

                  byte[] signature = Base64.decode(encSigValue);

                  Signature sig;
                  try {
                     sig = Signature.getInstance(algorithm.get().getAlgorithm());
                  } catch (NoSuchAlgorithmException nsae) {
                     throw new SAML2Exception(nsae);
                  }

                  try {
                     sig.initVerify(cert);
                  } catch (InvalidKeyException ike) {
                     throw new SAML2Exception(ike);
                  }

                  try {
                     sig.update(newQueryString.getBytes(StandardCharsets.UTF_8));
                  } catch (SignatureException se1) {
                     throw new SAML2Exception(se1);
                  }

                  try {
                     return sig.verify(signature);
                  } catch (SignatureException se2) {
                     throw new SAML2Exception(se2);
                  }
               } else {
                  SAML2Utils.DEBUG.message(classMethod + "Null Signature query parameter value.");
                  throw new SAML2Exception(SAML2Utils.BUNDLE.getString("nullSig"));
               }
            } else {
               SAML2Utils.DEBUG.error(classMethod + "Null SigAlg query parameter value.");
               throw new SAML2Exception(SAML2Utils.BUNDLE.getString("nullSigAlg"));
            }
         } else {
            SAML2Utils.DEBUG.error(classMethod + "Null Signature query parameter.");
            throw new SAML2Exception(SAML2Utils.BUNDLE.getString("nullSig"));
         }
      } else {
         SAML2Utils.DEBUG.error(classMethod + "Input query string or certificate is null");
         throw new SAML2Exception(SAML2Utils.BUNDLE.getString("nullInput"));
      }
   }
}