/*
 * 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 org.junit.Assert.assertArrayEquals;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertTrue;
import static org.mockito.Matchers.any;
import static org.mockito.Matchers.anyString;
import static org.mockito.Matchers.eq;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

import java.security.KeyManagementException;
import java.security.KeyStore.PasswordProtection;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.Security;
import java.security.cert.X509Certificate;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.SQLIntegrityConstraintViolationException;
import java.sql.Statement;

import javax.sql.DataSource;

import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.junit.Before;
import org.junit.BeforeClass;
import org.junit.Test;

import com.scytl.cryptolib.api.elgamal.ElGamalServiceAPI;
import com.scytl.cryptolib.api.exceptions.GeneralCryptoLibException;
import com.scytl.cryptolib.elgamal.bean.ElGamalEncryptionParameters;
import com.scytl.cryptolib.elgamal.bean.ElGamalKeyPair;
import com.scytl.cryptolib.elgamal.bean.ElGamalPrivateKey;
import com.scytl.cryptolib.elgamal.bean.ElGamalPublicKey;
import com.scytl.cryptolib.elgamal.service.ElGamalService;

/**
 * Tests of {@link DatabaseImpl}.
 */
public class DatabaseImplTest {
    private static final String ENCRYPTION_PARAMETERS_JSON =
            "{\"encryptionParams\":{\"p\":\"AMoJqsCm6UG4B+FGO1Vh6qJcCJcMAVyZFIbZnTIZyD3XAQSTLdxVZxrXm5ZNQMajKAmjUfb3/BV84OgDZaNMXQziV2Uw6hQ70jUcswZ6ZfHAQBy+kwSozDM1LA6ndqp0JIGmPD7LqFIhv9ly7427ksgX1LjNR1tUInVaCS26dfP5wkZyCxwa9/3TapFpFugek+WSJqozaDSd89WAuqXrbBu0FMJWvBVvJapsjfJLEOXTHzdeWef9X5M4FY4Uk+/0WaDVZXy5JhcdoxuNU6Oy0/9buDrPfzlg5a26gpV5ei7XHfykYxdBIPUQoAmraSyUp/byhSDcQu6GJxkfknfAORU=\",\"q\":\"AJT/piLUjl6zN9D/Pc4GYsAG0KAS57H7xJ/mqb+ycbIl\",\"g\":\"ALNkOUZfTDbTMsCYUAKBAAzPdl7422rYGhca6GUAwEPQMlexuIZfIoft+ww54NDqGrUb4TFyw8Yy+y+9Ryb7AXNockro4XYfUjN6eiwOrPxfWG3r9BxqReWZRp74uQ2Ec1toVBrpHaXQNETOJpPBThMKKqrgIuPeAdMRM9U9viku8X0IDmHcqoNaI2o2fvFLFaxlzt5h4kzeJ0A5Twn8wwp5Uln0KcTBJf3H904Icq1u9WTvCTa1jQkoSzkNDQ1VABNFWBa/hW/RGAG7YqTurYHQB+9sMRpOiN10X+6ihdF5VgGpCgXfquKCSqa7I44vq2Vr9+Bm34Gjy+uDOJWrilg=\"}}";

    private static final String NODE_ID = "nodeId";

    private static final String ELECTION_EVENT_ID = "electionEventId";

    private static final String ELECTORAL_AUTHORITY_ID =
        "electoralAuthorityId";

    private static final String VERIFICATION_CARD_SET_ID =
        "verificationCardSetId";

    private static ElGamalPrivateKey elGamalPrivateKey;

    private static ElGamalPublicKey elGamalPublicKey;

    private ResultSet resultSet;

    private Statement statement;

    private PreparedStatement preparedStatement;

    private Connection connection;

    private DataSource dataSource;

    private Codec codec;

    private PasswordProtection password;

    private Generator generator;

    private PrivateKey encryptionPrivateKey;

    private PublicKey encryptionPublicKey;

    private DatabaseImpl database;

    @BeforeClass
    public static void beforeClass() throws GeneralCryptoLibException {
        if (Security
            .getProvider(BouncyCastleProvider.PROVIDER_NAME) == null) {
            Security.addProvider(new BouncyCastleProvider());
        }
        ElGamalServiceAPI elGamalService = new ElGamalService();
        ElGamalEncryptionParameters parameters = ElGamalEncryptionParameters
                .fromJson(ENCRYPTION_PARAMETERS_JSON);
        ElGamalKeyPair pair = elGamalService.getElGamalKeyPairGenerator()
            .generateKeys(parameters, 1);
        elGamalPrivateKey = pair.getPrivateKeys();
        elGamalPublicKey = pair.getPublicKeys();
    }

    @Before
    public void setUp() throws KeyManagementException, SQLException {
        resultSet = mock(ResultSet.class);
        when(resultSet.next()).thenReturn(true);
        when(resultSet.getInt(1)).thenReturn(0);

        statement = mock(Statement.class);

        preparedStatement = mock(PreparedStatement.class);
        when(preparedStatement.executeQuery()).thenReturn(resultSet);

        connection = mock(Connection.class);
        when(connection.createStatement()).thenReturn(statement);
        when(connection.prepareStatement(anyString()))
            .thenReturn(preparedStatement);

        dataSource = mock(DataSource.class);
        when(dataSource.getConnection()).thenReturn(connection);

        codec = mock(Codec.class);

        password = new PasswordProtection("password".toCharArray());

        generator = mock(Generator.class);
        when(generator.generatePassword()).thenReturn(password);

        encryptionPrivateKey = mock(PrivateKey.class);
        encryptionPublicKey = mock(PublicKey.class);

        database = new DatabaseImpl(dataSource, codec, generator, NODE_ID);
        database.setEncryptionKeys(encryptionPrivateKey,
            encryptionPublicKey);
    }

