/*
 * Copyright 2018 Scytl Secure Electronic Voting SA
 *
 * All rights reserved
 *
 * See our extended copyright notice in *file 'Copyright.txt' which is part of this source code package
 */
package com.scytl.products.ov.channel.cc.commons.keymanagement;

import static java.text.MessageFormat.format;
import static java.util.Arrays.fill;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.nio.ByteBuffer;
import java.nio.CharBuffer;
import java.nio.charset.StandardCharsets;
import java.security.GeneralSecurityException;
import java.security.KeyManagementException;
import java.security.KeyStore;
import java.security.KeyStore.PasswordProtection;
import java.security.KeyStore.PrivateKeyEntry;
import java.security.KeyStoreException;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.UnrecoverableEntryException;
import java.security.UnrecoverableKeyException;
import java.security.cert.X509Certificate;

import com.scytl.cryptolib.api.asymmetric.AsymmetricServiceAPI;
import com.scytl.cryptolib.api.exceptions.GeneralCryptoLibException;
import com.scytl.cryptolib.api.stores.StoresServiceAPI;
import com.scytl.cryptolib.api.stores.bean.KeyStoreType;
import com.scytl.cryptolib.elgamal.bean.ElGamalPrivateKey;
import com.scytl.cryptolib.elgamal.bean.ElGamalPublicKey;

/**
 * Implementation of {@link Codec}.
 */
class CodecImpl implements Codec {
    private static final String NODE_CA_ALIAS = "CA";

    private static final String NODE_ENCRYPTION_ALIAS = "encryption";

    private static final String NODE_LOG_SIGNING_ALIAS = "logSigning";

    private static final String NODE_LOG_ENCRYPTION_ALIAS =
        "logEncryption";

    private static final String ELECTION_SIGNING_ALIAS = "signing";

    private final StoresServiceAPI storesService;

    private final AsymmetricServiceAPI asymmetricService;

    /**
     * Constructor.
     * 
     * @param storesService
     *            the stores service
     * @param asymmetricService
     *            the asymmetric service.
     */
    public CodecImpl(StoresServiceAPI storesService,
            AsymmetricServiceAPI asymmetricService) {
        this.storesService = storesService;
        this.asymmetricService = asymmetricService;
    }

    private static PrivateKeyEntry decodePrivateKeyEntry(KeyStore store,
            String alias, PasswordProtection protection)
            throws InvalidPasswordException, KeyManagementException {
        try {
            if (!store.entryInstanceOf(alias, PrivateKeyEntry.class)) {
                throw new KeyManagementException(format(
                    "Key entry ''{0}'' is missing or invalid.", alias));
            }
            return (PrivateKeyEntry) store.getEntry(alias, protection);
        } catch (UnrecoverableKeyException e) {
            throw new InvalidPasswordException("Invalid key password.", e);
        } catch (KeyStoreException | NoSuchAlgorithmException
                | UnrecoverableEntryException e) {
            throw new KeyManagementException(
                "Failed to decode private key entry.", e);
        }
    }

    @Override
    public ElectionSigningKeys decodeElectionSigningKeys(byte[] bytes,
            PasswordProtection protection)
            throws InvalidPasswordException, KeyManagementException {
        KeyStore store = decodeKeyStore(bytes, protection);
        PrivateKeyEntry entry = decodePrivateKeyEntry(store,
            ELECTION_SIGNING_ALIAS, protection);
        return new ElectionSigningKeys(entry.getPrivateKey(),
            (X509Certificate[]) entry.getCertificateChain());
    }

    @Override
	public ElGamalPrivateKey decodeElGamalPrivateKey(byte[] bytes, PrivateKey encryptionKey)
			throws KeyManagementException {
    	
    	try {
        	byte[] decryptedBytes = asymmetricService.decrypt(encryptionKey, bytes);
			String json = new String(decryptedBytes, StandardCharsets.UTF_8);
			return ElGamalPrivateKey.fromJson(json);
		} catch (GeneralCryptoLibException e) {
			throw new KeyManagementException("Failed to decode ElGamal private key.", e);
		}
	}

    @Override
    public ElGamalPublicKey decodeElGamalPublicKey(byte[] bytes)
            throws KeyManagementException {
        String json = new String(bytes, StandardCharsets.UTF_8);
        try {
            return ElGamalPublicKey.fromJson(json);
        } catch (GeneralCryptoLibException e) {
            throw new KeyManagementException(
                "Failed to decode ElGamal public key.", e);
        }
    }

    @Override
    public NodeKeys decodeNodeKeys(byte[] bytes,
            PasswordProtection password)
            throws InvalidPasswordException, KeyManagementException {
        KeyStore store = decodeKeyStore(bytes, password);
        NodeKeys.Builder builder = new NodeKeys.Builder();
        builder.setCAKeys(
            decodePrivateKeyEntry(store, NODE_CA_ALIAS, password));
        builder.setEncryptionKeys(
            decodePrivateKeyEntry(store, NODE_ENCRYPTION_ALIAS, password));
        builder.setLogSigningKeys(decodePrivateKeyEntry(store,
            NODE_LOG_SIGNING_ALIAS, password));
        builder.setLogEncryptionKeys(decodePrivateKeyEntry(store,
            NODE_LOG_ENCRYPTION_ALIAS, password));
        return builder.build();
    }

