/*
 * Copyright 2012-2017 Brian Campbell
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.jose4j.jws;


import org.jose4j.jca.ProviderContext;
import org.jose4j.jwa.CryptoPrimitive;
import org.jose4j.jwk.EllipticCurveJsonWebKey;
import org.jose4j.jwk.PublicJsonWebKey;
import org.jose4j.keys.BigEndianBigInteger;
import org.jose4j.keys.EllipticCurves;
import org.jose4j.lang.ByteUtil;
import org.jose4j.lang.ExceptionHelp;
import org.jose4j.lang.InvalidKeyException;
import org.jose4j.lang.JoseException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.math.BigInteger;
import java.security.Key;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.interfaces.ECKey;
import java.security.spec.ECParameterSpec;
import java.security.spec.EllipticCurve;


/**
 */
public class EcdsaUsingShaAlgorithm extends BaseSignatureAlgorithm implements JsonWebSignatureAlgorithm
{
    private final String curveName;
    private final int signatureByteLength;

    public EcdsaUsingShaAlgorithm(String id, String javaAlgo, String curveName, int signatureByteLength)
    {
        super(id, javaAlgo, EllipticCurveJsonWebKey.KEY_TYPE);
        this.curveName = curveName;
        this.signatureByteLength = signatureByteLength;
    }

    public boolean verifySignature(byte[] signatureBytes, Key key, byte[] securedInputBytes, ProviderContext providerContext) throws JoseException
    {
        // some pre-validation before calling the JCA to verify the signature
        // inspired by CVE-2022-21449 https://neilmadden.blog/2022/04/19/psychic-signatures-in-java/
        if (signatureBytes.length > signatureByteLength)
        {
            return false;
        }

        final byte[] rb = ByteUtil.leftHalf(signatureBytes);
        final BigInteger r = BigEndianBigInteger.fromBytes(rb);
        final byte[] sb = ByteUtil.rightHalf(signatureBytes);
        final BigInteger s = BigEndianBigInteger.fromBytes(sb);
        ECParameterSpec ecParams = EllipticCurves.getSpec(curveName);
        final BigInteger orderN = ecParams.getOrder();

        if (r.mod(orderN).equals(BigInteger.ZERO) || s.mod(orderN).equals(BigInteger.ZERO))
        {
            return false;
        }

        byte[] derEncodedSignatureBytes;
        try
        {
            derEncodedSignatureBytes = convertConcatenatedToDer(signatureBytes);
        }
        catch (IOException e)
        {
            throw new JoseException("Unable to convert R and S as a concatenated byte array to DER encoding.", e);
        }

        return super.verifySignature(derEncodedSignatureBytes, key, securedInputBytes, providerContext);
    }

    public byte[] sign(CryptoPrimitive cryptoPrimitive, byte[] securedInputBytes) throws JoseException
    {
        byte[] derEncodedSignatureBytes = super.sign(cryptoPrimitive, securedInputBytes);
        try
        {
            return convertDerToConcatenated(derEncodedSignatureBytes, signatureByteLength);
        }
        catch (IOException e)
        {
            throw new JoseException("Unable to convert DER encoding to R and S as a concatenated byte array.", e);
        }
    }

    /*
        The result of an ECDSA signature is the EC point (R, S), where R and S are unsigned (very large) integers.
        The JCA ECDSA signature implementation (sun.security.ec.ECDSASignature) produces and expects a DER encoding
        of R and S while JOSE/JWS wants R and S as a concatenated byte array. XML signatures (best I can tell) treats
        ECDSA similarly to JOSE and the code for the two methods that convert to and from DER and concatenated
        R and S was originally taken from org.apache.xml.security.algorithms.implementations.SignatureECDSA in the
        (Apache 2 licensed) Apache Santuario XML Security library. Some minor changes have been made to ensure the
        concatenated output left zero pads R & S to consistent length - i.e. the "octet sequence representations
        MUST NOT be shortened to omit any leading zero octets" per http://tools.ietf.org/html/draft-ietf-jose-json-web-algorithms-25#section-3.4

        Which seemed like a better idea than trying to write it myself or using sun.security.util.Der[Input/Output]Stream
        as sun.security.ec.ECDSASignature does or some other half-arsed approach.
     */

