package com.jscape.inet.ssh.protocol.v2.marshaling.algorithms;

import com.jscape.util.ArrayTools;
import com.jscape.util.io.ByteArrayOutputStream;

import javax.crypto.spec.SecretKeySpec;
import java.io.OutputStream;
import java.security.Key;
import java.security.Provider;

import static com.jscape.util.Parameters.assertGreaterOrEqual;
import static com.jscape.util.Parameters.assertNotNull;

/**
 * @author Alex Usun
 */
public class JceMac
        implements Mac {

    private final javax.crypto.Mac base;
    private final int digestLength;
    private final byte[] baseBuffer;
    private final ByteArrayOutputStream macBuffer;

    public static Mac macFor(String name, String provider, byte[] keyData, int digestLength)
            throws OperationException {
        try {
            javax.crypto.Mac mac = provider != null ? javax.crypto.Mac.getInstance(name, provider) : javax.crypto.Mac.getInstance(name);
            return macFor(mac, keyData, digestLength);
        } catch (Exception e) {
            throw new OperationException(e);
        }
    }

    public static Mac macFor(String name, Provider provider, byte[] keyData, int digestLength)
            throws OperationException {
        try {
            javax.crypto.Mac mac = provider != null ? javax.crypto.Mac.getInstance(name, provider) : javax.crypto.Mac.getInstance(name);
            return macFor(mac, keyData, digestLength);
        } catch (Exception e) {
            throw new OperationException(e);
        }
    }

    private static Mac macFor(javax.crypto.Mac mac, byte[] keyData, int digestLength)
            throws Exception {
        Key key = new SecretKeySpec(keyData, mac.getAlgorithm());
        mac.init(key);

        return new JceMac(mac, digestLength);
    }

    public JceMac(javax.crypto.Mac base, int digestLength) {
        assertNotNull(base);
        this.base = base;

        assertGreaterOrEqual(base.getMacLength(), digestLength, "Bad digest length.");
        this.digestLength = digestLength;

        this.baseBuffer = new byte[this.base.getMacLength()];
        this.macBuffer = new ByteArrayOutputStream(this.digestLength);
    }

    @Override
    public int digestLength() {
        return this.digestLength;
    }

    @Override
    public void apply(byte[] sequenceNumber, byte[] data, int offset, int length, OutputStream out)
            throws OperationException {
        try {
            this.base.update(sequenceNumber);
            this.base.update(data, offset, length);
            this.base.doFinal(this.baseBuffer, 0);
            out.write(this.baseBuffer, 0, this.digestLength);
        } catch (Exception e) {
            throw new OperationException(e);
        }
    }

    @Override
    public void assertMacValid(byte[] sequenceNumber, byte[] data, int offset, int length, byte[] mac)
            throws OperationException {
        this.macBuffer.reset();
        apply(sequenceNumber, data, offset, length, this.macBuffer);

        if (!ArrayTools.equals(this.macBuffer.buffer(), mac, this.digestLength)) {
            throw new InvalidMacException();
        }
    }

    @Override
    public String toString() {
        return "JceMac {"
                + "base=" + this.base.getAlgorithm()
                + ", digestLength=" + this.digestLength
                + '}';
    }
}
