// Copyright (c) 2013-2023 xipki. All rights reserved.
// License Apache License 2.0

package org.xipki.qa.ca;

import org.bouncycastle.asn1.ASN1Encodable;
import org.bouncycastle.asn1.ASN1Integer;
import org.bouncycastle.asn1.ASN1ObjectIdentifier;
import org.bouncycastle.asn1.ASN1Sequence;
import org.bouncycastle.asn1.x509.SubjectPublicKeyInfo;
import org.bouncycastle.asn1.x9.X9ECParameters;
import org.bouncycastle.jcajce.provider.asymmetric.util.ECUtil;
import org.bouncycastle.math.ec.ECCurve;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.xipki.ca.api.profile.KeyParametersOption;
import org.xipki.ca.api.profile.KeyParametersOption.AllowAllParametersOption;
import org.xipki.ca.api.profile.KeyParametersOption.DSAParametersOption;
import org.xipki.ca.api.profile.KeyParametersOption.ECParamatersOption;
import org.xipki.ca.api.profile.KeyParametersOption.RSAParametersOption;
import org.xipki.qa.ValidationIssue;
import org.xipki.security.util.AlgorithmUtil;
import org.xipki.security.util.X509Util;
import org.xipki.util.LruCache;
import org.xipki.util.exception.BadCertTemplateException;

import java.security.spec.InvalidKeySpecException;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;

import static org.xipki.util.Args.notNull;
import static org.xipki.util.CollectionUtil.isEmpty;

/**
 * Public key checker.
 *
 * @author Lijun Liao (xipki)
 * @since 2.0.0
 */

public class PublicKeyChecker {

  private static final Logger LOG = LoggerFactory.getLogger(PublicKeyChecker.class);

  private static final LruCache<ASN1ObjectIdentifier, Integer> EC_CURVEFIELD_SIZES = new LruCache<>(100);

  private final Map<ASN1ObjectIdentifier, KeyParametersOption> keyAlgorithms;

  public PublicKeyChecker(Map<ASN1ObjectIdentifier, KeyParametersOption> keyAlgorithms) {
    this.keyAlgorithms = keyAlgorithms;
  }

  public List<ValidationIssue> checkPublicKey(SubjectPublicKeyInfo publicKey, SubjectPublicKeyInfo requestedPublicKey) {
    notNull(publicKey, "publicKey");
    notNull(requestedPublicKey, "requestedPublicKey");

    List<ValidationIssue> resultIssues = new LinkedList<>();
    if (keyAlgorithms != null) {
      ValidationIssue issue = new ValidationIssue("X509.PUBKEY.SYN",
          "whether the public key in certificate is permitted");
      resultIssues.add(issue);
      try {
        checkPublicKey(publicKey);
      } catch (BadCertTemplateException ex) {
        issue.setFailureMessage(ex.getMessage());
      }
    }

    ValidationIssue issue = new ValidationIssue("X509.PUBKEY.REQ", "whether public key matches the request one");
    resultIssues.add(issue);
    SubjectPublicKeyInfo c14nRequestedPublicKey;
    try {
      c14nRequestedPublicKey = X509Util.toRfc3279Style(requestedPublicKey);
      if (!c14nRequestedPublicKey.equals(publicKey)) {
        issue.setFailureMessage("public key in the certificate does not equal the requested one");
      }
    } catch (InvalidKeySpecException ex) {
      issue.setFailureMessage("public key in request is invalid");
    }

    return resultIssues;
  } // method checkPublicKey