    @Test
    public void testLoadChoiceCodeKeys()
            throws KeyManagementException, SQLException {
        byte[] publicKeySignature = {1, 2, 3 };

        byte[] generationPrivateKeyBytes = {1 };
        byte[] generationPublicKeyBytes = {2 };
        byte[] decryptionPrivateKeyBytes = {3 };
        byte[] decryptionPublicKeyBytes = {4 };

        when(resultSet.getBytes("generation_private_key"))
            .thenReturn(generationPrivateKeyBytes);
        when(resultSet.getBytes("generation_public_key"))
            .thenReturn(generationPublicKeyBytes);
        when(resultSet.getBytes("generation_signature"))
            .thenReturn(publicKeySignature);

        when(resultSet.getBytes("decryption_private_key"))
            .thenReturn(decryptionPrivateKeyBytes);
        when(resultSet.getBytes("decryption_public_key"))
            .thenReturn(decryptionPublicKeyBytes);
        when(resultSet.getBytes("decryption_signature"))
            .thenReturn(publicKeySignature);

        when(codec.decodeElGamalPrivateKey(generationPrivateKeyBytes,
            encryptionPrivateKey)).thenReturn(elGamalPrivateKey);
        when(codec.decodeElGamalPublicKey(generationPublicKeyBytes))
            .thenReturn(elGamalPublicKey);
        when(codec.decodeElGamalPrivateKey(decryptionPrivateKeyBytes,
            encryptionPrivateKey)).thenReturn(elGamalPrivateKey);
        when(codec.decodeElGamalPublicKey(decryptionPublicKeyBytes))
            .thenReturn(elGamalPublicKey);

        ChoiceCodeKeys keys = database.loadChoiceCodeKeys(
            ELECTION_EVENT_ID, VERIFICATION_CARD_SET_ID);

        assertEquals(elGamalPrivateKey, keys.generationPrivateKey());
        assertEquals(elGamalPublicKey, keys.generationPublicKey());
        assertArrayEquals(publicKeySignature,
            keys.generationPublicKeySignature());
        assertEquals(elGamalPrivateKey, keys.decryptionPrivateKey());
        assertEquals(elGamalPublicKey, keys.decryptionPublicKey());
        assertArrayEquals(publicKeySignature,
            keys.decryptionPublicKeySignature());

        verify(connection).close();
        verify(preparedStatement).close();
        verify(resultSet).close();
        verify(preparedStatement).setString(1, NODE_ID);
        verify(preparedStatement).setString(2, ELECTION_EVENT_ID);
        verify(preparedStatement).setString(3, VERIFICATION_CARD_SET_ID);
    }

    @Test(expected = KeyManagementException.class)
    public void testLoadChoiceCodeKeysCodecException()
            throws KeyManagementException, SQLException {
        when(resultSet.getBytes(anyString())).thenReturn(new byte[0]);
        when(codec.decodeElGamalPrivateKey(any(byte[].class),
            eq(encryptionPrivateKey)))
                .thenThrow(new KeyManagementException("test"));
        database.loadChoiceCodeKeys(ELECTION_EVENT_ID,
            VERIFICATION_CARD_SET_ID);
    }

    @Test(expected = KeyNotFoundException.class)
    public void testLoadChoiceCodeKeysNotFound()
            throws SQLException, KeyNotFoundException,
            KeyManagementException {
        when(resultSet.next()).thenReturn(false);
        try {
            database.loadChoiceCodeKeys(ELECTION_EVENT_ID,
                VERIFICATION_CARD_SET_ID);
        } finally {
            verify(connection).close();
            verify(preparedStatement).close();
            verify(resultSet).close();
        }
    }

    @Test(expected = KeyManagementException.class)
    public void testLoadChoiceCodeKeysSQLException()
            throws SQLException, KeyNotFoundException,
            KeyManagementException {
        when(resultSet.next()).thenThrow(new SQLException("test"));
        try {
            database.loadChoiceCodeKeys(ELECTION_EVENT_ID,
                VERIFICATION_CARD_SET_ID);
        } finally {
            verify(connection).close();
            verify(preparedStatement).close();
            verify(resultSet).close();
        }
    }

    @Test
    public void testLoadElectionSigningKeys()
            throws SQLException, KeyManagementException {
        PrivateKey privateKey = mock(PrivateKey.class);
        X509Certificate[] certificateChain =
            {mock(X509Certificate.class) };

        ElectionSigningKeys keys =
            new ElectionSigningKeys(privateKey, certificateChain);

        byte[] keysBytes = {1, 2, 3 };
        byte[] passwordBytes = {4, 5, 6 };

        when(resultSet.getBytes("keys")).thenReturn(keysBytes);
        when(resultSet.getBytes("password")).thenReturn(passwordBytes);

        when(codec.decodePassword(passwordBytes, encryptionPrivateKey))
            .thenReturn(password);
        when(codec.decodeElectionSigningKeys(keysBytes, password))
            .thenReturn(keys);

        assertEquals(keys,
            database.loadElectionSigningKeys(ELECTION_EVENT_ID));

        verify(connection).close();
        verify(preparedStatement).close();
        verify(resultSet).close();
        verify(preparedStatement).setString(1, NODE_ID);
        verify(preparedStatement).setString(2, ELECTION_EVENT_ID);
        assertTrue(password.isDestroyed());
    }