    // Convert the concatenation of R and S into DER encoding
    public static byte[] convertConcatenatedToDer(byte[] concatenatedSignatureBytes) throws IOException
    {
        int rawLen = concatenatedSignatureBytes.length/2;

        int i;

        for (i = rawLen; (i > 1) && (concatenatedSignatureBytes[rawLen - i] == 0); i--);

        int j = i;

        if (concatenatedSignatureBytes[rawLen - i] < 0)
        {
            j += 1;
        }

        int k;

        for (k = rawLen; (k > 1) && (concatenatedSignatureBytes[2*rawLen - k] == 0); k--);

        int l = k;

        if (concatenatedSignatureBytes[2*rawLen - k] < 0)
        {
            l += 1;
        }

        int len = 2 + j + 2 + l;
        if (len > 255)
        {
            throw new IOException("Invalid format of ECDSA signature");
        }
        int offset;
        byte derEncodedSignatureBytes[];
        if (len < 128)
        {
            derEncodedSignatureBytes = new byte[2 + 2 + j + 2 + l];
            offset = 1;
        }
        else
        {
            derEncodedSignatureBytes = new byte[3 + 2 + j + 2 + l];
            derEncodedSignatureBytes[1] = (byte) 0x81;
            offset = 2;
        }

        derEncodedSignatureBytes[0] = 48;
        derEncodedSignatureBytes[offset++] = (byte) len;
        derEncodedSignatureBytes[offset++] = 2;
        derEncodedSignatureBytes[offset++] = (byte) j;

        System.arraycopy(concatenatedSignatureBytes, rawLen - i, derEncodedSignatureBytes, (offset + j) - i, i);

        offset += j;

        derEncodedSignatureBytes[offset++] = 2;
        derEncodedSignatureBytes[offset++] = (byte) l;

        System.arraycopy(concatenatedSignatureBytes, 2*rawLen - k, derEncodedSignatureBytes, (offset + l) - k, k);

        return derEncodedSignatureBytes;
    }

    // Convert the DER encoding of R and S into a concatenation of R and S
    public static byte[] convertDerToConcatenated(byte derEncodedBytes[], int outputLength) throws IOException
    {

        if (derEncodedBytes.length < 8 || derEncodedBytes[0] != 48)
        {
            throw new IOException("Invalid format of ECDSA signature");
        }

        int offset;
        if (derEncodedBytes[1] > 0)
        {
            offset = 2;
        }
        else if (derEncodedBytes[1] == (byte) 0x81)
        {
            offset = 3;
        }
        else
        {
            throw new IOException("Invalid format of ECDSA signature");
        }

        byte rLength = derEncodedBytes[offset + 1];

        int i;
        for (i = rLength; (i > 0) && (derEncodedBytes[(offset + 2 + rLength) - i] == 0); i--);

        byte sLength = derEncodedBytes[offset + 2 + rLength + 1];

        int j;
        for (j = sLength; (j > 0) && (derEncodedBytes[(offset + 2 + rLength + 2 + sLength) - j] == 0); j--);

        int rawLen = Math.max(i, j);
        rawLen = Math.max(rawLen, outputLength/2);

        if ((derEncodedBytes[offset - 1] & 0xff) != derEncodedBytes.length - offset
            || (derEncodedBytes[offset - 1] & 0xff) != 2 + rLength + 2 + sLength
            || derEncodedBytes[offset] != 2
            || derEncodedBytes[offset + 2 + rLength] != 2)
        {
            throw new IOException("Invalid format of ECDSA signature");
        }
        
        byte concatenatedSignatureBytes[] = new byte[2*rawLen];

        System.arraycopy(derEncodedBytes, (offset + 2 + rLength) - i, concatenatedSignatureBytes, rawLen - i, i);
        System.arraycopy(derEncodedBytes, (offset + 2 + rLength + 2 + sLength) - j, concatenatedSignatureBytes, 2*rawLen - j, j);

        return concatenatedSignatureBytes;
    }

