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

package com.google.crypto.tink.aead;

import java.nio.BufferUnderflowException;
import com.google.crypto.tink.proto.KeyData;
import com.google.protobuf.ByteString;
import java.nio.ByteBuffer;
import com.google.crypto.tink.Key;
import com.google.crypto.tink.internal.MutablePrimitiveRegistry;
import com.google.crypto.tink.InsecureSecretKeyAccess;
import com.google.crypto.tink.internal.MutableSerializationRegistry;
import com.google.crypto.tink.internal.ProtoKeySerialization;
import com.google.crypto.tink.internal.MutableKeyCreationRegistry;
import com.google.protobuf.InvalidProtocolBufferException;
import com.google.protobuf.ExtensionRegistryLite;
import java.security.GeneralSecurityException;
import com.google.crypto.tink.TinkProtoParametersFormat;
import com.google.crypto.tink.proto.OutputPrefixType;
import com.google.crypto.tink.proto.KeyTemplate;
import java.util.Collections;
import java.util.HashSet;
import java.util.Set;
import com.google.crypto.tink.Parameters;
import com.google.crypto.tink.Aead;

public final class KmsEnvelopeAead implements Aead
{
    private static final byte[] EMPTY_AAD;
    private final String typeUrlForParsing;
    private final Parameters parametersForNewKeys;
    private final Aead remote;
    private static final int LENGTH_ENCRYPTED_DEK = 4;
    private static final int MAX_LENGTH_ENCRYPTED_DEK = 4096;
    private static final Set<String> supportedDekKeyTypes;
    
    private static Set<String> listSupportedDekKeyTypes() {
        final HashSet<String> dekKeyTypeUrls = new HashSet<String>();
        dekKeyTypeUrls.add("type.googleapis.com/google.crypto.tink.AesGcmKey");
        dekKeyTypeUrls.add("type.googleapis.com/google.crypto.tink.ChaCha20Poly1305Key");
        dekKeyTypeUrls.add("type.googleapis.com/google.crypto.tink.XChaCha20Poly1305Key");
        dekKeyTypeUrls.add("type.googleapis.com/google.crypto.tink.AesCtrHmacAeadKey");
        dekKeyTypeUrls.add("type.googleapis.com/google.crypto.tink.AesGcmSivKey");
        dekKeyTypeUrls.add("type.googleapis.com/google.crypto.tink.AesEaxKey");
        return Collections.unmodifiableSet((Set<? extends String>)dekKeyTypeUrls);
    }
    
    public static boolean isSupportedDekKeyType(final String dekKeyTypeUrl) {
        return KmsEnvelopeAead.supportedDekKeyTypes.contains(dekKeyTypeUrl);
    }
    
    private Parameters getRawParameters(final KeyTemplate dekTemplate) throws GeneralSecurityException {
        final KeyTemplate rawTemplate = KeyTemplate.newBuilder(dekTemplate).setOutputPrefixType(OutputPrefixType.RAW).build();
        return TinkProtoParametersFormat.parse(rawTemplate.toByteArray());
    }
    
    @Deprecated
    public KmsEnvelopeAead(final KeyTemplate dekTemplate, final Aead remote) throws GeneralSecurityException {
        if (!isSupportedDekKeyType(dekTemplate.getTypeUrl())) {
            throw new IllegalArgumentException("Unsupported DEK key type: " + dekTemplate.getTypeUrl() + ". Only Tink AEAD key types are supported.");
        }
        this.typeUrlForParsing = dekTemplate.getTypeUrl();
        this.parametersForNewKeys = this.getRawParameters(dekTemplate);
        this.remote = remote;
    }
    
    public static Aead create(final AeadParameters dekParameters, final Aead remote) throws GeneralSecurityException {
        KeyTemplate dekTemplate;
        try {
            dekTemplate = KeyTemplate.parseFrom(TinkProtoParametersFormat.serialize(dekParameters), ExtensionRegistryLite.getEmptyRegistry());
        }
        catch (final InvalidProtocolBufferException e) {
            throw new GeneralSecurityException(e);
        }
        return new KmsEnvelopeAead(dekTemplate, remote);
    }
    
    @Override
    public byte[] encrypt(final byte[] plaintext, final byte[] associatedData) throws GeneralSecurityException {
        final Key key = MutableKeyCreationRegistry.globalInstance().createKey(this.parametersForNewKeys, null);
        final ProtoKeySerialization serialization = MutableSerializationRegistry.globalInstance().serializeKey(key, ProtoKeySerialization.class, InsecureSecretKeyAccess.get());
        final byte[] dek = serialization.getValue().toByteArray();
        final byte[] encryptedDek = this.remote.encrypt(dek, KmsEnvelopeAead.EMPTY_AAD);
        if (encryptedDek.length > 4096) {
            throw new GeneralSecurityException("length of encrypted DEK too large");
        }
        final Aead aead = MutablePrimitiveRegistry.globalInstance().getPrimitive(key, Aead.class);
        final byte[] payload = aead.encrypt(plaintext, associatedData);
        return this.buildCiphertext(encryptedDek, payload);
    }
    
    @Override
    public byte[] decrypt(final byte[] ciphertext, final byte[] associatedData) throws GeneralSecurityException {
        try {
            final ByteBuffer buffer = ByteBuffer.wrap(ciphertext);
            final int encryptedDekSize = buffer.getInt();
            if (encryptedDekSize <= 0 || encryptedDekSize > 4096 || encryptedDekSize > ciphertext.length - 4) {
                throw new GeneralSecurityException("length of encrypted DEK too large");
            }
            final byte[] encryptedDek = new byte[encryptedDekSize];
            buffer.get(encryptedDek, 0, encryptedDekSize);
            final byte[] payload = new byte[buffer.remaining()];
            buffer.get(payload, 0, buffer.remaining());
            final byte[] dek = this.remote.decrypt(encryptedDek, KmsEnvelopeAead.EMPTY_AAD);
            final ProtoKeySerialization serialization = ProtoKeySerialization.create(this.typeUrlForParsing, ByteString.copyFrom(dek), KeyData.KeyMaterialType.SYMMETRIC, OutputPrefixType.RAW, null);
            final Key key = MutableSerializationRegistry.globalInstance().parseKey(serialization, InsecureSecretKeyAccess.get());
            final Aead aead = MutablePrimitiveRegistry.globalInstance().getPrimitive(key, Aead.class);
            return aead.decrypt(payload, associatedData);
        }
        catch (final IndexOutOfBoundsException | BufferUnderflowException | NegativeArraySizeException e) {
            throw new GeneralSecurityException("invalid ciphertext", e);
        }
    }
    
    private byte[] buildCiphertext(final byte[] encryptedDek, final byte[] payload) {
        return ByteBuffer.allocate(4 + encryptedDek.length + payload.length).putInt(encryptedDek.length).put(encryptedDek).put(payload).array();
    }
    
    static {
        EMPTY_AAD = new byte[0];
        supportedDekKeyTypes = listSupportedDekKeyTypes();
    }
}