// Read a DER-encoded PKCS#8 RSA private key, derive the public key,
// emit the Android adb_keys line on stdout.
//
//   java test-rig/Keygen.java <path/to/adbkey.der>
//
// The Android pubkey blob format (from system/core/libcrypto_utils):
//   uint32_t modulus_size_words;   // = 64 for RSA-2048
//   uint32_t n0inv;                // -1 / N[0] mod 2^32
//   uint8_t  modulus[256];         // little-endian
//   uint8_t  rr[256];              // (2^2048)^2 mod N, little-endian
//   uint32_t exponent;             // typically 65537
// = 524 bytes, then base64-encoded, then " scrcpy-android@test".

import java.io.IOException;
import java.math.BigInteger;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.nio.file.Files;
import java.nio.file.Path;
import java.security.KeyFactory;
import java.security.PrivateKey;
import java.security.interfaces.RSAPrivateCrtKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.RSAPublicKeySpec;
import java.util.Base64;

public final class Keygen {

    private static final int MODULUS_SIZE       = 256;             // bytes (2048 / 8)
    private static final int MODULUS_SIZE_WORDS = MODULUS_SIZE / 4;
    private static final int ENCODED_SIZE       = 3 * 4 + 2 * MODULUS_SIZE;

    public static void main(String[] args) throws Exception {
        if (args.length != 1) {
            System.err.println("usage: java Keygen.java <der private key path>");
            System.exit(2);
        }
        byte[] der = Files.readAllBytes(Path.of(args[0]));
        PrivateKey priv = KeyFactory.getInstance("RSA")
                .generatePrivate(new PKCS8EncodedKeySpec(der));
        RSAPublicKey pub = derivePublic(priv);

        byte[] blob = encode(pub);
        String b64 = Base64.getEncoder().encodeToString(blob);
        System.out.print(b64);
        System.out.println(" scrcpy-android@test");
    }

    private static RSAPublicKey derivePublic(PrivateKey priv) throws Exception {
        if (priv instanceof RSAPrivateCrtKey crt) {
            RSAPublicKeySpec spec = new RSAPublicKeySpec(crt.getModulus(), crt.getPublicExponent());
            return (RSAPublicKey) KeyFactory.getInstance("RSA").generatePublic(spec);
        }
        throw new IOException("private key is not an RSAPrivateCrtKey");
    }

    private static byte[] encode(RSAPublicKey pub) {
        ByteBuffer buf = ByteBuffer.allocate(ENCODED_SIZE).order(ByteOrder.LITTLE_ENDIAN);
        BigInteger n = pub.getModulus();

        buf.putInt(MODULUS_SIZE_WORDS);

        // n0inv = (2^32) - (N[0] mod 2^32)^-1 mod 2^32
        BigInteger r32   = BigInteger.ZERO.setBit(32);
        BigInteger n0inv = n.mod(r32).modInverse(r32);
        n0inv = r32.subtract(n0inv);
        buf.putInt(n0inv.intValue());

        buf.put(beToLe(MODULUS_SIZE, n));

        // rr = (2^(MODULUS_SIZE*8))^2 mod N
        BigInteger rr = BigInteger.ZERO.setBit(MODULUS_SIZE * 8).modPow(BigInteger.TWO, n);
        buf.put(beToLe(MODULUS_SIZE, rr));

        buf.putInt(pub.getPublicExponent().intValue());
        return buf.array();
    }

    private static byte[] beToLe(int len, BigInteger v) {
        byte[] be = v.toByteArray();
        // strip leading sign byte if present
        int start = 0;
        if (be.length > len) start = be.length - len;
        byte[] out = new byte[len];
        for (int i = 0; i < len; i++) {
            int srcIdx = be.length - 1 - i;
            if (srcIdx >= start) out[i] = be[srcIdx];
        }
        return out;
    }
}