    @Test(expected = KeyManagementException.class)
    public void testLoadElectionSigningKeysCodecException()
            throws SQLException, KeyManagementException {
        when(resultSet.getBytes("keys")).thenReturn(new byte[0]);
        when(resultSet.getBytes("password")).thenReturn(new byte[0]);

        when(codec.decodePassword(any(byte[].class),
            eq(encryptionPrivateKey)))
                .thenThrow(new KeyManagementException("test"));

        try {
            database.loadElectionSigningKeys(ELECTION_EVENT_ID);
        } finally {
            verify(connection).close();
            verify(preparedStatement).close();
            verify(resultSet).close();
        }
    }

    @Test(expected = KeyNotFoundException.class)
    public void testLoadElectionSigningKeysNotFound()
            throws KeyNotFoundException, KeyManagementException,
            SQLException {
        when(resultSet.next()).thenReturn(false);
        try {
            database.loadElectionSigningKeys(ELECTION_EVENT_ID);
        } finally {
            verify(connection).close();
            verify(preparedStatement).close();
            verify(resultSet).close();
        }
    }

    @Test(expected = KeyManagementException.class)
    public void testLoadElectionSigningKeysSQLException()
            throws KeyNotFoundException, KeyManagementException,
            SQLException {
        when(resultSet.next()).thenThrow(new SQLException("test"));
        try {
            database.loadElectionSigningKeys(ELECTION_EVENT_ID);
        } finally {
            verify(connection).close();
            verify(preparedStatement).close();
            verify(resultSet).close();
        }
    }

    @Test
    public void testLoadMixingKeys()
            throws SQLException, KeyManagementException {
        byte[] publicKeySignature = {1, 2, 3 };

        byte[] privateKeyBytes = {1 };
        byte[] publicKeyBytes = {2 };

        when(resultSet.getBytes("private_key"))
            .thenReturn(privateKeyBytes);
        when(resultSet.getBytes("public_key")).thenReturn(publicKeyBytes);
        when(resultSet.getBytes("signature"))
            .thenReturn(publicKeySignature);

        when(codec.decodeElGamalPrivateKey(privateKeyBytes,
            encryptionPrivateKey)).thenReturn(elGamalPrivateKey);
        when(codec.decodeElGamalPublicKey(publicKeyBytes))
            .thenReturn(elGamalPublicKey);

        MixingKeys keys = database.loadMixingKeys(ELECTION_EVENT_ID,
            ELECTORAL_AUTHORITY_ID);

        assertEquals(elGamalPrivateKey, keys.privateKey());
        assertEquals(elGamalPublicKey, keys.publicKey());
        assertArrayEquals(publicKeySignature, keys.publicKeySignature());

        verify(connection).close();
        verify(preparedStatement).close();
        verify(resultSet).close();
        verify(preparedStatement).setString(1, NODE_ID);
        verify(preparedStatement).setString(2, ELECTION_EVENT_ID);
        verify(preparedStatement).setString(3, ELECTORAL_AUTHORITY_ID);
    }

    @Test(expected = KeyManagementException.class)
    public void testLoadMixingKeysCodecException()
            throws SQLException, KeyManagementException {
        when(resultSet.getBytes("private_key")).thenReturn(new byte[0]);
        when(resultSet.getBytes("public_key")).thenReturn(new byte[0]);
        when(resultSet.getBytes("signature")).thenReturn(new byte[0]);

        when(codec.decodeElGamalPrivateKey(any(byte[].class),
            eq(encryptionPrivateKey)))
                .thenThrow(new KeyManagementException("test"));
        try {
            database.loadMixingKeys(ELECTION_EVENT_ID,
                ELECTORAL_AUTHORITY_ID);
        } finally {
            verify(connection).close();
            verify(preparedStatement).close();
            verify(resultSet).close();
        }
    }

    @Test(expected = KeyNotFoundException.class)
    public void testLoadMixingKeysNotFound()
            throws SQLException, KeyNotFoundException,
            KeyManagementException {
        when(resultSet.next()).thenReturn(false);
        try {
            database.loadMixingKeys(ELECTION_EVENT_ID,
                ELECTORAL_AUTHORITY_ID);
        } finally {
            verify(connection).close();
            verify(preparedStatement).close();
            verify(resultSet).close();
        }
    }

    @Test(expected = KeyManagementException.class)
    public void testLoadMixingKeysSQLException()
            throws SQLException, KeyNotFoundException,
            KeyManagementException {
        when(resultSet.next()).thenThrow(new SQLException("test"));
        try {
            database.loadMixingKeys(ELECTION_EVENT_ID,
                ELECTORAL_AUTHORITY_ID);
        } finally {
            verify(connection).close();
            verify(preparedStatement).close();
            verify(resultSet).close();
        }
    }

