import java.security.*;
import javax.crypto.EncryptedPrivateKeyInfo;
import javax.crypto.spec.PBEParameterSpec;

/** Encrypted private keys: write one, read one, read one openssl wrote, and every way it can go wrong. */
public class EncryptedKeys {
    static final char[] PW = "changeit".toCharArray();

    public static void main(String[] args) throws Exception {
        PEMDecoder d = PEMDecoder.of();
        PEMEncoder e = PEMEncoder.of();
        PrivateKey key = d.decode(Fx.read("rsa-private-pkcs8.pem"), PrivateKey.class);

        Fx.h("write an encrypted key with the defaults");
        String enc = e.withEncryption(PW).encodeToString(key);
        System.out.println(enc.lines().filter(l -> l.startsWith("-----")).toList());
        EncryptedPrivateKeyInfo info = (EncryptedPrivateKeyInfo) d.decode(enc);
        System.out.println("default algorithm: " + info.getAlgName());
        PBEParameterSpec spec = info.getAlgParameters().getParameterSpec(PBEParameterSpec.class);
        System.out.println("default iteration count: " + spec.getIterationCount() + ", salt bytes: " + spec.getSalt().length);

        Fx.h("read it back");
        Fx.check(d.withDecryption(PW).decode(enc, PrivateKey.class).equals(key), "withDecryption(password) returns the original key");

        Fx.h("choose the algorithm and cost yourself");
        PBEParameterSpec custom = new PBEParameterSpec(new SecureRandom().generateSeed(16), 600_000);
        EncryptedPrivateKeyInfo strong = EncryptedPrivateKeyInfo.encrypt(key, PW, "PBEWithHmacSHA512AndAES_256", custom, null);
        String strongPem = e.encodeToString(strong);
        EncryptedPrivateKeyInfo again = (EncryptedPrivateKeyInfo) d.decode(strongPem);
        System.out.println(again.getAlgName() + ", iterations " + again.getAlgParameters().getParameterSpec(PBEParameterSpec.class).getIterationCount());
        Fx.check(again.getKey(PW).equals(key), "getKey(password) on the decoded EncryptedPrivateKeyInfo works too");

        Fx.h("a file written by openssl pkcs8 -topk8 -v2 aes-256-cbc");
        EncryptedPrivateKeyInfo fromOpenssl = (EncryptedPrivateKeyInfo) d.decode(Fx.read("rsa-private-encrypted.pem"));
        System.out.println(fromOpenssl.getAlgName());
        Fx.check(d.withDecryption(PW).decode(Fx.read("rsa-private-encrypted.pem"), PrivateKey.class).equals(key), "decrypts to the same key as the unencrypted file");

        Fx.h("failures");
        String f = Fx.read("rsa-private-encrypted.pem");
        Fx.tryIt("no password, ask for PrivateKey", () -> d.decode(f, PrivateKey.class));
        Fx.tryIt("no password, no class", () -> d.decode(f));
        Fx.tryIt("wrong password", () -> d.withDecryption("nope".toCharArray()).decode(f, PrivateKey.class));
        Fx.tryIt("password on a key that is not encrypted", () -> d.withDecryption(PW).decode(Fx.read("rsa-private-pkcs8.pem"), PrivateKey.class));
        Fx.tryIt("openssl 'traditional' encrypted (Proc-Type: 4,ENCRYPTED)", () -> d.withDecryption(PW).decode(Fx.read("rsa-private-legacy-encrypted.pem"), PrivateKey.class));
    }
}