/*
 * Copyright (C) 2019 The Android Open Source Project
 *
 * 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 com.android.internal.net.ipsec.ike.message;

import android.annotation.StringDef;
import android.net.ipsec.ike.exceptions.IkeProtocolException;

import com.android.internal.annotations.VisibleForTesting;
import com.android.internal.net.ipsec.ike.crypto.IkeMacPrf;
import com.android.internal.net.ipsec.ike.exceptions.AuthenticationFailedException;
import com.android.internal.net.ipsec.ike.message.IkeAuthPayload.AuthMethod;

import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.nio.ByteBuffer;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.ProviderException;
import java.security.Signature;
import java.security.SignatureException;
import java.security.cert.X509Certificate;
import java.util.Arrays;

/**
 * IkeAuthDigitalSignPayload represents Authentication Payload using a specific or generic digital
 * signature authentication method.
 *
 * <p>If AUTH_METHOD_RSA_DIGITAL_SIGN is used, then the hash algorithm is SHA1. If
 * AUTH_METHOD_GENERIC_DIGITAL_SIGN is used, the signature algorithm and hash algorithm are
 * extracted from authentication data.
 *
 * @see <a href="https://tools.ietf.org/html/rfc7296#section-3.8">RFC 7296, Internet Key Exchange
 *     Protocol Version 2 (IKEv2)</a>
 * @see <a href="https://tools.ietf.org/html/rfc7427">RFC 7427, Signature Authentication in the
 *     Internet Key Exchange Version 2 (IKEv2)</a>
 */
public class IkeAuthDigitalSignPayload extends IkeAuthPayload {
    private static final String KEY_ALGO_NAME = "RSA";
    private static final byte SIGNATURE_ALGO_ASN1_BYTES_LEN = (byte) 15;
    private static final byte SIGNATURE_ALGO_ASN1_BYTES_LEN_LEN = (byte) 1;

    // Byte arrays of DER encoded identifier ASN.1 objects that indicates the algorithm used to
    // generate the signature, extracted from
    // <a href="https://tools.ietf.org/html/rfc7427#appendix-A"> RFC 7427. There is no need to
    // understand the encoding process. They are just constants to indicate the algorithm type.
    private static final byte[] PKI_ALGO_ID_DER_BYTES_RSA_SHA1 = {
        (byte) 0x30, (byte) 0x0d, (byte) 0x06, (byte) 0x09,
        (byte) 0x2a, (byte) 0x86, (byte) 0x48, (byte) 0x86,
        (byte) 0xf7, (byte) 0x0d, (byte) 0x01, (byte) 0x01,
        (byte) 0x05, (byte) 0x05, (byte) 0x00
    };
    private static final byte[] PKI_ALGO_ID_DER_BYTES_RSA_SHA2_256 = {
        (byte) 0x30, (byte) 0x0d, (byte) 0x06, (byte) 0x09,
        (byte) 0x2a, (byte) 0x86, (byte) 0x48, (byte) 0x86,
        (byte) 0xf7, (byte) 0x0d, (byte) 0x01, (byte) 0x01,
        (byte) 0x0b, (byte) 0x05, (byte) 0x00
    };
    private static final byte[] PKI_ALGO_ID_DER_BYTES_RSA_SHA2_384 = {
        (byte) 0x30, (byte) 0x0d, (byte) 0x06, (byte) 0x09,
        (byte) 0x2a, (byte) 0x86, (byte) 0x48, (byte) 0x86,
        (byte) 0xf7, (byte) 0x0d, (byte) 0x01, (byte) 0x01,
        (byte) 0x0c, (byte) 0x05, (byte) 0x00
    };
    private static final byte[] PKI_ALGO_ID_DER_BYTES_RSA_SHA2_512 = {
        (byte) 0x30, (byte) 0x0d, (byte) 0x06, (byte) 0x09,
        (byte) 0x2a, (byte) 0x86, (byte) 0x48, (byte) 0x86,
        (byte) 0xf7, (byte) 0x0d, (byte) 0x01, (byte) 0x01,
        (byte) 0x0d, (byte) 0x05, (byte) 0x00
    };

    // Length of ASN.1 object length field.
    private static final int SIGNATURE_ALGO_ASN1_LEN_LEN = 1;

    // Currently we only support RSA for signature algorithm.
    @Retention(RetentionPolicy.SOURCE)
    @StringDef({
        SIGNATURE_ALGO_RSA_SHA1,
        SIGNATURE_ALGO_RSA_SHA2_256,
        SIGNATURE_ALGO_RSA_SHA2_384,
        SIGNATURE_ALGO_RSA_SHA2_512
    })
    @VisibleForTesting
    @interface SignatureAlgo {}