    @Test
    public void testLoadNodeKeys()
            throws SQLException, InvalidPasswordException,
            KeyManagementException {
        PrivateKey privateKey = mock(PrivateKey.class);
        X509Certificate[] certificateChain =
            {mock(X509Certificate.class) };
        NodeKeys keys = new NodeKeys.Builder()
            .setCAKeys(privateKey, certificateChain)
            .setEncryptionKeys(privateKey, certificateChain)
            .setLogSigningKeys(privateKey, certificateChain)
            .setLogEncryptionKeys(privateKey, certificateChain).build();

        byte[] keysBytes = {1, 2, 3 };
        when(resultSet.getBytes("keys")).thenReturn(keysBytes);
        when(codec.decodeNodeKeys(keysBytes, password)).thenReturn(keys);

        assertEquals(keys, database.loadNodeKeys(password));

        verify(connection).close();
        verify(preparedStatement).close();
        verify(resultSet).close();
        verify(preparedStatement).setString(1, NODE_ID);
    }

    @Test(expected = KeyManagementException.class)
    public void testLoadNodeKeysCodecException()
            throws SQLException, KeyNotFoundException,
            InvalidPasswordException, KeyManagementException {
        when(resultSet.getBytes("keys")).thenReturn(new byte[0]);
        when(codec.decodeNodeKeys(any(byte[].class), eq(password)))
            .thenThrow(new KeyManagementException("test"));
        try {
            database.loadNodeKeys(password);
        } finally {
            verify(connection).close();
            verify(preparedStatement).close();
            verify(resultSet).close();
        }
    }

    @Test(expected = KeyNotFoundException.class)
    public void testLoadNodeKeysNotFound()
            throws SQLException, KeyNotFoundException,
            InvalidPasswordException, KeyManagementException {
        when(resultSet.next()).thenReturn(false);
        try {
            database.loadNodeKeys(password);
        } finally {
            verify(connection).close();
            verify(preparedStatement).close();
            verify(resultSet).close();
        }
    }

    @Test(expected = KeyManagementException.class)
    public void testLoadNodeKeysSQLException()
            throws SQLException, KeyNotFoundException,
            InvalidPasswordException, KeyManagementException {
        when(resultSet.next()).thenThrow(new SQLException("test"));
        try {
            database.loadNodeKeys(password);
        } finally {
            verify(connection).close();
            verify(preparedStatement).close();
            verify(resultSet).close();
        }
    }

    @Test
    public void testSaveChoiceCodeKeys()
            throws GeneralCryptoLibException, KeyManagementException,
            SQLException {
        byte[] generationSignature = {1, 2, 3 };
        byte[] decryptionSignature = {4, 5, 6 };
        ChoiceCodeKeys keys = new ChoiceCodeKeys.Builder()
            .setGenerationKeys(elGamalPrivateKey, elGamalPublicKey,
                generationSignature)
            .setDecryptionKeys(elGamalPrivateKey, elGamalPublicKey,
                decryptionSignature)
            .build();
        byte[] generationPrivateKeyBytes = {1 };
        byte[] generationPublicKeyBytes = {2 };
        byte[] decryptionPrivateKeyBytes = {3 };
        byte[] decryptionPublicKeyBytes = {4 };

        when(codec.encodeElGamalPrivateKey(elGamalPrivateKey,
            encryptionPublicKey)).thenReturn(generationPrivateKeyBytes,
                decryptionPrivateKeyBytes);
        when(codec.encodeElGamalPublicKey(elGamalPublicKey)).thenReturn(
            generationPublicKeyBytes, decryptionPublicKeyBytes);
        database.saveChoiceCodeKeys(ELECTION_EVENT_ID,
            VERIFICATION_CARD_SET_ID, keys);

        verify(connection).close();
        verify(preparedStatement, times(2)).close();
        verify(resultSet).close();
        verify(preparedStatement, times(2)).setString(1, NODE_ID);
        verify(preparedStatement, times(2)).setString(2,
            ELECTION_EVENT_ID);
        verify(preparedStatement, times(2)).setString(3,
            VERIFICATION_CARD_SET_ID);
        verify(preparedStatement).setBytes(4, generationPrivateKeyBytes);
        verify(preparedStatement).setBytes(5, generationPublicKeyBytes);
        verify(preparedStatement).setBytes(6, generationSignature);
        verify(preparedStatement).setBytes(7, decryptionPrivateKeyBytes);
        verify(preparedStatement).setBytes(8, decryptionPublicKeyBytes);
        verify(preparedStatement).setBytes(9, decryptionSignature);
        verify(preparedStatement).executeUpdate();
    }

    @Test(expected = KeyAlreadyExistsException.class)
    public void testSaveChoiceCodeKeysAlreadyExist()
            throws KeyManagementException, SQLException {
        byte[] generationSignature = {1, 2, 3 };
        byte[] decryptionSignature = {4, 5, 6 };
        ChoiceCodeKeys keys = new ChoiceCodeKeys.Builder()
            .setGenerationKeys(elGamalPrivateKey, elGamalPublicKey,
                generationSignature)
            .setDecryptionKeys(elGamalPrivateKey, elGamalPublicKey,
                decryptionSignature)
            .build();
        byte[] generationPrivateKeyBytes = {1 };
        byte[] generationPublicKeyBytes = {2 };
        byte[] decryptionPrivateKeyBytes = {3 };
        byte[] decryptionPublicKeyBytes = {4 };

        when(codec.encodeElGamalPrivateKey(elGamalPrivateKey,
            encryptionPublicKey)).thenReturn(generationPrivateKeyBytes,
                decryptionPrivateKeyBytes);
        when(codec.encodeElGamalPublicKey(elGamalPublicKey)).thenReturn(
            generationPublicKeyBytes, decryptionPublicKeyBytes);

        when(resultSet.getInt(1)).thenReturn(1);

        database.saveChoiceCodeKeys(ELECTION_EVENT_ID,
            VERIFICATION_CARD_SET_ID, keys);
    }

