/*
 * JBoss, Home of Professional Open Source
 * Copyright 2016 Red Hat, Inc., and individual contributors
 * as indicated by the @author tags.
 *
 * 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.wildfly.security.util;


import static org.wildfly.security.util.ElytronMessages.log;

import java.io.ByteArrayOutputStream;
import java.nio.charset.StandardCharsets;
import java.security.AlgorithmParameters;
import java.security.GeneralSecurityException;
import java.security.NoSuchAlgorithmException;
import java.security.Provider;
import java.security.Security;
import java.security.spec.AlgorithmParameterSpec;
import java.security.spec.InvalidParameterSpecException;
import java.text.Normalizer;

import javax.crypto.Cipher;
import javax.crypto.SecretKey;
import javax.crypto.SecretKeyFactory;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.PBEKeySpec;
import javax.crypto.spec.PBEParameterSpec;
import javax.crypto.spec.SecretKeySpec;

import org.wildfly.common.Assert;
import org.wildfly.common.codec.Alphabet;
import org.wildfly.common.codec.Base32Alphabet;
import org.wildfly.common.codec.Base64Alphabet;
import org.wildfly.common.iteration.ByteIterator;
import org.wildfly.common.iteration.CodePointIterator;

/**
 * Password Based Encryption utility class for tooling.
 * It provides builder to build PBE masked strings for usage with {@link org.wildfly.security.credential.store.CredentialStore}.
 *
 * @author <a href="mailto:pskopek@redhat.com">Peter Skopek</a>
 */
public final class PasswordBasedEncryptionUtil {

    // PicketBox compatibility related
    private static final char PAD = '_';
    private static final String REGEX = "^" + PAD + "{0,2}[0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz./]*$";
    private static final String DEFAULT_PICKETBOX_ALGORITHM = "PBEWithMD5AndDES";
    private static final String DEFAULT_PICKETBOX_INITIAL_KEY_MATERIAL = "somearbitrarycrazystringthatdoesnotmatter";

    private static final String DEFAULT_PBE_ALGORITHM = "PBEWithHmacSHA1andAES_128";

    private final Cipher cipher;
    private final AlgorithmParameters algorithmParameters;
    private final Alphabet alphabet;
    private final boolean picketBoxCompatibility;
    private final boolean usePadding;

    PasswordBasedEncryptionUtil(Cipher cipher, AlgorithmParameters algorithmParameters, Alphabet alphabet, boolean usePadding, boolean picketBoxCompatibility) {
        this.cipher = cipher;
        this.alphabet = alphabet;
        this.algorithmParameters = algorithmParameters;
        this.usePadding = usePadding;
        this.picketBoxCompatibility = picketBoxCompatibility;
    }

    PasswordBasedEncryptionUtil(Cipher cipher, AlgorithmParameters algorithmParameters, Alphabet alphabet) {
        this(cipher, algorithmParameters, alphabet, false, false);
    }

    /**
     * Encrypt a payload and encode the result using {@link Alphabet} given to builder.
     * All necessary parameters are supplied through {@link Builder}.
     * @param payload secret to encrypt
     * @return String encrypted and encoded using given parameters
     * @throws GeneralSecurityException when problem occurs like non-existent algorithm or similar problems
     */
    public String encryptAndEncode(char[] payload) throws GeneralSecurityException {
        return encodeUsingAlphabet(encrypt(charArrayEncode(payload)));
    }

    /**
     * Decode given payload and decrypt it to original.
     * All necessary parameters are supplied through {@link Builder}.
     * @param encodedPayload text to decode and decrypt
     * @return decrypted secret
     * @throws GeneralSecurityException when problem occurs like non-existent algorithm or similar problems
     */
    public char[] decodeAndDecrypt(String encodedPayload) throws GeneralSecurityException {
        return byteArrayDecode(decrypt(decodeUsingAlphabet(encodedPayload)));
    }

    /**
     * Returns algorithm parameters used in the process of encryption.
     * Might be useful to store them separately after encryption happened. It depends on used algorithm.
     * @return {@link AlgorithmParameters} as generated by encryption process
     */
    public AlgorithmParameters getAlgorithmParameters() {
        return algorithmParameters;
    }