    public static final String SIGNATURE_ALGO_RSA_SHA1 = "SHA1withRSA";
    public static final String SIGNATURE_ALGO_RSA_SHA2_256 = "SHA256withRSA";
    public static final String SIGNATURE_ALGO_RSA_SHA2_384 = "SHA384withRSA";
    public static final String SIGNATURE_ALGO_RSA_SHA2_512 = "SHA512withRSA";

    // IKEv2 types for hash algorithms.
    public static final short HASH_ALGORITHM_RSA_SHA1 = 1;
    public static final short HASH_ALGORITHM_RSA_SHA2_256 = 2;
    public static final short HASH_ALGORITHM_RSA_SHA2_384 = 3;
    public static final short HASH_ALGORITHM_RSA_SHA2_512 = 4;
    public static final short[] ALL_SIGNATURE_ALGO_TYPES =
            new short[] {
                HASH_ALGORITHM_RSA_SHA1,
                HASH_ALGORITHM_RSA_SHA2_256,
                HASH_ALGORITHM_RSA_SHA2_384,
                HASH_ALGORITHM_RSA_SHA2_512
            };

    public final String signatureAndHashAlgos;
    public final byte[] signature;

    protected IkeAuthDigitalSignPayload(
            boolean critical, @AuthMethod int authMethod, byte[] authData)
            throws IkeProtocolException {
        super(critical, authMethod);
        switch (authMethod) {
            case AUTH_METHOD_RSA_DIGITAL_SIGN:
                signatureAndHashAlgos = SIGNATURE_ALGO_RSA_SHA1;
                signature = authData;
                break;
            case AUTH_METHOD_GENERIC_DIGITAL_SIGN:
                ByteBuffer inputBuffer = ByteBuffer.wrap(authData);

                // Get signature algorithm.
                int signAlgoLen = Byte.toUnsignedInt(inputBuffer.get());
                byte[] signAlgoBytes = new byte[signAlgoLen];
                inputBuffer.get(signAlgoBytes);
                signatureAndHashAlgos = bytesToJavaStandardSignAlgoName(signAlgoBytes);

                // Get signature.
                signature = new byte[authData.length - SIGNATURE_ALGO_ASN1_LEN_LEN - signAlgoLen];
                inputBuffer.get(signature);
                break;
            default:
                throw new IllegalArgumentException("Unrecognized authentication method.");
        }
    }

    /**
     * Construct IkeAuthDigitalSignPayload for an outbound IKE packet.
     *
     * <p>Since IKE library is always a client, outbound IkeAuthDigitalSignPayload always signs IKE
     * initiator's SignedOctets, which is concatenation of the IKE_INIT request message, the Nonce
     * of IKE responder and the signed ID-Initiator payload body.
     *
     * <p>Caller MUST validate that the signatureAlgoName is supported by IKE library.
     *
     * @param signatureAlgoName the name of the algorithm requested. See the Signature section in
     *     the <a href= "{@docRoot}/../technotes/guides/security/StandardNames.html#Signature"> Java
     *     Cryptography Architecture Standard Algorithm Name Documentation</a> for information about
     *     standard algorithm names.
     * @param privateKey the private key of the identity whose signature is going to be generated.
     * @param ikeInitBytes IKE_INIT request for calculating IKE initiator's SignedOctets.
     * @param nonce nonce of IKE responder for calculating IKE initiator's SignedOctets.
     * @param idPayloadBodyBytes ID-Initiator payload body for calculating IKE initiator's
     *     SignedOctets.
     * @param ikePrf the negotiated PRF.
     * @param prfKeyBytes the negotiated PRF initiator key.
     */
    public IkeAuthDigitalSignPayload(
            String signatureAlgoName,
            PrivateKey privateKey,
            byte[] ikeInitBytes,
            byte[] nonce,
            byte[] idPayloadBodyBytes,
            IkeMacPrf ikePrf,
            byte[] prfKeyBytes) {
        super(false, IkeAuthPayload.AUTH_METHOD_GENERIC_DIGITAL_SIGN);
        byte[] dataToSignBytes =
                getSignedOctets(ikeInitBytes, nonce, idPayloadBodyBytes, ikePrf, prfKeyBytes);

        try {
            Signature signGen = Signature.getInstance(signatureAlgoName);
            signGen.initSign(privateKey);
            signGen.update(dataToSignBytes);

            signature = signGen.sign();
            signatureAndHashAlgos = signatureAlgoName;
        } catch (SignatureException | InvalidKeyException e) {
            throw new IllegalArgumentException("Signature generation failed", e);
        } catch (NoSuchAlgorithmException e) {
            throw new ProviderException(
                    "Security Provider does not support "
                            + KEY_ALGO_NAME
                            + " or "
                            + signatureAlgoName);
        }
    }