    @Override
	public PasswordProtection decodePassword(byte[] bytes, PrivateKey encryptionKey) throws KeyManagementException {

		byte[] decryptedBytes;
		try {
			decryptedBytes = asymmetricService.decrypt(encryptionKey, bytes);
		} catch (GeneralCryptoLibException e) {
			throw new KeyManagementException("Failed to decode password.", e);
		}
		ByteBuffer byteBuffer = ByteBuffer.wrap(decryptedBytes);
		try {
			CharBuffer charBuffer = StandardCharsets.UTF_8.decode(byteBuffer);
			try {
				return new PasswordProtection(charBuffer.array());
			} finally {
				fill(charBuffer.array(), '\u0000');
			}
		} finally {
			fill(byteBuffer.array(), (byte) 0);
		}
	}

    @Override
    public byte[] encodeElectionSigningKeys(ElectionSigningKeys keys,
            PasswordProtection protection)
            throws KeyManagementException {
        try {
            KeyStore store =
                storesService.createKeyStore(KeyStoreType.PKCS12);
            store.setKeyEntry(ELECTION_SIGNING_ALIAS, keys.privateKey(),
                protection.getPassword(), keys.certificateChain());
            ByteArrayOutputStream stream = new ByteArrayOutputStream();
            try {
                store.store(stream, protection.getPassword());
            } finally {
                stream.close();
            }
            return stream.toByteArray();
        } catch (IOException | GeneralCryptoLibException
                | GeneralSecurityException e) {
            throw new KeyManagementException(
                "Failed to encode election signing keys.", e);
        }
    }

    @Override
    public byte[] encodeElGamalPrivateKey(ElGamalPrivateKey key,
            PublicKey encryptionKey)
            throws KeyManagementException {
        try {
            byte[] bytes = key.toJson().getBytes(StandardCharsets.UTF_8);
            return asymmetricService.encrypt(encryptionKey, bytes);
        } catch (GeneralCryptoLibException e) {
            throw new KeyManagementException(
                "Failed to encode ElGamal private key.", e);
        }
    }

    @Override
    public byte[] encodeElGamalPublicKey(ElGamalPublicKey key)
            throws KeyManagementException {
        try {
            return key.toJson().getBytes(StandardCharsets.UTF_8);
        } catch (GeneralCryptoLibException e) {
            throw new KeyManagementException(
                "Failed to encode ElGamal public key.", e);
        }
    }

    @Override
    public byte[] encodeNodeKeys(NodeKeys keys,
            PasswordProtection protection)
            throws KeyManagementException {
        try {
            KeyStore store =
                storesService.createKeyStore(KeyStoreType.PKCS12);
            store.setKeyEntry(NODE_CA_ALIAS, keys.caPrivateKey(),
                protection.getPassword(), keys.caCertificateChain());
            store.setKeyEntry(NODE_ENCRYPTION_ALIAS,
                keys.encryptionPrivateKey(), protection.getPassword(),
                keys.encryptionCertificateChain());
            store.setKeyEntry(NODE_LOG_SIGNING_ALIAS,
                keys.logSigningPrivateKey(), protection.getPassword(),
                keys.logSigningCertificateChain());
            store.setKeyEntry(NODE_LOG_ENCRYPTION_ALIAS,
                keys.logEncryptionPrivateKey(), protection.getPassword(),
                keys.logEncryptionCertificateChain());
            ByteArrayOutputStream stream = new ByteArrayOutputStream();
            try {
                store.store(stream, protection.getPassword());
            } finally {
                stream.close();
            }
            return stream.toByteArray();
        } catch (IOException | GeneralCryptoLibException
                | GeneralSecurityException e) {
            throw new KeyManagementException(
                "Failed to encode election signing keys.", e);
        }
    }

    @Override
    public byte[] encodePassword(PasswordProtection protection,
            PublicKey encryptionKey)
            throws KeyManagementException {
        CharBuffer charBuffer = CharBuffer.wrap(protection.getPassword());
        ByteBuffer byteBuffer = StandardCharsets.UTF_8.encode(charBuffer);
        byte[] bytes = new byte[byteBuffer.limit()];
        byteBuffer.get(bytes);
        try {
            return asymmetricService.encrypt(encryptionKey, bytes);
        } catch (GeneralCryptoLibException e) {
            throw new KeyManagementException("Failed to encode password.",
                e);
        } finally {
            fill(bytes, (byte) 0);
            fill(byteBuffer.array(), (byte) 0);
            fill(charBuffer.array(), '\u0000');
        }
    }

    private KeyStore decodeKeyStore(byte[] bytes,
            PasswordProtection protection)
            throws InvalidPasswordException, KeyManagementException {
        try (InputStream stream = new ByteArrayInputStream(bytes)) {
            return storesService.loadKeyStore(KeyStoreType.PKCS12, stream,
                protection.getPassword());
        } catch (GeneralCryptoLibException e) {
            if (e.getCause() instanceof IOException && e.getCause()
                .getCause() instanceof UnrecoverableKeyException) {
                throw new InvalidPasswordException(
                    "Key store password is invalid.", e);
            } else {
                throw new KeyManagementException(
                    "Failed to decode the key store.", e);
            }
        } catch (IOException e) {
            throw new KeyManagementException(
                "Failed to decode the key store.", e);
        }
    }
}
