package meerkat.crypto.concrete;

import com.google.protobuf.ByteString;
import com.google.protobuf.Message;
import meerkat.protobuf.Crypto;
import meerkat.protobuf.Crypto.Signature;
import org.spongycastle.asn1.ASN1Integer;
import org.spongycastle.asn1.DERSequence;
import org.spongycastle.crypto.params.AsymmetricKeyParameter;
import org.spongycastle.crypto.params.ECPrivateKeyParameters;
import org.spongycastle.crypto.signers.DSAKCalculator;
import org.spongycastle.crypto.signers.ECDSASigner;
import org.spongycastle.crypto.signers.HMacDSAKCalculator;
import org.spongycastle.jcajce.provider.asymmetric.util.ECUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.io.InputStream;
import java.math.BigInteger;
import java.security.*;
import java.security.cert.*;
import java.security.cert.Certificate;
import java.security.interfaces.ECPrivateKey;


/**
 * Sign and verify digital signatures.
 *
 * Uses deterministic ECDSA signatures as per RFC 6979
 *
 * This class uses BouncyCastle directly, so will not work with arbitrary PKCS11 providers.
 *
 * This class is not thread-safe (each thread should have its own instance).
 */
public class ECDSADeterministicSignature extends ECDSASignature {
    final Logger logger = LoggerFactory.getLogger(getClass());

    /**
     * Digest of message contents for deterministic signing.
     */
    SHA256Digest msgDigest = new SHA256Digest();

    /**
     * The actual signing implementation. (used only signing -- superclass is used for verification)
     */
    ECDSASigner deterministicSigner;

    /**
     * Output the DER encoding of the ASN.1 sequence r,s
     * @param r
     * @param s
     * @return
     */
    public static byte[] derEncodeSignature(BigInteger r, BigInteger s) {
        ASN1Integer[] rs = {new ASN1Integer(r), new ASN1Integer(s)};
        DERSequence seq = new DERSequence(rs);

        try {
            return seq.getEncoded();
        } catch (IOException e) {
            throw new RuntimeException("Should never happen! DER Encoding exception", e);
        }
    }

    public ECDSADeterministicSignature() {
        DSAKCalculator kCalk = new HMacDSAKCalculator(new org.spongycastle.crypto.digests.SHA256Digest());
        deterministicSigner = new ECDSASigner(kCalk);
    }

    @Override
    public void loadSigningCertificate(Certificate cert, PrivateKey key) throws InvalidKeyException, CertificateException {
        super.loadSigningCertificate(cert, key);

        if (!(loadedSigningKey instanceof ECPrivateKey)) {
            logger.error("Wrong private key type (expected ECPrivateKey, got {})", loadedSigningKey.getClass());
            throw new CertificateException("Wrong signing key type!");
        }
        ECPrivateKey eckey = (ECPrivateKey) loadedSigningKey;

        AsymmetricKeyParameter baseParams;
        baseParams = ECUtil.generatePrivateKeyParameter(eckey);

        if (!(baseParams instanceof ECPrivateKeyParameters)) {
            logger.error("Error converting to bouncycastle type! (got {})", baseParams.getClass());
            throw new InvalidKeyException("Wrong signing key type!");
        }

        ECPrivateKeyParameters params = (ECPrivateKeyParameters) baseParams;

        deterministicSigner.init(true, params);
    }

    /**
     * Add the list of messages to the stream that is being verified/signed.
     * Messages are prepended with their length in 32-bit big-endian format.
     *
     * @param msg
     * @throws SignatureException
     */
    @Override
    public void updateContent(Message msg) throws SignatureException {
        assert msg != null;

        // We're doing twice the digest work so that we also support verification with the same update.
        // If this becomes a problem, we can decide which way to update based on which init was called.
        super.updateContent(msg);
        msgDigest.update(msg);
    }

    @Override
    public void updateContent(InputStream in) throws IOException, SignatureException {
        ByteString inStr = ByteString.readFrom(in);
        signer.update(inStr.asReadOnlyByteBuffer());
        msgDigest.update(inStr);
    }

    @Override
    public Signature sign() throws SignatureException {
        Signature.Builder sig = Signature.newBuilder();
        sig.setType(Crypto.SignatureType.ECDSA);

        BigInteger[] rawSig = deterministicSigner.generateSignature(msgDigest.digest());

        sig.setData(ByteString.copyFrom(derEncodeSignature(rawSig[0], rawSig[1])));

        sig.setSignerId(loadedSigningKeyId);
        return sig.build();
    }
}
