import java.io.ByteArrayOutputStream;
import java.io.DataOutputStream;
import java.io.StringReader;
import java.security.KeyFactory;
import java.security.PublicKey;
import java.security.interfaces.DSAParams;
import java.security.interfaces.DSAPublicKey;
import java.security.interfaces.ECPublicKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.X509EncodedKeySpec;
import org.apache.commons.codec.binary.Base64;
import org.bouncycastle.openssl.PEMKeyPair;
import org.bouncycastle.openssl.PEMParser;

public class GenPubSSHKey {

    /**
     * Generates a public key from a SSH private key.
     * @param privateKey A RSA/DSA/EC private key
     * @return SSH encoded public key
     */
    public static String getPublicKey(String privateKey) throws Exception {
        PEMParser pemParser = new PEMParser(new StringReader(privateKey));
        PEMKeyPair kp = (PEMKeyPair) pemParser.readObject();
        X509EncodedKeySpec x509 = new X509EncodedKeySpec(kp.getPublicKeyInfo().getEncoded());
        KeyFactory keyFactory;
        if (privateKey.startsWith("-----BEGIN RSA PRIVATE KEY-----")) {
            keyFactory = KeyFactory.getInstance("RSA");
        } else if (privateKey.startsWith("-----BEGIN DSA PRIVATE KEY-----")) {
            keyFactory = KeyFactory.getInstance("DSA");
        } else if (privateKey.startsWith("-----BEGIN EC PRIVATE KEY-----")) {
            keyFactory = KeyFactory.getInstance("EC");
        } else {
            throw new Exception("Incompatible SSH key algorithm");
        }
        return encodePublicKey(keyFactory.generatePublic(x509), "ozone");
    }

    private static String encodePublicKey(PublicKey publicKey, String user) throws Exception {
		String publicKeyEncoded;
		if ("RSA".equals(publicKey.getAlgorithm())) {
			RSAPublicKey rsaPublicKey = (RSAPublicKey) publicKey;
			StdByteArrayOutputStream byteOs = new StdByteArrayOutputStream();
			DataOutputStream dos = new DataOutputStream(byteOs);
			dos.writeInt("ssh-rsa".getBytes().length);
			dos.write("ssh-rsa".getBytes());
			dos.writeInt(rsaPublicKey.getPublicExponent().toByteArray().length);
			dos.write(rsaPublicKey.getPublicExponent().toByteArray());
			dos.writeInt(rsaPublicKey.getModulus().toByteArray().length);
			dos.write(rsaPublicKey.getModulus().toByteArray());
			publicKeyEncoded = Base64.getEncoder().encodeToString(byteOs.toByteArray());
			return String.format("ssh-rsa %s %s", publicKeyEncoded, user);
		} else if ("DSA".equals(publicKey.getAlgorithm())) {
			DSAPublicKey dsaPublicKey = (DSAPublicKey) publicKey;
			DSAParams dsaParams = dsaPublicKey.getParams();
			StdByteArrayOutputStream byteOs = new StdByteArrayOutputStream();
			DataOutputStream dos = new DataOutputStream(byteOs);
			dos.writeInt("ssh-dss".getBytes().length);
			dos.write("ssh-dss".getBytes());
			dos.writeInt(dsaParams.getP().toByteArray().length);
			dos.write(dsaParams.getP().toByteArray());
			dos.writeInt(dsaParams.getQ().toByteArray().length);
			dos.write(dsaParams.getQ().toByteArray());
			dos.writeInt(dsaParams.getG().toByteArray().length);
			dos.write(dsaParams.getG().toByteArray());
			dos.writeInt(dsaPublicKey.getY().toByteArray().length);
			dos.write(dsaPublicKey.getY().toByteArray());
			publicKeyEncoded = new String(Base64.getDecoder().decode(byteOs.toByteArray()));
			return String.format("ssh-dss %s %s", publicKeyEncoded, user);
		} else if ("ECDSA".equals(publicKey.getAlgorithm())) {
			ECPublicKey ecPublicKey = (ECPublicKey) publicKey;
			ECParameterSpec ecParams = ecPublicKey.getParams();
			StdByteArrayOutputStream byteOs = new StdByteArrayOutputStream();
			DataOutputStream dos = new DataOutputStream(byteOs);
			int fieldSize = ecParams.getCurve().getField().getFieldSize();
			String curveName = switch (fieldSize) {
				case 256 -> "nistp256";
				case 384 -> "nistp384";
				case 521 -> "nistp521";
				default -> throw new IllegalArgumentException("Unknown curve size: " + fieldSize);
			};
			String fullName = "ecdsa-sha2-" + curveName;
			dos.writeInt(fullName.getBytes().length);
			dos.write(fullName.getBytes());
			dos.writeInt(curveName.getBytes().length);
			dos.write(curveName.getBytes());

			ECPoint group = ecPublicKey.getW();
			int elementSize = (fieldSize + 7) / 8;
			byte[] M = new byte[2 * elementSize + 1];
			M[0] = 0x04;

			byte[] affineX = group.getAffineX().toByteArray();
			affineX = dropLeadingZeroBytes(affineX);
			System.arraycopy(affineX, 0, M, 1 + elementSize - affineX.length, affineX.length);

			byte[] affineY = group.getAffineY().toByteArray();
			affineY = dropLeadingZeroBytes(affineY);
			System.arraycopy(affineY, 0, M, 1 + elementSize + elementSize - affineY.length, affineY.length);

			dos.writeInt(M.length);
			dos.write(M);

			publicKeyEncoded = new String(Base64.getEncoder().encode(byteOs.toByteArray()));
			return String.format("%s %s %s", fullName, publicKeyEncoded, user);
		} else {
			throw new IllegalArgumentException("Unknown public key encoding: " + publicKey.getAlgorithm());
		}
	}

    private static byte[] dropLeadingZeroBytes(byte[] input) {
        int firstNonZero = 0;
        while (firstNonZero < input.length && input[firstNonZero] == 0x00) {
            firstNonZero++;
        }
        byte[] result = new byte[input.length - firstNonZero];
        System.arraycopy(input, firstNonZero, result, 0, result.length);
        return result;
    }
}