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

package com.google.crypto.tink.daead;

import java.util.Iterator;
import com.google.crypto.tink.daead.internal.LegacyFullDeterministicAead;
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.DeterministicAead;
import com.google.crypto.tink.internal.PrimitiveWrapper;

public class DeterministicAeadWrapper implements PrimitiveWrapper<DeterministicAead, DeterministicAead>
{
    private static final DeterministicAeadWrapper WRAPPER;
    private static final PrimitiveConstructor<LegacyProtoKey, DeterministicAead> LEGACY_FULL_DAEAD_PRIMITIVE_CONSTRUCTOR;
    
    private static Bytes getOutputPrefix(final Key key) throws GeneralSecurityException {
        if (key instanceof DeterministicAeadKey) {
            return ((DeterministicAeadKey)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());
    }
    
    DeterministicAeadWrapper() {
    }
    
    @Override
    public DeterministicAead wrap(final KeysetHandleInterface handle, final MonitoringAnnotations annotations, final PrimitiveFactory<DeterministicAead> factory) throws GeneralSecurityException {
        final PrefixMap.Builder<DeterministicAeadWithId> builder = new PrefixMap.Builder<DeterministicAeadWithId>();
        for (int i = 0; i < handle.size(); ++i) {
            final KeysetHandleInterface.Entry entry = handle.getAt(i);
            if (entry.getStatus().equals(KeyStatus.ENABLED)) {
                final DeterministicAead deterministicAead = factory.create(entry);
                builder.put(getOutputPrefix(entry.getKey()), new DeterministicAeadWithId(deterministicAead, entry.getId()));
            }
        }
        MonitoringClient.Logger encLogger;
        MonitoringClient.Logger decLogger;
        if (!annotations.isEmpty()) {
            final MonitoringClient client = MutableMonitoringRegistry.globalInstance().getMonitoringClient();
            encLogger = client.createLogger(handle, annotations, "daead", "encrypt");
            decLogger = client.createLogger(handle, annotations, "daead", "decrypt");
        }
        else {
            encLogger = MonitoringUtil.DO_NOTHING_LOGGER;
            decLogger = MonitoringUtil.DO_NOTHING_LOGGER;
        }
        return new WrappedDeterministicAead(new DeterministicAeadWithId(factory.create(handle.getPrimary()), handle.getPrimary().getId()), builder.build(), encLogger, decLogger);
    }
    
    @Override
    public Class<DeterministicAead> getPrimitiveClass() {
        return DeterministicAead.class;
    }
    
    @Override
    public Class<DeterministicAead> getInputPrimitiveClass() {
        return DeterministicAead.class;
    }
    
    public static void register() throws GeneralSecurityException {
        MutablePrimitiveRegistry.globalInstance().registerPrimitiveWrapper((PrimitiveWrapper<Object, Object>)DeterministicAeadWrapper.WRAPPER);
        MutablePrimitiveRegistry.globalInstance().registerPrimitiveConstructor(DeterministicAeadWrapper.LEGACY_FULL_DAEAD_PRIMITIVE_CONSTRUCTOR);
    }
    
    public static void registerToInternalPrimitiveRegistry(final PrimitiveRegistry.Builder primitiveRegistryBuilder) throws GeneralSecurityException {
        primitiveRegistryBuilder.registerPrimitiveWrapper((PrimitiveWrapper<Object, Object>)DeterministicAeadWrapper.WRAPPER);
    }
    
    static {
        WRAPPER = new DeterministicAeadWrapper();
        LEGACY_FULL_DAEAD_PRIMITIVE_CONSTRUCTOR = PrimitiveConstructor.create(LegacyFullDeterministicAead::create, LegacyProtoKey.class, DeterministicAead.class);
    }
    
    private static class DeterministicAeadWithId
    {
        public final DeterministicAead daead;
        public final int id;
        
        public DeterministicAeadWithId(final DeterministicAead daead, final int id) {
            this.daead = daead;
            this.id = id;
        }
    }
    
    private static class WrappedDeterministicAead implements DeterministicAead
    {
        private final DeterministicAeadWithId primary;
        private final PrefixMap<DeterministicAeadWithId> allDaeads;
        private final MonitoringClient.Logger encLogger;
        private final MonitoringClient.Logger decLogger;
        
        public WrappedDeterministicAead(final DeterministicAeadWithId primary, final PrefixMap<DeterministicAeadWithId> allDaeads, final MonitoringClient.Logger encLogger, final MonitoringClient.Logger decLogger) {
            this.primary = primary;
            this.allDaeads = allDaeads;
            this.encLogger = encLogger;
            this.decLogger = decLogger;
        }
        
        @Override
        public byte[] encryptDeterministically(final byte[] plaintext, final byte[] associatedData) throws GeneralSecurityException {
            try {
                final byte[] result = this.primary.daead.encryptDeterministically(plaintext, associatedData);
                this.encLogger.log(this.primary.id, plaintext.length);
                return result;
            }
            catch (final GeneralSecurityException e) {
                this.encLogger.logFailure();
                throw e;
            }
        }
        
        @Override
        public byte[] decryptDeterministically(final byte[] ciphertext, final byte[] associatedData) throws GeneralSecurityException {
            for (final DeterministicAeadWithId aeadWithId : this.allDaeads.getAllWithMatchingPrefix(ciphertext)) {
                try {
                    final byte[] result = aeadWithId.daead.decryptDeterministically(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");
        }
    }
}