package com.nimbusds.common.id;


import java.security.InvalidKeyException;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import javax.crypto.Mac;
import javax.crypto.SecretKey;
import javax.crypto.spec.SecretKeySpec;

import org.apache.commons.lang3.ArrayUtils;

import com.nimbusds.jose.util.Base64URL;


/**
 * Identifier with Hash-based (SHA-256) Message Authentication Code (HMAC).
 */
public class IdentifierWithHMAC extends BaseIdentifier {


        /**
         * The default byte length of generated identifiers.
         */
        public static final int DEFAULT_BYTE_LENGTH = 16;


        /**
         * Computes a SHA-256 based HMAC for the specified message, truncated
         * to the 128 left-most bits.
         *
         * @param message The message. Must not be {@code null}.
         * @param hmacKey The HMAC key. Must be at least 128 bits long.
         *
         * @return The HMAC, truncated to the 128 left-most bits.
         *
         * @throws RuntimeException If HMAC computation failed.
         */
        private static byte[] computeHMAC(final byte[] message, final SecretKey hmacKey) {

                if (hmacKey.getEncoded().length < 32) {
                        throw new IllegalArgumentException("The HMAC key must be at least 256 bits long");
                }

                Mac hmacComputer;

                try {
                        hmacComputer = Mac.getInstance("HmacSHA256");
                } catch (NoSuchAlgorithmException e) {
                        throw new RuntimeException(e.getMessage(), e);
                }

                SecretKeySpec secret_key = new SecretKeySpec(hmacKey.getEncoded(), "HmacSHA256");

                try {
                        hmacComputer.init(secret_key);
                } catch (InvalidKeyException e) {
                        throw new RuntimeException(e.getMessage(), e);
                }

                byte[] hmac256bit = hmacComputer.doFinal(message);

                // Truncate to 128 bits (keep left most)
                // https://tools.ietf.org/html/rfc2104#section-5
                return ArrayUtils.subarray(hmac256bit, 0, 16);
        }


        /**
         * Generates a new 128-bite secure random identifier and protects it
         * with a SHA-256 based HMAC.
         *
         * @param hmacKey The HMAC key. Must be at least 128 bits long.
         *
         * @return The generated identifier, BASE64-URL encoded, where a dot
         *         delimits the identifier from the HMAC.
         *
         * @throws RuntimeException If HMAC computation failed.
         */
        private static String generate(final SecretKey hmacKey) {

                byte[] n = new byte[DEFAULT_BYTE_LENGTH];
                SECURE_RANDOM.nextBytes(n);
                byte[] hmac = computeHMAC(n, hmacKey);
                return Base64URL.encode(n) + "." + Base64URL.encode(hmac);
        }


        /**
         * Creates a new identifier protected with a SHA-256 based HMAC.
         *
         * @param value   The identifier value. Must not empty or {@code null}.
         * @param hmacKey The HMAC key. Must be at least 128 bits long.
         */
        public IdentifierWithHMAC(final byte[] value, final SecretKey hmacKey) {
                super(Base64URL.encode(value) + "." + Base64URL.encode(computeHMAC(value, hmacKey)));
        }


        /**
         * Generates a new 128-bite secure random identifier and protects it
         * with a SHA-256 based HMAC.
         *
         * @param hmacKey The HMAC key. Must be at least 128 bits long.
         *
         * @throws RuntimeException If HMAC computation failed.
         */
        public IdentifierWithHMAC(final SecretKey hmacKey) {
                super(generate(hmacKey));
        }


        /**
         * Creates a new 128-bite identifier protected with a SHA-256 based
         * HMAC.
         *
         * @param value The identifier value. Must have been validated by the
         *              caller.
         */
        private IdentifierWithHMAC(final String value) {
                super(value);
        }


        @Override
        public boolean equals(final Object object) {

                return object instanceof IdentifierWithHMAC && this.toString().equals(object.toString());
        }


        /**
         * Parses and validates the specified identifier with HMAC protection.
         *
         * @param value   The identifier value to parse and validate. Must not
         *                be {@code null}.
         * @param hmacKey The HMAC key. Must be at least 128 bits long.
         *
         * @return The validated identifier.
         *
         * @throws InvalidIdentifierException If the identifier format is
         *                                    illegal or the HMAC is invalid.
         * @throws RuntimeException           If HMAC computation failed.
         */
        public static IdentifierWithHMAC parseAndValidate(final String value, final SecretKey hmacKey)
                throws InvalidIdentifierException {

                String[] parts = value.split("\\.");

                if (parts.length != 2) {
                        throw new InvalidIdentifierException("Illegal identifier with HMAC format");
                }

                if (parts[0].trim().isEmpty()) {
                        throw new InvalidIdentifierException("Missing identifier value");
                }

                if (parts[1].trim().isEmpty()) {
                        throw new InvalidIdentifierException("Missing HMAC for the identifier value");
                }

                byte[] n = new Base64URL(parts[0]).decode();
                
                if (n.length == 0) {
                        throw new InvalidIdentifierException("Illegal identifier value");
                }
                
                byte[] hmac = new Base64URL(parts[1]).decode();
                
                if (hmac.length == 0) {
                        throw new InvalidIdentifierException("Illegal HMAC value");
                }

                byte[] expectedHMAC = computeHMAC(n, hmacKey);

                if (! MessageDigest.isEqual(expectedHMAC, hmac)) {
                        throw new InvalidIdentifierException("Invalid HMAC");
                }

                return new IdentifierWithHMAC(value);
        }
}