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

package com.google.crypto.tink.aead;

import java.util.Iterator;
import com.google.crypto.tink.aead.internal.LegacyFullAead;
import com.google.crypto.tink.internal.PrimitiveRegistry;
import com.google.crypto.tink.internal.MutablePrimitiveRegistry;
import com.google.crypto.tink.internal.MonitoringClient;
import com.google.crypto.tink.internal.MonitoringUtil;
import com.google.crypto.tink.internal.MutableMonitoringRegistry;
import com.google.crypto.tink.KeyStatus;
import com.google.crypto.tink.internal.PrefixMap;
import com.google.crypto.tink.internal.MonitoringAnnotations;
import com.google.crypto.tink.internal.KeysetHandleInterface;
import java.security.GeneralSecurityException;
import com.google.crypto.tink.util.Bytes;
import com.google.crypto.tink.Key;
import com.google.crypto.tink.internal.LegacyProtoKey;
import com.google.crypto.tink.internal.PrimitiveConstructor;
import com.google.crypto.tink.Aead;
import com.google.crypto.tink.internal.PrimitiveWrapper;

public class AeadWrapper implements PrimitiveWrapper<Aead, Aead>
{
    private static final AeadWrapper WRAPPER;
    private static final PrimitiveConstructor<LegacyProtoKey, Aead> LEGACY_FULL_AEAD_PRIMITIVE_CONSTRUCTOR;
    
    private static Bytes getOutputPrefix(final Key key) throws GeneralSecurityException {
        if (key instanceof AeadKey) {
            return ((AeadKey)key).getOutputPrefix();
        }
        if (key instanceof LegacyProtoKey) {
            return ((LegacyProtoKey)key).getOutputPrefix();
        }
        throw new GeneralSecurityException("Cannot get output prefix for key of class " + key.getClass().getName() + " with parameters " + key.getParameters());
    }
    
    AeadWrapper() {
    }
    
    @Override
    public Aead wrap(final KeysetHandleInterface keysetHandle, final MonitoringAnnotations annotations, final PrimitiveFactory<Aead> factory) throws GeneralSecurityException {
        final PrefixMap.Builder<AeadWithId> builder = new PrefixMap.Builder<AeadWithId>();
        for (int i = 0; i < keysetHandle.size(); ++i) {
            final KeysetHandleInterface.Entry entry = keysetHandle.getAt(i);
            if (entry.getStatus().equals(KeyStatus.ENABLED)) {
                builder.put(getOutputPrefix(entry.getKey()), new AeadWithId(factory.create(entry), entry.getId()));
            }
        }
        MonitoringClient.Logger encLogger;
        MonitoringClient.Logger decLogger;
        if (!annotations.isEmpty()) {
            final MonitoringClient client = MutableMonitoringRegistry.globalInstance().getMonitoringClient();
            encLogger = client.createLogger(keysetHandle, annotations, "aead", "encrypt");
            decLogger = client.createLogger(keysetHandle, annotations, "aead", "decrypt");
        }
        else {
            encLogger = MonitoringUtil.DO_NOTHING_LOGGER;
            decLogger = MonitoringUtil.DO_NOTHING_LOGGER;
        }
        return new WrappedAead(new AeadWithId(factory.create(keysetHandle.getPrimary()), keysetHandle.getPrimary().getId()), (PrefixMap)builder.build(), encLogger, decLogger);
    }
    
    @Override
    public Class<Aead> getPrimitiveClass() {
        return Aead.class;
    }
    
    @Override
    public Class<Aead> getInputPrimitiveClass() {
        return Aead.class;
    }
    
    public static void register() throws GeneralSecurityException {
        MutablePrimitiveRegistry.globalInstance().registerPrimitiveWrapper((PrimitiveWrapper<Object, Object>)AeadWrapper.WRAPPER);
        MutablePrimitiveRegistry.globalInstance().registerPrimitiveConstructor(AeadWrapper.LEGACY_FULL_AEAD_PRIMITIVE_CONSTRUCTOR);
    }
    
    public static void registerToInternalPrimitiveRegistry(final PrimitiveRegistry.Builder primitiveRegistryBuilder) throws GeneralSecurityException {
        primitiveRegistryBuilder.registerPrimitiveWrapper((PrimitiveWrapper<Object, Object>)AeadWrapper.WRAPPER);
    }
    
    static {
        WRAPPER = new AeadWrapper();
        LEGACY_FULL_AEAD_PRIMITIVE_CONSTRUCTOR = PrimitiveConstructor.create(LegacyFullAead::create, LegacyProtoKey.class, Aead.class);
    }
    
    private static class AeadWithId
    {
        public final Aead aead;
        public final int id;
        
        public AeadWithId(final Aead aead, final int id) {
            this.aead = aead;
            this.id = id;
        }
    }
    
    private static class WrappedAead implements Aead
    {
        private final AeadWithId primary;
        private final PrefixMap<AeadWithId> allAeads;
        private final MonitoringClient.Logger encLogger;
        private final MonitoringClient.Logger decLogger;
        
        private WrappedAead(final AeadWithId primary, final PrefixMap<AeadWithId> allAeads, final MonitoringClient.Logger encLogger, final MonitoringClient.Logger decLogger) {
            this.primary = primary;
            this.allAeads = allAeads;
            this.encLogger = encLogger;
            this.decLogger = decLogger;
        }
        
        @Override
        public byte[] encrypt(final byte[] plaintext, final byte[] associatedData) throws GeneralSecurityException {
            try {
                final byte[] result = this.primary.aead.encrypt(plaintext, associatedData);
                this.encLogger.log(this.primary.id, plaintext.length);
                return result;
            }
            catch (final GeneralSecurityException e) {
                this.encLogger.logFailure();
                throw e;
            }
        }
        
        @Override
        public byte[] decrypt(final byte[] ciphertext, final byte[] associatedData) throws GeneralSecurityException {
            for (final AeadWithId aeadWithId : this.allAeads.getAllWithMatchingPrefix(ciphertext)) {
                try {
                    final byte[] result = aeadWithId.aead.decrypt(ciphertext, associatedData);
                    this.decLogger.log(aeadWithId.id, ciphertext.length);
                    return result;
                }
                catch (final GeneralSecurityException ex) {
                    continue;
                }
                break;
            }
            this.decLogger.logFailure();
            throw new GeneralSecurityException("decryption failed");
        }
    }
}