  private void checkPublicKey(SubjectPublicKeyInfo publicKey) throws BadCertTemplateException {
    if (isEmpty(keyAlgorithms)) {
      return;
    }

    ASN1ObjectIdentifier keyType = publicKey.getAlgorithm().getAlgorithm();
    if (!keyAlgorithms.containsKey(keyType)) {
      throw new BadCertTemplateException("key type " + keyType.getId() + " is not permitted");
    }

    KeyParametersOption keyParamsOption = keyAlgorithms.get(keyType);
    if (keyParamsOption instanceof AllowAllParametersOption) {
      return;
    } else if (keyParamsOption instanceof ECParamatersOption) {
      ECParamatersOption ecOption = (ECParamatersOption) keyParamsOption;
      // parameters
      ASN1Encodable algParam = publicKey.getAlgorithm().getParameters();
      ASN1ObjectIdentifier curveOid;

      if (algParam instanceof ASN1ObjectIdentifier) {
        curveOid = (ASN1ObjectIdentifier) algParam;
        if (!ecOption.allowsCurve(curveOid)) {
          throw new BadCertTemplateException("EC curve " + AlgorithmUtil.getCurveName(curveOid)
              + " (OID: " + curveOid.getId() + ") is not allowed");
        }
      } else {
        throw new BadCertTemplateException("only namedCurve EC public key is supported");
      }

      // point encoding
      if (ecOption.getPointEncodings() != null) {
        byte[] keyData = publicKey.getPublicKeyData().getBytes();
        if (keyData.length < 1) {
          throw new BadCertTemplateException("invalid publicKeyData");
        }
        byte pointEncoding = keyData[0];
        if (!ecOption.getPointEncodings().contains(pointEncoding)) {
          throw new BadCertTemplateException("not-accepted EC point encoding " + pointEncoding);
        }
      }

      try {
        checkECSubjectPublicKeyInfo(curveOid, publicKey.getPublicKeyData().getBytes());
      } catch (BadCertTemplateException ex) {
        throw ex;
      } catch (Exception ex) {
        LOG.debug("checkECSubjectPublicKeyInfo", ex);
        throw new BadCertTemplateException("invalid public key: " + ex.getMessage());
      }

      return;
    } else if (keyParamsOption instanceof RSAParametersOption) {
      RSAParametersOption rsaOption = (RSAParametersOption) keyParamsOption;

      ASN1Integer modulus;
      try {
        ASN1Sequence seq = ASN1Sequence.getInstance(publicKey.getPublicKeyData().getBytes());
        modulus = ASN1Integer.getInstance(seq.getObjectAt(0));
      } catch (IllegalArgumentException ex) {
        throw new BadCertTemplateException("invalid publicKeyData");
      }

      int modulusLength = modulus.getPositiveValue().bitLength();
      if ((rsaOption.allowsModulusLength(modulusLength))) {
        return;
      }
    } else if (keyParamsOption instanceof DSAParametersOption) {
      DSAParametersOption dsaOption = (DSAParametersOption) keyParamsOption;
      ASN1Encodable params = publicKey.getAlgorithm().getParameters();
      if (params == null) {
        throw new BadCertTemplateException("null Dss-Parms is not permitted");
      }

      int plength;
      int qlength;

      try {
        ASN1Sequence seq = ASN1Sequence.getInstance(params);
        plength = ASN1Integer.getInstance(seq.getObjectAt(0)).getPositiveValue().bitLength();
        qlength = ASN1Integer.getInstance(seq.getObjectAt(1)).getPositiveValue().bitLength();
      } catch (IllegalArgumentException | ArrayIndexOutOfBoundsException ex) {
        throw new BadCertTemplateException("illegal Dss-Parms");
      }

      if (dsaOption.allowsPlength(plength) && dsaOption.allowsQlength(qlength)) {
        return;
      }
    } else {
      String txt = (keyParamsOption == null) ? "null" : keyParamsOption.getClass().getName();
      throw new IllegalStateException("should not reach here, unknown keyParamsOption " + txt);
    }

    throw new BadCertTemplateException("the given publicKey is not permitted");
  } // method checkPublicKey

  private static void checkECSubjectPublicKeyInfo(ASN1ObjectIdentifier curveOid, byte[] encoded)
      throws BadCertTemplateException {
    Integer expectedLength = EC_CURVEFIELD_SIZES.get(curveOid);
    if (expectedLength == null) {
      X9ECParameters ecP = ECUtil.getNamedCurveByOid(curveOid);
      ECCurve curve = ecP.getCurve();
      expectedLength = (curve.getFieldSize() + 7) / 8;
      EC_CURVEFIELD_SIZES.put(curveOid, expectedLength);
    }

    switch (encoded[0]) {
      case 0x02: // compressed
      case 0x03: // compressed
        if (encoded.length != (expectedLength + 1)) {
          throw new BadCertTemplateException("incorrect length for compressed encoding");
        }
        break;
      case 0x04: // uncompressed
      case 0x06: // hybrid
      case 0x07: // hybrid
        if (encoded.length != (2 * expectedLength + 1)) {
          throw new BadCertTemplateException("incorrect length for uncompressed/hybrid encoding");
        }
        break;
      default:
        throw new BadCertTemplateException("invalid point encoding 0x" + Integer.toString(encoded[0], 16));
    } // end switch
  } // method checkECSubjectPublicKeyInfo

}