    /**
     * Returns encrypted IV (initial vector) as generated by AES algorithm in the process of encryption.
     * Other algorithms are not using it.
     * In case of no such data available it returns {@code null}.
     * It uses already set {@link Alphabet} to encode it.
     * @return encoded form of IV or {@code null} when not available
     */
    public String getEncodedIV() {
        if (algorithmParameters != null) {
            try {
                PBEParameterSpec spec = algorithmParameters.getParameterSpec(PBEParameterSpec.class);
                AlgorithmParameterSpec algSpec = spec.getParameterSpec();
                if (algSpec instanceof IvParameterSpec) {
                    return encodeUsingAlphabet(((IvParameterSpec) algSpec).getIV());
                }
            } catch (InvalidParameterSpecException e) {
                return null;
            }
        }
        return null;
    }

    private byte[] decodeUsingAlphabet(String payload) {
        if (picketBoxCompatibility) {
            return picketBoxBase64Decode(payload);
        } else {
            ByteIterator byteIterator = isBase64(alphabet) ? CodePointIterator.ofString(payload).base64Decode(getAlphabet64(alphabet), usePadding)
                    : CodePointIterator.ofString(payload).base32Decode(getAlphabet32(alphabet));
            return byteIterator.drain();
        }
    }

    private String encodeUsingAlphabet(byte[] payload) {
        if (picketBoxCompatibility) {
            return picketBoxBased64Encode(payload);
        } else {
            CodePointIterator codePointIterator = isBase64(alphabet) ? ByteIterator.ofBytes(payload).base64Encode(getAlphabet64(alphabet), usePadding)
                    : ByteIterator.ofBytes(payload).base32Encode(getAlphabet32(alphabet));
            return codePointIterator.drainToString();
        }
    }

    private static boolean isBase64(Alphabet alphabet) {
        return alphabet instanceof Base64Alphabet;
    }

    private static Base64Alphabet getAlphabet64(Alphabet alphabet) {
        return (Base64Alphabet) alphabet;
    }

    private static Base32Alphabet getAlphabet32(Alphabet alphabet) {
        return (Base32Alphabet) alphabet;
    }

    private byte[] encrypt(byte[] payload) throws GeneralSecurityException {
        return cipher.doFinal(payload);
    }

    private byte[] decrypt(byte[] payload) throws GeneralSecurityException {
        return cipher.doFinal(payload);
    }

    private static char[] byteArrayDecode(byte[] buffer) {
        return new String(buffer, StandardCharsets.UTF_8).toCharArray();
    }

    private static byte[] charArrayEncode(char[] buffer) {
        return Normalizer.normalize(new String(buffer), Normalizer.Form.NFKC).getBytes(StandardCharsets.UTF_8);
    }

    private static byte[] picketBoxBase64Decode(String picketBoxBase64) {
        if (picketBoxBase64.length() == 0) {
            return new byte[0];
        }

        while (picketBoxBase64.length() % 4 != 0) {
            picketBoxBase64 = PAD + picketBoxBase64;
        }

        if (!picketBoxBase64.matches(REGEX)) {
            throw log.wrongBase64InPBCompatibleMode(picketBoxBase64);
        }

        ByteArrayOutputStream bos = new ByteArrayOutputStream((picketBoxBase64.length() * 3) / 4);
        for (int i = 0, n = picketBoxBase64.length(); i < n;) {
            int pos0 = PICKETBOX_COMPATIBILITY.decode(picketBoxBase64.charAt(i++));
            int pos1 = PICKETBOX_COMPATIBILITY.decode(picketBoxBase64.charAt(i++));
            int pos2 = PICKETBOX_COMPATIBILITY.decode(picketBoxBase64.charAt(i++));
            int pos3 = PICKETBOX_COMPATIBILITY.decode(picketBoxBase64.charAt(i++));
            if (pos0 > -1) {
                bos.write(((pos1 & 0x30) >>> 4) | (pos0 << 2));
            }
            if (pos1 > -1) {
                bos.write(((pos2 & 0x3c) >>> 2) | ((pos1 & 0xf) << 4));
            }
            bos.write(((pos2 & 3) << 6) | pos3);
        }
        return bos.toByteArray();
    }