    @Test(expected = KeyManagementException.class)
    public void testSaveChoiceCodeKeysCodecException()
            throws KeyManagementException {
        byte[] generationSignature = {1, 2, 3 };
        byte[] decryptionSignature = {4, 5, 6 };
        ChoiceCodeKeys keys = new ChoiceCodeKeys.Builder()
            .setGenerationKeys(elGamalPrivateKey, elGamalPublicKey,
                generationSignature)
            .setDecryptionKeys(elGamalPrivateKey, elGamalPublicKey,
                decryptionSignature)
            .build();

        when(codec.encodeElGamalPrivateKey(elGamalPrivateKey,
            encryptionPublicKey))
                .thenThrow(new KeyManagementException("test"));
        database.saveChoiceCodeKeys(ELECTION_EVENT_ID,
            VERIFICATION_CARD_SET_ID, keys);
    }

    @Test(expected = KeyAlreadyExistsException.class)
    public void testSaveChoiceCodeKeysDuplicates()
            throws KeyManagementException, SQLException {
        byte[] generationSignature = {1, 2, 3 };
        byte[] decryptionSignature = {4, 5, 6 };
        ChoiceCodeKeys keys = new ChoiceCodeKeys.Builder()
            .setGenerationKeys(elGamalPrivateKey, elGamalPublicKey,
                generationSignature)
            .setDecryptionKeys(elGamalPrivateKey, elGamalPublicKey,
                decryptionSignature)
            .build();
        byte[] generationPrivateKeyBytes = {1 };
        byte[] generationPublicKeyBytes = {2 };
        byte[] decryptionPrivateKeyBytes = {3 };
        byte[] decryptionPublicKeyBytes = {4 };

        when(codec.encodeElGamalPrivateKey(elGamalPrivateKey,
            encryptionPublicKey)).thenReturn(generationPrivateKeyBytes,
                decryptionPrivateKeyBytes);
        when(codec.encodeElGamalPublicKey(elGamalPublicKey)).thenReturn(
            generationPublicKeyBytes, decryptionPublicKeyBytes);

        when(preparedStatement.executeUpdate()).thenThrow(
            new SQLIntegrityConstraintViolationException("test"));

        try {
            database.saveChoiceCodeKeys(ELECTION_EVENT_ID,
                VERIFICATION_CARD_SET_ID, keys);
        } finally {
            verify(connection).close();
            verify(preparedStatement, times(2)).close();
            verify(resultSet).close();
        }
    }

    @Test(expected = KeyManagementException.class)
    public void testSaveChoiceCodeKeysSQLException()
            throws KeyManagementException, SQLException {
        byte[] generationSignature = {1, 2, 3 };
        byte[] decryptionSignature = {4, 5, 6 };
        ChoiceCodeKeys keys = new ChoiceCodeKeys.Builder()
            .setGenerationKeys(elGamalPrivateKey, elGamalPublicKey,
                generationSignature)
            .setDecryptionKeys(elGamalPrivateKey, elGamalPublicKey,
                decryptionSignature)
            .build();
        byte[] generationPrivateKeyBytes = {1 };
        byte[] generationPublicKeyBytes = {2 };
        byte[] decryptionPrivateKeyBytes = {3 };
        byte[] decryptionPublicKeyBytes = {4 };

        when(codec.encodeElGamalPrivateKey(elGamalPrivateKey,
            encryptionPublicKey)).thenReturn(generationPrivateKeyBytes,
                decryptionPrivateKeyBytes);
        when(codec.encodeElGamalPublicKey(elGamalPublicKey)).thenReturn(
            generationPublicKeyBytes, decryptionPublicKeyBytes);

        when(preparedStatement.executeUpdate())
            .thenThrow(new SQLException("test"));

        try {
            database.saveChoiceCodeKeys(ELECTION_EVENT_ID,
                VERIFICATION_CARD_SET_ID, keys);
        } finally {
            verify(connection).close();
            verify(preparedStatement, times(2)).close();
            verify(resultSet).close();
        }
    }

    @Test
    public void testSaveElectionSigningKeys()
            throws KeyManagementException, SQLException {
        PrivateKey privateKey = mock(PrivateKey.class);
        X509Certificate[] certificateChain =
            {mock(X509Certificate.class) };
        ElectionSigningKeys keys =
            new ElectionSigningKeys(privateKey, certificateChain);

        byte[] passwordBytes = {1 };
        byte[] keysBytes = {2 };
        when(codec.encodePassword(password, encryptionPublicKey))
            .thenReturn(passwordBytes);
        when(codec.encodeElectionSigningKeys(keys, password))
            .thenReturn(keysBytes);

        database.saveElectionSigningKeys(ELECTION_EVENT_ID, keys);

        verify(connection).close();
        verify(preparedStatement, times(2)).close();
        verify(resultSet).close();
        verify(preparedStatement, times(2)).setString(1, NODE_ID);
        verify(preparedStatement, times(2)).setString(2,
            ELECTION_EVENT_ID);
        verify(preparedStatement).setBytes(3, keysBytes);
        verify(preparedStatement).setBytes(4, passwordBytes);
        verify(preparedStatement).executeUpdate();
        assertTrue(password.isDestroyed());
    }