    private byte[] javaStandardSignAlgoNameToAsn1Bytes(String javaSignatureAndHashAlgo) {
        switch (javaSignatureAndHashAlgo) {
            case SIGNATURE_ALGO_RSA_SHA1:
                return PKI_ALGO_ID_DER_BYTES_RSA_SHA1;
            case SIGNATURE_ALGO_RSA_SHA2_256:
                return PKI_ALGO_ID_DER_BYTES_RSA_SHA2_256;
            case SIGNATURE_ALGO_RSA_SHA2_384:
                return PKI_ALGO_ID_DER_BYTES_RSA_SHA2_384;
            case SIGNATURE_ALGO_RSA_SHA2_512:
                return PKI_ALGO_ID_DER_BYTES_RSA_SHA2_512;
            default:
                throw new IllegalArgumentException("Impossible! We used an unsupported algo");
        }
    }

    private String bytesToJavaStandardSignAlgoName(byte[] signAlgoBytes)
            throws AuthenticationFailedException {
        if (Arrays.equals(PKI_ALGO_ID_DER_BYTES_RSA_SHA1, signAlgoBytes)) {
            return SIGNATURE_ALGO_RSA_SHA1;
        } else if (Arrays.equals(PKI_ALGO_ID_DER_BYTES_RSA_SHA2_256, signAlgoBytes)) {
            return SIGNATURE_ALGO_RSA_SHA2_256;
        } else if (Arrays.equals(PKI_ALGO_ID_DER_BYTES_RSA_SHA2_384, signAlgoBytes)) {
            return SIGNATURE_ALGO_RSA_SHA2_384;
        } else if (Arrays.equals(PKI_ALGO_ID_DER_BYTES_RSA_SHA2_512, signAlgoBytes)) {
            return SIGNATURE_ALGO_RSA_SHA2_512;
        } else {
            throw new AuthenticationFailedException(
                    "Unrecognized ASN.1 objects for Signature algorithm and Hash");
        }
    }

    /**
     * Verify received signature in an inbound IKE packet.
     *
     * <p>Since IKE library is always a client, inbound IkeAuthDigitalSignPayload always signs IKE
     * responder's SignedOctets, which is concatenation of the IKE_INIT response message, the Nonce
     * of IKE initiator and the signed ID-Responder payload body.
     *
     * @param certificate received end certificate to verify the signature.
     * @param ikeInitBytes IKE_INIT response for calculating IKE responder's SignedOctets.
     * @param nonce nonce of IKE initiator for calculating IKE responder's SignedOctets.
     * @param idPayloadBodyBytes ID-Responder payload body for calculating IKE responder's
     *     SignedOctets.
     * @param ikePrf the negotiated PRF.
     * @param prfKeyBytes the negotiated PRF responder key.
     * @throws AuthenticationFailedException if received signature verification failed.
     */
    public void verifyInboundSignature(
            X509Certificate certificate,
            byte[] ikeInitBytes,
            byte[] nonce,
            byte[] idPayloadBodyBytes,
            IkeMacPrf ikePrf,
            byte[] prfKeyBytes)
            throws AuthenticationFailedException {
        byte[] dataToSignBytes =
                getSignedOctets(ikeInitBytes, nonce, idPayloadBodyBytes, ikePrf, prfKeyBytes);

        try {
            Signature signValidator = Signature.getInstance(signatureAndHashAlgos);
            signValidator.initVerify(certificate);
            signValidator.update(dataToSignBytes);

            if (!signValidator.verify(signature)) {
                throw new AuthenticationFailedException("Signature verification failed.");
            }
        } catch (SignatureException | InvalidKeyException e) {
            throw new AuthenticationFailedException(e);
        } catch (NoSuchAlgorithmException e) {
            throw new ProviderException(
                    "Security Provider does not support " + signatureAndHashAlgos);
        }
    }

    @Override
    protected void encodeAuthDataToByteBuffer(ByteBuffer byteBuffer) {
        if (authMethod == AUTH_METHOD_GENERIC_DIGITAL_SIGN) {
            byteBuffer.put(SIGNATURE_ALGO_ASN1_BYTES_LEN);
            byteBuffer.put(javaStandardSignAlgoNameToAsn1Bytes(signatureAndHashAlgos));
        }
        byteBuffer.put(signature);
    }

    @Override
    protected int getAuthDataLength() {
        if (authMethod == AUTH_METHOD_GENERIC_DIGITAL_SIGN) {
            return SIGNATURE_ALGO_ASN1_BYTES_LEN_LEN
                    + SIGNATURE_ALGO_ASN1_BYTES_LEN
                    + signature.length;
        }
        return signature.length;
    }

    @Override
    public String getTypeString() {
        return "Auth(Digital Sign)";
    }
}