    private String picketBoxBased64Encode(byte[] buffer) {
        int len = buffer.length, pos = len % 3, c;
        byte b0 = 0, b1 = 0, b2 = 0;
        StringBuffer sb = new StringBuffer();

        int i = 0;
        switch (pos) {
            case 2:
                b1 = buffer[i++];
                c = ((b0 & 3) << 4) | ((b1 & 0xf0) >>> 4);
                sb.appendCodePoint(PICKETBOX_COMPATIBILITY.encode(c));
            case 1:
                b2 = buffer[i++];
                c = ((b1 & 0xf) << 2) | ((b2 & 0xc0) >>> 6);
                sb.appendCodePoint(PICKETBOX_COMPATIBILITY.encode(c));
                c = b2 & 0x3f;
                sb.appendCodePoint(PICKETBOX_COMPATIBILITY.encode(c));
                break;
        }

        while (pos < len) {
            b0 = buffer[pos++];
            b1 = buffer[pos++];
            b2 = buffer[pos++];
            c = (b0 & 0xfc) >>> 2;
            sb.appendCodePoint(PICKETBOX_COMPATIBILITY.encode(c));
            c = ((b0 & 3) << 4) | ((b1 & 0xf0) >>> 4);
            sb.appendCodePoint(PICKETBOX_COMPATIBILITY.encode(c));
            c = ((b1 & 0xf) << 2) | ((b2 & 0xc0) >>> 6);
            sb.appendCodePoint(PICKETBOX_COMPATIBILITY.encode(c));
            c = b2 & 0x3f;
            sb.appendCodePoint(PICKETBOX_COMPATIBILITY.encode(c));
        }

        return sb.toString();
    }

    /**
     * Builder class to build {@link PasswordBasedEncryptionUtil} class with all necessary parameters to support
     * password based encryption algorithms.
     */
    public static class Builder {

        // algorithm names
        private String keyAlgorithm;
        private String transformation;
        private String parametersAlgorithm;
        // key parameters
        private int iteration = -1;
        private byte[] salt;
        private int keyLength = 0;
        private char[] password;   // actually initial key
        // cipher parameters
        private int cipherMode;
        private int cipherIteration = -1;
        private byte[] cipherSalt;

        private Provider provider;
        private Alphabet alphabet = Base64Alphabet.STANDARD;
        private boolean usePadding = false;
        private IvParameterSpec ivSpec;
        private String encodedIV;
        private AlgorithmParameters algorithmParameters;

        private boolean picketBoxCompatibility = false;

        /**
         * Set password to use to generate the encryption key
         * @param password the password
         * @return this Builder
         */
        public Builder password(char[] password) {
            this.password = password;
            return this;
        }

        /**
         * Set password to use to generate the encryption key
         * @param password the password
         * @return this Builder
         */
        public Builder password(String password) {
            this.password = password.toCharArray();
            return this;
        }

        /**
         * Set initialization vector for use with AES algorithms
         * @param iv the raw IV
         * @return this Builder
         */
        public Builder iv(byte[] iv) {
            ivSpec = new IvParameterSpec(iv);
            return this;
        }

        /**
         * Set initialization vector for use with AES algorithms
         * @param encodedIV IV encoded using {@link Alphabet} set in this builder (or default)
         * @return this Builder
         */
        public Builder iv(String encodedIV) {
            this.encodedIV = encodedIV;
            return this;
        }

        /**
         * Transformation name to use as {@code Cipher} parameter.
         * @param transformation the name of transformation
         * @return this Builder
         */
        public Builder transformation(String transformation) {
            this.transformation = transformation;
            return this;
        }

        /**
         * Set the name of parameter's algorithm to initialize the {@code Cipher}
         * @param parametersAlgorithm the name of parameter's algorithm
         * @return this Builder
         */
        public Builder parametersAlgorithm(String parametersAlgorithm) {
            this.parametersAlgorithm = parametersAlgorithm;
            return this;
        }

        /**
         * Set salt for key derivation.
         * @param salt the salt
         * @return this Builder
         */
        public Builder salt(String salt) {
            this.salt = salt.getBytes(StandardCharsets.UTF_8);
            return this;
        }

        /**
         * Set salt for key derivation.
         * @param salt the salt
         * @return this Builder
         */
        public Builder salt(byte[] salt) {
            this.salt = salt;
            return this;
        }

        /**
         * Use PicketBox compatibility mode for producing exact output as using PBE for MASK- purpose.
         *
         * Problem is that PicketBox is using different base64 than standard.
         * Default is {@code false}.
         *
         * @return this Builder
         */
        public Builder picketBoxCompatibility() {
            picketBoxCompatibility = true;
            return this;
        }