    @Test(expected = KeyAlreadyExistsException.class)
    public void testSaveElectionSigningKeysAlreadyExists()
            throws KeyManagementException, SQLException {
        PrivateKey privateKey = mock(PrivateKey.class);
        X509Certificate[] certificateChain =
            {mock(X509Certificate.class) };
        ElectionSigningKeys keys =
            new ElectionSigningKeys(privateKey, certificateChain);

        byte[] passwordBytes = {1 };
        byte[] keysBytes = {2 };
        when(codec.encodePassword(password, encryptionPublicKey))
            .thenReturn(passwordBytes);
        when(codec.encodeElectionSigningKeys(keys, password))
            .thenReturn(keysBytes);

        when(resultSet.getInt(1)).thenReturn(1);

        try {
            database.saveElectionSigningKeys(ELECTION_EVENT_ID, keys);
        } finally {
            verify(connection).close();
            verify(preparedStatement).close();
            verify(resultSet).close();
        }
    }

    @Test(expected = KeyManagementException.class)
    public void testSaveElectionSigningKeysCodecException()
            throws KeyManagementException, SQLException {
        PrivateKey privateKey = mock(PrivateKey.class);
        X509Certificate[] certificateChain =
            {mock(X509Certificate.class) };
        ElectionSigningKeys keys =
            new ElectionSigningKeys(privateKey, certificateChain);

        byte[] passwordBytes = {1 };
        when(codec.encodePassword(password, encryptionPublicKey))
            .thenReturn(passwordBytes);
        when(codec.encodeElectionSigningKeys(keys, password))
            .thenThrow(new KeyManagementException("test"));

        try {
            database.saveElectionSigningKeys(ELECTION_EVENT_ID, keys);
        } finally {
            verify(connection).close();
            verify(preparedStatement, times(2)).close();
            verify(resultSet).close();
            assertTrue(password.isDestroyed());
        }
    }

    @Test(expected = KeyAlreadyExistsException.class)
    public void testSaveElectionSigningKeysDuplicates()
            throws KeyManagementException, SQLException {
        PrivateKey privateKey = mock(PrivateKey.class);
        X509Certificate[] certificateChain =
            {mock(X509Certificate.class) };
        ElectionSigningKeys keys =
            new ElectionSigningKeys(privateKey, certificateChain);

        byte[] passwordBytes = {1 };
        byte[] keysBytes = {2 };
        when(codec.encodePassword(password, encryptionPublicKey))
            .thenReturn(passwordBytes);
        when(codec.encodeElectionSigningKeys(keys, password))
            .thenReturn(keysBytes);

        when(preparedStatement.executeUpdate()).thenThrow(
            new SQLIntegrityConstraintViolationException("test"));

        try {
            database.saveElectionSigningKeys(ELECTION_EVENT_ID, keys);
        } finally {
            verify(connection).close();
            verify(preparedStatement, times(2)).close();
            verify(resultSet).close();
            assertTrue(password.isDestroyed());
        }
    }

    @Test(expected = KeyManagementException.class)
    public void testSaveElectionSigningKeysGeneratorException()
            throws KeyManagementException {
        PrivateKey privateKey = mock(PrivateKey.class);
        X509Certificate[] certificateChain =
            {mock(X509Certificate.class) };
        ElectionSigningKeys keys =
            new ElectionSigningKeys(privateKey, certificateChain);

        when(generator.generatePassword())
            .thenThrow(new KeyManagementException("test"));

        database.saveElectionSigningKeys(ELECTION_EVENT_ID, keys);
    }

    @Test(expected = KeyManagementException.class)
    public void testSaveElectionSigningKeysSQLException()
            throws KeyManagementException, SQLException {
        PrivateKey privateKey = mock(PrivateKey.class);
        X509Certificate[] certificateChain =
            {mock(X509Certificate.class) };
        ElectionSigningKeys keys =
            new ElectionSigningKeys(privateKey, certificateChain);

        byte[] passwordBytes = {1 };
        byte[] keysBytes = {2 };
        when(codec.encodePassword(password, encryptionPublicKey))
            .thenReturn(passwordBytes);
        when(codec.encodeElectionSigningKeys(keys, password))
            .thenReturn(keysBytes);

        when(preparedStatement.executeUpdate())
            .thenThrow(new SQLException("test"));

        try {
            database.saveElectionSigningKeys(ELECTION_EVENT_ID, keys);
        } finally {
            verify(connection).close();
            verify(preparedStatement, times(2)).close();
            verify(resultSet).close();
            assertTrue(password.isDestroyed());
        }
    }

