import java.nio.charset.StandardCharsets;
import java.security.*;
import java.util.Base64;

/** Sign and verify a JWT with keys loaded from PEM files. No JWT library: the point is the key handling. */
public class JwtSigning {
    static final Base64.Encoder B64 = Base64.getUrlEncoder().withoutPadding();

    static String signingInput(String alg) {
        String header = "{\"alg\":\"" + alg + "\",\"typ\":\"JWT\"}";
        String payload = "{\"sub\":\"ankur\",\"iss\":\"https://ankurm.com\"}";
        return B64.encodeToString(header.getBytes(StandardCharsets.UTF_8)) + "." + B64.encodeToString(payload.getBytes(StandardCharsets.UTF_8));
    }

    static byte[] sign(String jcaAlg, PrivateKey key, String input) throws Exception {
        Signature s = Signature.getInstance(jcaAlg);
        s.initSign(key);
        s.update(input.getBytes(StandardCharsets.US_ASCII));
        return s.sign();
    }

    static boolean verify(String jcaAlg, PublicKey key, String input, byte[] sig) throws Exception {
        Signature s = Signature.getInstance(jcaAlg);
        s.initVerify(key);
        s.update(input.getBytes(StandardCharsets.US_ASCII));
        return s.verify(sig);
    }

    public static void main(String[] args) throws Exception {
        PEMDecoder d = PEMDecoder.of();

        Fx.h("RS256 with an RSA key from PEM");
        PrivateKey rsaPriv = d.decode(Fx.read("rsa-private-pkcs8.pem"), PrivateKey.class);
        PublicKey rsaPub = d.decode(Fx.read("rsa-public.pem"), PublicKey.class);
        String in = signingInput("RS256");
        byte[] sig = sign("SHA256withRSA", rsaPriv, in);
        String jwt = in + "." + B64.encodeToString(sig);
        System.out.println("header.payload = " + in);
        System.out.println("signature bytes = " + sig.length + ", jwt parts = " + jwt.split("\\.").length);
        Fx.check(verify("SHA256withRSA", rsaPub, in, sig), "RS256 signature verifies with the public PEM");

        Fx.h("ES256 with an EC key from PEM (the trap is the signature format)");
        PrivateKey ecPriv = d.decode(Fx.read("ec-private-pkcs8.pem"), PrivateKey.class);
        PublicKey ecPub = d.decode(Fx.read("ec-public.pem"), PublicKey.class);
        String in2 = signingInput("ES256");
        byte[] der = sign("SHA256withECDSA", ecPriv, in2);
        byte[] raw = sign("SHA256withECDSAinP1363Format", ecPriv, in2);
        System.out.println("SHA256withECDSA              -> DER, starts 0x" + Integer.toHexString(der[0] & 0xff) + ", length varies run to run, 70-72: " + (der.length >= 70 && der.length <= 72) + "   (JWT readers reject this)");
        System.out.println("SHA256withECDSAinP1363Format -> R||S, length " + raw.length + "   (what RFC 7518 requires for ES256)");
        Fx.check(raw.length == 64, "ES256 signature is exactly 64 bytes");
        Fx.check(verify("SHA256withECDSAinP1363Format", ecPub, in2, raw), "ES256 signature verifies with the public PEM");

        Fx.h("EdDSA with an Ed25519 key from PEM");
        PrivateKey edPriv = d.decode(Fx.read("ed25519-private.pem"), PrivateKey.class);
        byte[] ed = sign("Ed25519", edPriv, signingInput("EdDSA"));
        System.out.println(edPriv.getAlgorithm() + " signature bytes = " + ed.length
                + "  (the PKCS#8 file holds no public key, so there is nothing to decode for verifying here)");
    }
}