        /**
         * Use padding when encoding/decoding binary data.
         *
         * @return this Builder
         */
        public Builder encodingPadded() {
            usePadding = true;
            return this;
        }

        /**
         * Set number of iteration for key derivation.
         * @param iteration the number of iterations
         * @return this Builder
         */
        public Builder iteration(int iteration) {
            this.iteration = iteration;
            return this;
        }

        /**
         * Set the key derivation algorithm.
         * @param keyAlgorithm the algorithm
         * @return this Builder
         */
        public Builder keyAlgorithm(String keyAlgorithm) {
            this.keyAlgorithm = keyAlgorithm;
            return this;
        }

        /**
         * Set the key length.
         * @param keyLength the length
         * @return this Builder
         */
        public Builder keyLength(int keyLength) {
            this.keyLength = keyLength;
            return this;
        }

        /**
         * Set the number of iterations for {@code Cipher}
         * @param cipherIteration number of iterations
         * @return this Builder
         */
        public Builder cipherIteration(int cipherIteration) {
            this.cipherIteration = cipherIteration;
            return this;
        }

        /**
         * Set salt for the {@code Cipher}
         * @param cipherSalt the salt
         * @return this Builder
         */
        public Builder cipherSalt (byte[] cipherSalt) {
            this.cipherSalt = cipherSalt;
            return this;
        }

        /**
         * Set salt for the {@code Cipher}
         * @param cipherSalt the salt
         * @return this Builder
         */
        public Builder cipherSalt (String cipherSalt) {
            this.cipherSalt = cipherSalt.getBytes(StandardCharsets.UTF_8);
            return this;
        }

        /**
         * Set the JCA provider which contains all classes needed by built utility class.
         * @param provider the provider
         * @return this Builder
         */
        public Builder provider(Provider provider) {
            this.provider = provider;
            return this;
        }

        /**
         * Set the JCA provider name which contains all classes needed by built utility class.
         * @param providerName the provider name
         * @return this Builder
         */
        public Builder provider(String providerName) {
            Assert.checkNotNullParam("providerName", providerName);
            provider = Security.getProvider(providerName);
            if (provider == null) {
                throw log.securityProviderDoesnExist(providerName);
            }
            return this;
        }

        /**
         * Set the alphabet to encode/decode result of encryption/decryption.
         * @param alphabet the {@link Alphabet} instance
         * @return this Builder
         */
        public Builder alphabet(Alphabet alphabet) {
            this.alphabet = alphabet;
            return this;
        }

        /**
         * Set encryption mode for chosen {@code Cipher}
         * @return this Builder
         */
        public Builder encryptMode() {
            cipherMode = Cipher.ENCRYPT_MODE;
            return this;
        }

        /**
         * Set decryption mode for chosen {@code Cipher}
         * @return this Builder
         */
        public Builder decryptMode() {
            cipherMode = Cipher.DECRYPT_MODE;
            return this;
        }

        /**
         * Set algorithm parameters for {@code Cipher} initialization.
         * @param algorithmParameters the algorithm parameters instance in form required by the used {@code Cipher}
         * @return this Builder
         */
        public Builder algorithmParameters(AlgorithmParameters algorithmParameters) {
            if (this.algorithmParameters == null) {
                this.algorithmParameters = algorithmParameters;
            }
            return this;
        }

        private Cipher createAndInitCipher(SecretKey secretKey) throws GeneralSecurityException {
            Cipher cipher = provider == null ? Cipher.getInstance(transformation)
                    : Cipher.getInstance(transformation, provider);
            if (cipherMode == Cipher.ENCRYPT_MODE) {
                cipher.init(cipherMode, secretKey, generateAlgorithmParameters(parametersAlgorithm, cipherIteration, cipherSalt, null, provider));
                algorithmParameters = cipher.getParameters();
            } else {
                if (algorithmParameters != null) {
                    cipher.init(cipherMode, secretKey, algorithmParameters);
                } else {
                        cipher.init(cipherMode, secretKey, generateAlgorithmParameters(parametersAlgorithm, cipherIteration, cipherSalt, ivSpec, provider));
                }
            }
            return cipher;
        }