    public void validatePrivateKey(PrivateKey privateKey) throws InvalidKeyException
    {
        validateKeySpec(privateKey);
    }

    public void validatePublicKey(PublicKey publicKey) throws InvalidKeyException
    {
        validateKeySpec(publicKey);
    }

    private void validateKeySpec(Key key) throws InvalidKeyException
    {
        // some keys are valid for EC operations with the provider
        // (like with Sun PKCS11 provider https://bitbucket.org/b_c/jose4j/issues/77 )
        // but aren't actually an instance of ECKey so only check the spec when we can
        if (key instanceof ECKey)
        {
            ECKey ecKey = (ECKey) key;
            ECParameterSpec spec = ecKey.getParams();
            EllipticCurve curve = spec.getCurve();

            String name = EllipticCurves.getName(curve);

            if (!getCurveName().equals(name))
            {
                throw new InvalidKeyException(getAlgorithmIdentifier() + "/" + getJavaAlgorithm() + " expects a key using " +
                        getCurveName() + " but was " + name);
            }
        }
    }

    public String getCurveName()
    {
        return curveName;
    }

    public static class EcdsaP256UsingSha256 extends EcdsaUsingShaAlgorithm
    {
        public EcdsaP256UsingSha256()
        {
            super(AlgorithmIdentifiers.ECDSA_USING_P256_CURVE_AND_SHA256, "SHA256withECDSA", EllipticCurves.P_256, 64);
        }
    }

    public static class EcdsaP384UsingSha384 extends EcdsaUsingShaAlgorithm
    {
        public EcdsaP384UsingSha384()
        {
            super(AlgorithmIdentifiers.ECDSA_USING_P384_CURVE_AND_SHA384, "SHA384withECDSA", EllipticCurves.P_384, 96);
        }
    }

    public static class EcdsaP521UsingSha512 extends EcdsaUsingShaAlgorithm
    {
        public EcdsaP521UsingSha512()
        {
            super(AlgorithmIdentifiers.ECDSA_USING_P521_CURVE_AND_SHA512, "SHA512withECDSA", EllipticCurves.P_521, 132);
        }
    }

    public static class EcdsaSECP256K1UsingSha256 extends EcdsaUsingShaAlgorithm
    {
        Logger log = LoggerFactory.getLogger(getClass());

        public EcdsaSECP256K1UsingSha256()
        {
            super(AlgorithmIdentifiers.ECDSA_USING_SECP256K1_CURVE_AND_SHA256, "SHA256withECDSA", EllipticCurves.SECP_256K1, 64);
        }

        @Override
        public boolean isAvailable()
        {
            /*
                The SunEC provider deprecated (~Java 14) and removed (~Java 16) support for the "Legacy Elliptic Curves",
                which includes secp256k1.
                I was unable to figure out a nice provider-agnostic way of determining if a curve is supported
                so going with just attempting to sign some junk and see if it works.
             */

            if (super.isAvailable())
            {
                try
                {
                    PublicJsonWebKey jwk = PublicJsonWebKey.Factory.newPublicJwk("{\"kty\":\"EC\"," +
                            "\"x\":\"gi0g9DzM2SvjVV7iD_upIU0urmZRjpoIc4Efu8563y8\"," +
                            "\"y\":\"Y5K6GofrdlWNLlfT8-AEyJyVZ3yJJcGgkGroHQCAhmk\"," +
                            "\"crv\":\"secp256k1\"," +
                            "\"d\":\"Vd99BKh6pxt3mXSDJzHuVrCq52xBXAKVahbuFb6dqBc\"}");

                    CryptoPrimitive cryptoPrimitive = prepareForSign(jwk.getPrivateKey(), new ProviderContext());
                    byte[] sig = sign(cryptoPrimitive, new byte[]{2,6});
                    return sig != null;
                }
                catch (JoseException e)
                {
                    log.debug(getAlgorithmIdentifier() + " is not available due to " + ExceptionHelp.toStringWithCauses(e));
                    return false;
                }
            }
            return false;
        }
    }
}
