import java.io.OutputStream;
import java.nio.file.Files;
import java.nio.file.Path;
import java.security.DomainLoadStoreParameter;
import java.security.KeyStore;
import java.util.Enumeration;
import java.util.LinkedHashMap;
import java.util.Map;
import javax.crypto.KeyGenerator;

public class DksProblem {
    private static void write(Path p, KeyStore keystore) throws Exception {
        try (OutputStream outputStream = Files.newOutputStream(p)) {
            keystore.store(outputStream, new char[] { 'x' });
        }
    }

    public static void main(String...args) throws Exception {
        final KeyGenerator keyGenerator = KeyGenerator.getInstance("AES");
        keyGenerator.init(256);

        // Create keystore with one key:
        //
        final KeyStore nonEmptyKeystore = KeyStore.getInstance("PKCS12");
        nonEmptyKeystore.load(null, null);

        final Path nonEmptyPath = Path.of("non_empty.p12");
        nonEmptyKeystore.setKeyEntry("a", keyGenerator.generateKey(), new char[] { 'a' }, null);
        write(nonEmptyPath, nonEmptyKeystore);

        // Create empty keystore:
        //
        final KeyStore emptyKeystore = KeyStore.getInstance("PKCS12");
        emptyKeystore.load(null, null);

        final Path emptyPath = Path.of("empty.p12");
        write(emptyPath, emptyKeystore);

        // Create domain keystore with two non-empty parts:
        //
        final Path dksWithTwoPartsPath = Path.of("two-parts.dks");

        final var twoPartsConfiguration = "domain Combo { keystore a keystoreURI=\"%s\"; keystore b keystoreURI=\"%s\"; };";

        Files.writeString(dksWithTwoPartsPath, String.format(twoPartsConfiguration, nonEmptyPath.toUri(), nonEmptyPath.toUri()));


        final Map<String,KeyStore.ProtectionParameter> protectionParameters = new LinkedHashMap<>();

        final KeyStore dksKeystore = KeyStore.getInstance("DKS");
        dksKeystore.load(new DomainLoadStoreParameter(dksWithTwoPartsPath.toUri(), protectionParameters));

        System.out.println(String.format("%s size: %d", dksWithTwoPartsPath, dksKeystore.size()));
        int index = 0;

        for (final Enumeration<String> enumeration = dksKeystore.aliases(); enumeration.hasMoreElements(); ) {
            System.out.println(String.format("%d: %s", index, enumeration.nextElement()));
            index++;
        }

        System.out.println(String.format("enumerated aliases from %s: %d", dksWithTwoPartsPath, index));


        // Create domain keystore with two non-empty parts with an empty one in between:
        //
        final Path dksWithThreePartsPath = Path.of("three-parts.dks");

        final var fmt = "domain Combo { keystore a keystoreURI=\"%s\"; keystore b keystoreURI=\"%s\"; keystore c keystoreURI=\"%s\"; };";

        Files.writeString(dksWithThreePartsPath, String.format(fmt, nonEmptyPath.toUri(), emptyPath.toUri(), nonEmptyPath.toUri()));
        dksKeystore.load(new DomainLoadStoreParameter(dksWithThreePartsPath.toUri(), protectionParameters));

        System.out.println(String.format("%s size: %d", dksWithThreePartsPath, dksKeystore.size()));
        index = 0;

        for (final Enumeration<String> enumeration = dksKeystore.aliases(); enumeration.hasMoreElements(); ) {
            System.out.println(String.format("%d: %s", index, enumeration.nextElement()));
            index++;
        }

        System.out.println(String.format("enumerated aliases from %s: %d", dksWithThreePartsPath, index));
    }
} 