        private static AlgorithmParameters generateAlgorithmParameters(String algorithm, int iterationCount, byte[] salt, IvParameterSpec ivSpec, Provider provider) throws GeneralSecurityException {
            AlgorithmParameters tempParams = provider == null ? AlgorithmParameters.getInstance(algorithm)
                    : AlgorithmParameters.getInstance(algorithm, provider);
            PBEParameterSpec pbeParameterSpec = ivSpec != null ? new PBEParameterSpec(salt, iterationCount, ivSpec)
                    : new PBEParameterSpec(salt, iterationCount);
            tempParams.init(pbeParameterSpec);
            return tempParams;
        }

        private SecretKey deriveSecretKey() throws GeneralSecurityException {
            SecretKeyFactory secretKeyFactory;
            try {
                if (provider != null) {
                    secretKeyFactory = SecretKeyFactory.getInstance(keyAlgorithm, provider);
                } else {
                    secretKeyFactory = SecretKeyFactory.getInstance(keyAlgorithm);
                }
            } catch (NoSuchAlgorithmException e) {
                throw log.noSuchKeyAlgorithm(keyAlgorithm, e);
            }

            PBEKeySpec pbeKeySpec = keyLength == 0 ? new PBEKeySpec(password, salt, iteration)
                    : new PBEKeySpec(password, salt, iteration, keyLength);
            SecretKey partialKey = secretKeyFactory.generateSecret(pbeKeySpec);
            return new SecretKeySpec(partialKey.getEncoded(), transformation);
        }

        /**
         * Builds PBE utility class instance
         * @return PBE utility class instance {@link PasswordBasedEncryptionUtil}
         * @throws GeneralSecurityException when something goes wrong while initializing encryption related objects
         */
        public PasswordBasedEncryptionUtil build() throws GeneralSecurityException {
            if (picketBoxCompatibility) {
                // in compatible mode use Alphabet.PICKETBOX_COMPATIBILITY and no padding with proper key material and algorithm
                alphabet = PICKETBOX_COMPATIBILITY;
                usePadding = false;
                keyAlgorithm = DEFAULT_PICKETBOX_ALGORITHM;
                password = DEFAULT_PICKETBOX_INITIAL_KEY_MATERIAL.toCharArray();
            }

            if (iteration <= -1) {
                throw log.iterationCountNotSpecified();
            }
            if (salt == null) {
                throw log.saltNotSpecified();
            }
            if (password == null || password.length == 0) {
                throw log.initialKeyNotSpecified();
            }
            if (keyAlgorithm == null) {
                keyAlgorithm = DEFAULT_PBE_ALGORITHM;
            }
            if (transformation == null) {
                    transformation = keyAlgorithm;
            }
            if (parametersAlgorithm == null) {
                parametersAlgorithm = keyAlgorithm;
            }
            if (cipherSalt == null) {
                cipherSalt = salt;
            }
            if (cipherIteration == -1) {
                cipherIteration = iteration;
            }
            if (ivSpec == null && encodedIV != null) {
                ByteIterator byteIterator = isBase64(alphabet) ? CodePointIterator.ofString(encodedIV).base64Decode(getAlphabet64(alphabet), false)
                        : CodePointIterator.ofString(encodedIV).base32Decode(getAlphabet32(alphabet), false);
                ivSpec = new IvParameterSpec(byteIterator.drain());
            }

            return new PasswordBasedEncryptionUtil(createAndInitCipher(deriveSecretKey()), algorithmParameters, alphabet, usePadding, picketBoxCompatibility);
        }
    }

    /**
     * The alphabet used by PicketBox project base 64 encoding.
     * {@code 0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz./}
     */
    public static final Base64Alphabet PICKETBOX_COMPATIBILITY = new Base64Alphabet(false) {
        public int encode(int val) {
            if (val <= 9) {
                return '0' + val;
            } else if (val <= 35) {
                return 'A' + val - 10;
            } else if (val <= 61) {
                return 'a' + val - 36;
            } else if (val == 62) {
                return '.';
            } else {
                assert val == 63;
                return '/';
            }
        }

        public int decode(int codePoint) {
            if ('0' <= codePoint && codePoint <= '9') {
                return codePoint - '0';
            } else if ('A' <= codePoint && codePoint <= 'Z') {
                return codePoint - 'A' + 10;
            } else if ('a' <= codePoint && codePoint <= 'z') {
                return codePoint - 'a' + 36;
            } else if (codePoint == '.') {
                return 62;
            } else if (codePoint == '/') {
                return 63;
            } else {
                return -1;
            }
        }
    };
}
