package com.ttstd.adbloopback.tls;

import javax.crypto.Cipher;
import javax.crypto.Mac;
import javax.crypto.spec.GCMParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.security.MessageDigest;

/** TLS 1.2 所需的密码学原语：SHA-256、HMAC-SHA256、TLS PRF、AES-128-GCM。 */
public final class TlsCrypto {

    public static byte[] sha256(byte[] data) throws Exception {
        return MessageDigest.getInstance("SHA-256").digest(data);
    }

    public static byte[] hmacSha256(byte[] key, byte[] data) throws Exception {
        Mac mac = Mac.getInstance("HmacSHA256");
        mac.init(new SecretKeySpec(key, "HmacSHA256"));
        return mac.doFinal(data);
    }

    /**
     * TLS 1.2 PRF，使用 SHA-256：
     * PRF(secret, label, seed) = P_SHA256(secret, label + seed)
     */
    public static byte[] prf(byte[] secret, String label, byte[] seed, int outLen) throws Exception {
        byte[] labelSeed = concat(label.getBytes("ASCII"), seed);
        byte[] result = new byte[outLen];
        byte[] prev = hmacSha256(secret, labelSeed); // A(1)
        int generated = 0;
        while (generated < outLen) {
            byte[] block = hmacSha256(secret, concat(prev, labelSeed));
            int take = Math.min(block.length, outLen - generated);
            System.arraycopy(block, 0, result, generated, take);
            generated += take;
            prev = hmacSha256(secret, prev); // A(i+1)
        }
        return result;
    }

    public static byte[] aesGcmEncrypt(byte[] key, byte[] nonce12, byte[] plaintext, byte[] aad) throws Exception {
        Cipher c = Cipher.getInstance("AES/GCM/NoPadding");
        c.init(Cipher.ENCRYPT_MODE, new SecretKeySpec(key, "AES"), new GCMParameterSpec(128, nonce12));
        c.updateAAD(aad);
        return c.doFinal(plaintext);
    }

    public static byte[] aesGcmDecrypt(byte[] key, byte[] nonce12, byte[] ciphertextWithTag, byte[] aad) throws Exception {
        Cipher c = Cipher.getInstance("AES/GCM/NoPadding");
        c.init(Cipher.DECRYPT_MODE, new SecretKeySpec(key, "AES"), new GCMParameterSpec(128, nonce12));
        c.updateAAD(aad);
        return c.doFinal(ciphertextWithTag);
    }

    public static byte[] concat(byte[]... arrays) {
        int len = 0;
        for (byte[] a : arrays) {
            if (a != null) len += a.length;
        }
        byte[] r = new byte[len];
        int off = 0;
        for (byte[] a : arrays) {
            if (a != null) {
                System.arraycopy(a, 0, r, off, a.length);
                off += a.length;
            }
        }
        return r;
    }

    private TlsCrypto() {}
}
