// 
// Decompiled by Procyon v0.6.0
// 

package com.google.crypto.tink.hybrid.internal;

import java.util.Arrays;
import com.google.crypto.tink.subtle.Bytes;
import java.security.GeneralSecurityException;
import com.google.errorprone.annotations.Immutable;

@Immutable
final class X25519HpkeKem implements HpkeKem
{
    private final HkdfHpkeKdf hkdf;
    private final X25519 x25519;
    
    X25519HpkeKem(final HkdfHpkeKdf hkdf) {
        this.hkdf = hkdf;
        X25519 x25519 = null;
        try {
            x25519 = X25519Conscrypt.create();
        }
        catch (final GeneralSecurityException e) {
            x25519 = new X25519Java();
        }
        this.x25519 = x25519;
    }
    
    private byte[] deriveKemSharedSecret(final byte[] dhSharedSecret, final byte[] senderEphemeralPublicKey, final byte[] recipientPublicKey) throws GeneralSecurityException {
        final byte[] kemContext = Bytes.concat(new byte[][] { senderEphemeralPublicKey, recipientPublicKey });
        return this.extractAndExpand(dhSharedSecret, kemContext);
    }
    
    private byte[] deriveKemSharedSecret(final byte[] dhSharedSecret, final byte[] senderEphemeralPublicKey, final byte[] recipientPublicKey, final byte[] senderPublicKey) throws GeneralSecurityException {
        final byte[] kemContext = Bytes.concat(new byte[][] { senderEphemeralPublicKey, recipientPublicKey, senderPublicKey });
        return this.extractAndExpand(dhSharedSecret, kemContext);
    }
    
    private byte[] extractAndExpand(final byte[] dhSharedSecret, final byte[] kemContext) throws GeneralSecurityException {
        final byte[] kemSuiteId = HpkeUtil.kemSuiteId(HpkeUtil.X25519_HKDF_SHA256_KEM_ID);
        return this.hkdf.extractAndExpand(null, dhSharedSecret, "eae_prk", kemContext, "shared_secret", kemSuiteId, this.hkdf.getMacLength());
    }
    
    HpkeKemEncapOutput encapsulateWithFixedEphemeralKey(final byte[] recipientPublicKey, final byte[] ephemeralPrivateKey, final byte[] ephemeralPublicKey) throws GeneralSecurityException {
        final byte[] dhSharedSecret = this.x25519.computeSharedSecret(ephemeralPrivateKey, recipientPublicKey);
        final byte[] kemSharedSecret = this.deriveKemSharedSecret(dhSharedSecret, ephemeralPublicKey, recipientPublicKey);
        return new HpkeKemEncapOutput(kemSharedSecret, ephemeralPublicKey);
    }
    
    @Override
    public HpkeKemEncapOutput encapsulate(final byte[] recipientPublicKey) throws GeneralSecurityException {
        final X25519.KeyPair ephemeral = this.x25519.generateKeyPair();
        return this.encapsulateWithFixedEphemeralKey(recipientPublicKey, ephemeral.privateKey, ephemeral.publicKey);
    }
    
    HpkeKemEncapOutput authEncapsulateWithFixedEphemeralKey(final byte[] recipientPublicKey, final byte[] ephemeralPrivateKey, final byte[] ephemeralPublicKey, final HpkeKemPrivateKey senderPrivateKey) throws GeneralSecurityException {
        final byte[] dhSharedSecret = Bytes.concat(new byte[][] { this.x25519.computeSharedSecret(ephemeralPrivateKey, recipientPublicKey), this.x25519.computeSharedSecret(senderPrivateKey.getSerializedPrivate().toByteArray(), recipientPublicKey) });
        final byte[] senderPublicKey = senderPrivateKey.getSerializedPublic().toByteArray();
        final byte[] kemSharedSecret = this.deriveKemSharedSecret(dhSharedSecret, ephemeralPublicKey, recipientPublicKey, senderPublicKey);
        return new HpkeKemEncapOutput(kemSharedSecret, ephemeralPublicKey);
    }
    
    @Override
    public HpkeKemEncapOutput authEncapsulate(final byte[] recipientPublicKey, final HpkeKemPrivateKey senderPrivateKey) throws GeneralSecurityException {
        final X25519.KeyPair ephemeral = this.x25519.generateKeyPair();
        return this.authEncapsulateWithFixedEphemeralKey(recipientPublicKey, ephemeral.privateKey, ephemeral.publicKey, senderPrivateKey);
    }
    
    @Override
    public byte[] decapsulate(final byte[] encapsulatedKey, final HpkeKemPrivateKey recipientPrivateKey) throws GeneralSecurityException {
        final byte[] dhSharedSecret = this.x25519.computeSharedSecret(recipientPrivateKey.getSerializedPrivate().toByteArray(), encapsulatedKey);
        return this.deriveKemSharedSecret(dhSharedSecret, encapsulatedKey, recipientPrivateKey.getSerializedPublic().toByteArray());
    }
    
    @Override
    public byte[] authDecapsulate(final byte[] encapsulatedKey, final HpkeKemPrivateKey recipientPrivateKey, final byte[] senderPublicKey) throws GeneralSecurityException {
        final byte[] privateKey = recipientPrivateKey.getSerializedPrivate().toByteArray();
        final byte[] dhSharedSecret = Bytes.concat(new byte[][] { this.x25519.computeSharedSecret(privateKey, encapsulatedKey), this.x25519.computeSharedSecret(privateKey, senderPublicKey) });
        final byte[] recipientPublicKey = recipientPrivateKey.getSerializedPublic().toByteArray();
        return this.deriveKemSharedSecret(dhSharedSecret, encapsulatedKey, recipientPublicKey, senderPublicKey);
    }
    
    @Override
    public byte[] getKemId() throws GeneralSecurityException {
        if (Arrays.equals(this.hkdf.getKdfId(), HpkeUtil.HKDF_SHA256_KDF_ID)) {
            return HpkeUtil.X25519_HKDF_SHA256_KEM_ID;
        }
        throw new GeneralSecurityException("Could not determine HPKE KEM ID");
    }
    
    @Immutable
    private static final class X25519Java implements X25519
    {
        @Override
        public KeyPair generateKeyPair() throws GeneralSecurityException {
            final byte[] privateKey = com.google.crypto.tink.subtle.X25519.generatePrivateKey();
            final byte[] publicKey = com.google.crypto.tink.subtle.X25519.publicFromPrivate(privateKey);
            return new KeyPair(privateKey, publicKey);
        }
        
        @Override
        public byte[] computeSharedSecret(final byte[] privateKey, final byte[] publicKey) throws GeneralSecurityException {
            return com.google.crypto.tink.subtle.X25519.computeSharedSecret(privateKey, publicKey);
        }
    }
}