package xyz.whcself.whcshell;

import org.bouncycastle.crypto.AsymmetricCipherKeyPair;
import org.bouncycastle.crypto.generators.Ed25519KeyPairGenerator;
import org.bouncycastle.crypto.params.Ed25519KeyGenerationParameters;
import org.bouncycastle.crypto.util.OpenSSHPrivateKeyUtil;
import org.bouncycastle.crypto.util.OpenSSHPublicKeyUtil;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.security.SecureRandom;
import java.util.Base64;

/** The client's Ed25519 identity, stored in OpenSSH format in a private directory. */
final class SshKeys {

    final File privateKey;
    final File publicKey;

    private SshKeys(File dir) {
        privateKey = new File(dir, "id_ed25519");
        publicKey = new File(dir, "id_ed25519.pub");
    }

    /** Loads the key pair from {@code dir}, generating it first if missing. */
    static SshKeys ensure(File dir, String comment) throws IOException {
        SshKeys k = new SshKeys(dir);
        if (k.privateKey.isFile() && k.publicKey.isFile()) return k;
        if (!dir.isDirectory() && !dir.mkdirs()) throw new IOException("cannot create " + dir);

        Ed25519KeyPairGenerator gen = new Ed25519KeyPairGenerator();
        gen.init(new Ed25519KeyGenerationParameters(new SecureRandom()));
        AsymmetricCipherKeyPair pair = gen.generateKeyPair();

        Base64.Encoder b64 = Base64.getEncoder();
        String priv = "-----BEGIN OPENSSH PRIVATE KEY-----\n"
            + Base64.getMimeEncoder(70, new byte[]{'\n'}).encodeToString(
                OpenSSHPrivateKeyUtil.encodePrivateKey(pair.getPrivate()))
            + "\n-----END OPENSSH PRIVATE KEY-----\n";
        String pub = "ssh-ed25519 "
            + b64.encodeToString(OpenSSHPublicKeyUtil.encodePublicKey(pair.getPublic()))
            + " " + comment.replaceAll("[^A-Za-z0-9._@-]", "_") + "\n";

        write(k.privateKey, priv);
        write(k.publicKey, pub);
        return k;
    }

    /** The single authorized_keys line, without a trailing newline. */
    String publicLine() throws IOException {
        return new String(Files.readAllBytes(publicKey.toPath()), StandardCharsets.UTF_8).trim();
    }

    private static void write(File f, String s) throws IOException {
        File tmp = new File(f.getPath() + ".tmp");
        try (FileOutputStream out = new FileOutputStream(tmp)) {
            out.write(s.getBytes(StandardCharsets.UTF_8));
            out.getFD().sync();
        }
        tmp.setReadable(false, false);
        tmp.setReadable(true, true);
        tmp.setWritable(false, false);
        tmp.setWritable(true, true);
        if (!tmp.renameTo(f)) throw new IOException("cannot rename " + tmp);
    }
}