    @Test
    public void testSaveMixingKeys()
            throws KeyManagementException, SQLException {
        byte[] publicKeySignature = {1, 2, 3 };
        MixingKeys keys = new MixingKeys(elGamalPrivateKey,
            elGamalPublicKey, publicKeySignature);
        byte[] privateKeyBytes = {1 };
        byte[] publicKeyBytes = {2 };

        when(codec.encodeElGamalPrivateKey(elGamalPrivateKey,
            encryptionPublicKey)).thenReturn(privateKeyBytes,
                privateKeyBytes);
        when(codec.encodeElGamalPublicKey(elGamalPublicKey))
            .thenReturn(publicKeyBytes, publicKeyBytes);

        database.saveMixingKeys(ELECTION_EVENT_ID,
            VERIFICATION_CARD_SET_ID, keys);

        verify(connection).close();
        verify(preparedStatement, times(2)).close();
        verify(resultSet).close();
        verify(preparedStatement, times(2)).setString(1, NODE_ID);
        verify(preparedStatement, times(2)).setString(2,
            ELECTION_EVENT_ID);
        verify(preparedStatement, times(2)).setString(3,
            VERIFICATION_CARD_SET_ID);
        verify(preparedStatement).setBytes(4, privateKeyBytes);
        verify(preparedStatement).setBytes(5, publicKeyBytes);
        verify(preparedStatement).setBytes(6, publicKeySignature);
        verify(preparedStatement).executeUpdate();
    }

    @Test(expected = KeyAlreadyExistsException.class)
    public void testSaveMixingKeysAlreadyExist()
            throws KeyManagementException, SQLException {
        byte[] publicKeySignature = {1, 2, 3 };
        MixingKeys keys = new MixingKeys(elGamalPrivateKey,
            elGamalPublicKey, publicKeySignature);
        byte[] privateKeyBytes = {1 };
        byte[] publicKeyBytes = {2 };

        when(codec.encodeElGamalPrivateKey(elGamalPrivateKey,
            encryptionPublicKey)).thenReturn(privateKeyBytes,
                privateKeyBytes);
        when(codec.encodeElGamalPublicKey(elGamalPublicKey))
            .thenReturn(publicKeyBytes, publicKeyBytes);

        when(resultSet.getInt(1)).thenReturn(1);

        try {
            database.saveMixingKeys(ELECTION_EVENT_ID,
                VERIFICATION_CARD_SET_ID, keys);
        } finally {
            verify(connection).close();
            verify(preparedStatement).close();
            verify(resultSet).close();
        }
    }

    @Test(expected = KeyManagementException.class)
    public void testSaveMixingKeysCodecException()
            throws KeyManagementException, SQLException {
        byte[] publicKeySignature = {1, 2, 3 };
        MixingKeys keys = new MixingKeys(elGamalPrivateKey,
            elGamalPublicKey, publicKeySignature);

        when(codec.encodeElGamalPrivateKey(elGamalPrivateKey,
            encryptionPublicKey))
                .thenThrow(new KeyManagementException("test"));

        try {
            database.saveMixingKeys(ELECTION_EVENT_ID,
                VERIFICATION_CARD_SET_ID, keys);
        } finally {
            verify(connection).close();
            verify(preparedStatement, times(2)).close();
            verify(resultSet).close();
        }
    }

    @Test(expected = KeyAlreadyExistsException.class)
    public void testSaveMixingKeysDuplicates()
            throws KeyManagementException, SQLException {
        byte[] publicKeySignature = {1, 2, 3 };
        MixingKeys keys = new MixingKeys(elGamalPrivateKey,
            elGamalPublicKey, publicKeySignature);
        byte[] privateKeyBytes = {1 };
        byte[] publicKeyBytes = {2 };

        when(codec.encodeElGamalPrivateKey(elGamalPrivateKey,
            encryptionPublicKey)).thenReturn(privateKeyBytes,
                privateKeyBytes);
        when(codec.encodeElGamalPublicKey(elGamalPublicKey))
            .thenReturn(publicKeyBytes, publicKeyBytes);

        when(preparedStatement.executeUpdate()).thenThrow(
            new SQLIntegrityConstraintViolationException("test"));

        try {
            database.saveMixingKeys(ELECTION_EVENT_ID,
                VERIFICATION_CARD_SET_ID, keys);
        } finally {
            verify(connection).close();
            verify(preparedStatement, times(2)).close();
            verify(resultSet).close();
        }
    }

    @Test(expected = KeyManagementException.class)
    public void testSaveMixingKeysSQLException()
            throws KeyManagementException, SQLException {
        byte[] publicKeySignature = {1, 2, 3 };
        MixingKeys keys = new MixingKeys(elGamalPrivateKey,
            elGamalPublicKey, publicKeySignature);
        byte[] privateKeyBytes = {1 };
        byte[] publicKeyBytes = {2 };

        when(codec.encodeElGamalPrivateKey(elGamalPrivateKey,
            encryptionPublicKey)).thenReturn(privateKeyBytes,
                privateKeyBytes);
        when(codec.encodeElGamalPublicKey(elGamalPublicKey))
            .thenReturn(publicKeyBytes, publicKeyBytes);

        when(preparedStatement.executeUpdate())
            .thenThrow(new SQLException("test"));

        try {
            database.saveMixingKeys(ELECTION_EVENT_ID,
                VERIFICATION_CARD_SET_ID, keys);
        } finally {
            verify(connection).close();
            verify(preparedStatement, times(2)).close();
            verify(resultSet).close();
        }
    }

