/*
 * Copyright (C) 2022 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.server.nearby.util.encryption;

import static com.android.server.nearby.NearbyService.TAG;

import android.util.Log;

import androidx.annotation.Nullable;

import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;

import javax.crypto.Mac;
import javax.crypto.spec.SecretKeySpec;

/** Class for encryption/decryption functionality. */
public abstract class Cryptor {

    /** AES only supports key sizes of 16, 24 or 32 bytes. */
    static final int AUTHENTICITY_KEY_BYTE_SIZE = 16;

    private static final String HMAC_SHA256_ALGORITHM = "HmacSHA256";

    /**
     * Encrypt the provided data blob.
     *
     * @param data data blob to be encrypted.
     * @param salt used for IV
     * @param secretKeyBytes secrete key accessed from credentials
     * @return encrypted data, {@code null} if failed to encrypt.
     */
    @Nullable
    public byte[] encrypt(byte[] data, byte[] salt, byte[] secretKeyBytes) {
        return data;
    }

    /**
     * Decrypt the original data blob from the provided byte array.
     *
     * @param encryptedData data blob to be decrypted.
     * @param salt used for IV
     * @param secretKeyBytes secrete key accessed from credentials
     * @return decrypted data, {@code null} if failed to decrypt.
     */
    @Nullable
    public byte[] decrypt(byte[] encryptedData, byte[] salt, byte[] secretKeyBytes) {
        return encryptedData;
    }

    /**
     * Generates a digital signature for the data.
     *
     * @return signature {@code null} if failed to sign
     */
    @Nullable
    public byte[] sign(byte[] data, byte[] key) {
        return new byte[0];
    }

    /**
     * Verifies the signature generated by data and key, with the original signed data
     */
    public boolean verify(byte[] data, byte[] key, byte[] signature) {
        return true;
    }

    /**
     * @return length of the signature generated
     */
    public int getSignatureLength() {
        return 0;
    }

    /**
     * A HAMC sha256 based HKDF algorithm to pseudo randomly hash data and salt into a byte array of
     * given size.
     */
    // Based on google3/third_party/tink/java/src/main/java/com/google/crypto/tink/subtle/Hkdf.java
    @Nullable
    static byte[] computeHkdf(byte[] ikm, byte[] salt, int size) {
        Mac mac;
        try {
            mac = Mac.getInstance(HMAC_SHA256_ALGORITHM);
        } catch (NoSuchAlgorithmException e) {
            Log.w(TAG, "HMAC_SHA256_ALGORITHM is not supported.", e);
            return null;
        }

        if (size > 255 * mac.getMacLength()) {
            Log.w(TAG, "Size too large.");
            return null;
        }

        if (salt.length == 0) {
            Log.w(TAG, "Salt cannot be empty.");
            return null;
        }

        try {
            mac.init(new SecretKeySpec(salt, HMAC_SHA256_ALGORITHM));
        } catch (InvalidKeyException e) {
            Log.w(TAG, "Invalid key.", e);
            return null;
        }

        byte[] prk = mac.doFinal(ikm);
        byte[] result = new byte[size];
        try {
            mac.init(new SecretKeySpec(prk, HMAC_SHA256_ALGORITHM));
        } catch (InvalidKeyException e) {
            Log.w(TAG, "Invalid key.", e);
            return null;
        }

        byte[] digest = new byte[0];
        int ctr = 1;
        int pos = 0;
        while (true) {
            mac.update(digest);
            mac.update((byte) ctr);
            digest = mac.doFinal();
            if (pos + digest.length < size) {
                System.arraycopy(digest, 0, result, pos, digest.length);
                pos += digest.length;
                ctr++;
            } else {
                System.arraycopy(digest, 0, result, pos, size - pos);
                break;
            }
        }

        return result;
    }
}