    @Test
    public void testSaveNodeKeys()
            throws KeyManagementException, SQLException {
        PrivateKey privateKey = mock(PrivateKey.class);
        X509Certificate[] certificateChain =
            {mock(X509Certificate.class) };
        NodeKeys keys = new NodeKeys.Builder()
            .setCAKeys(privateKey, certificateChain)
            .setEncryptionKeys(privateKey, certificateChain)
            .setLogSigningKeys(privateKey, certificateChain)
            .setLogEncryptionKeys(privateKey, certificateChain).build();

        byte[] keysBytes = {1, 2, 3 };

        when(codec.encodeNodeKeys(keys, password)).thenReturn(keysBytes);

        database.saveNodeKeys(keys, password);

        verify(connection).close();
        verify(statement).close();
        verify(preparedStatement, times(2)).close();
        verify(resultSet).close();
        verify(statement)
            .executeUpdate("lock table cc_node_keys in exclusive mode");
        verify(preparedStatement, times(2)).setString(1, NODE_ID);
        verify(preparedStatement).setBytes(2, keysBytes);
        verify(preparedStatement).executeUpdate();
    }

    @Test(expected = KeyAlreadyExistsException.class)
    public void testSaveNodeKeysAlreadyExist()
            throws KeyManagementException, SQLException {
        PrivateKey privateKey = mock(PrivateKey.class);
        X509Certificate[] certificateChain =
            {mock(X509Certificate.class) };
        NodeKeys keys = new NodeKeys.Builder()
            .setCAKeys(privateKey, certificateChain)
            .setEncryptionKeys(privateKey, certificateChain)
            .setLogSigningKeys(privateKey, certificateChain)
            .setLogEncryptionKeys(privateKey, certificateChain).build();

        byte[] keysBytes = {1, 2, 3 };

        when(codec.encodeNodeKeys(keys, password)).thenReturn(keysBytes);

        when(resultSet.getInt(1)).thenReturn(1);

        try {
            database.saveNodeKeys(keys, password);
        } finally {
            verify(connection).close();
            verify(statement).close();
            verify(preparedStatement).close();
            verify(resultSet).close();
            verify(statement).executeUpdate(
                "lock table cc_node_keys in exclusive mode");
        }
    }

    @Test(expected = KeyManagementException.class)
    public void testSaveNodeKeysCodecException()
            throws KeyManagementException, SQLException {
        PrivateKey privateKey = mock(PrivateKey.class);
        X509Certificate[] certificateChain =
            {mock(X509Certificate.class) };
        NodeKeys keys = new NodeKeys.Builder()
            .setCAKeys(privateKey, certificateChain)
            .setEncryptionKeys(privateKey, certificateChain)
            .setLogSigningKeys(privateKey, certificateChain)
            .setLogEncryptionKeys(privateKey, certificateChain).build();

        when(codec.encodeNodeKeys(keys, password))
            .thenThrow(new KeyManagementException("test"));

        try {
            database.saveNodeKeys(keys, password);
        } finally {
            verify(connection).close();
            verify(statement).close();
            verify(preparedStatement, times(2)).close();
            verify(resultSet).close();
            verify(statement).executeUpdate(
                "lock table cc_node_keys in exclusive mode");
        }
    }

    @Test(expected = KeyAlreadyExistsException.class)
    public void testSaveNodeKeysDuplicates()
            throws KeyManagementException, SQLException {
        PrivateKey privateKey = mock(PrivateKey.class);
        X509Certificate[] certificateChain =
            {mock(X509Certificate.class) };
        NodeKeys keys = new NodeKeys.Builder()
            .setCAKeys(privateKey, certificateChain)
            .setEncryptionKeys(privateKey, certificateChain)
            .setLogSigningKeys(privateKey, certificateChain)
            .setLogEncryptionKeys(privateKey, certificateChain).build();

        byte[] keysBytes = {1, 2, 3 };

        when(codec.encodeNodeKeys(keys, password)).thenReturn(keysBytes);

        when(preparedStatement.executeUpdate()).thenThrow(
            new SQLIntegrityConstraintViolationException("test"));

        try {
            database.saveNodeKeys(keys, password);
        } finally {
            verify(connection).close();
            verify(statement).close();
            verify(preparedStatement, times(2)).close();
            verify(resultSet).close();
            verify(statement).executeUpdate(
                "lock table cc_node_keys in exclusive mode");
        }
    }

    @Test(expected = KeyManagementException.class)
    public void testSaveNodeKeysSQLException()
            throws KeyManagementException, SQLException {
        PrivateKey privateKey = mock(PrivateKey.class);
        X509Certificate[] certificateChain =
            {mock(X509Certificate.class) };
        NodeKeys keys = new NodeKeys.Builder()
            .setCAKeys(privateKey, certificateChain)
            .setEncryptionKeys(privateKey, certificateChain)
            .setLogSigningKeys(privateKey, certificateChain)
            .setLogEncryptionKeys(privateKey, certificateChain).build();

        byte[] keysBytes = {1, 2, 3 };

        when(codec.encodeNodeKeys(keys, password)).thenReturn(keysBytes);

        when(preparedStatement.executeUpdate())
            .thenThrow(new SQLException("test"));

        try {
            database.saveNodeKeys(keys, password);
        } finally {
            verify(connection).close();
            verify(statement).close();
            verify(preparedStatement, times(2)).close();
            verify(resultSet).close();
            verify(statement).executeUpdate(
                "lock table cc_node_keys in exclusive mode");
        }
    }
}
