package com.trilead.ssh2.crypto.digest;

import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;

import javax.crypto.Mac;
import javax.crypto.ShortBufferException;
import javax.crypto.spec.SecretKeySpec;

/**
 * MAC.
 *
 * @author Christian Plattner, plattner@trilead.com
 * @version $Id: MAC.java,v 1.1 2007/10/15 12:49:57 cplattne Exp $
 */
public final class MAC {
    /**
     * From http://tools.ietf.org/html/rfc4253
     */
    private static final String HMAC_MD5 = "hmac-md5";

    /**
     * From http://tools.ietf.org/html/rfc4253
     */
    private static final String HMAC_MD5_96 = "hmac-md5-96";

    /**
     * From http://tools.ietf.org/html/rfc4253
     */
    private static final String HMAC_SHA1 = "hmac-sha1";

    /**
     * From http://tools.ietf.org/html/rfc4253
     */
    private static final String HMAC_SHA1_96 = "hmac-sha1-96";

    /**
     * From http://tools.ietf.org/html/rfc6668
     */
    private static final String HMAC_SHA2_256 = "hmac-sha2-256";

    /**
     * From http://tools.ietf.org/html/rfc6668
     */
    private static final String HMAC_SHA2_512 = "hmac-sha2-512";

    Mac mac;
    int outSize;
    int macSize;
    byte[] buffer;

    /* Higher Priority First */
    private static final String[] MAC_LIST = {
        HMAC_SHA2_256, HMAC_SHA2_512,
        HMAC_SHA1_96, HMAC_SHA1, HMAC_MD5_96, HMAC_MD5
    };

    public final static String[] getMacList() {
        return MAC_LIST;
    }

    public final static void checkMacList(String[] macs) {
        for (int i = 0; i < macs.length; i++)
            getKeyLen(macs[i]);
    }

    public final static int getKeyLen(String type) {
        if (HMAC_SHA1.equals(type) || HMAC_SHA1_96.equals(type))
            return 20;

        if (HMAC_MD5.equals(type) || HMAC_MD5_96.equals(type))
            return 16;

        if (HMAC_SHA2_256.equals(type))
            return 32;

        if (HMAC_SHA2_512.equals(type))
            return 64;

        throw new IllegalArgumentException("Unkown algorithm " + type);
    }

    public MAC(String type, byte[] key) {
        try {
            if (HMAC_SHA1.equals(type) || HMAC_SHA1_96.equals(type)) {
                mac = Mac.getInstance("HmacSHA1");
            }
            else if (HMAC_MD5.equals(type) || HMAC_MD5_96.equals(type)) {
                mac = Mac.getInstance("HmacMD5");
            }
            else if (HMAC_SHA2_256.equals(type)) {
                mac = Mac.getInstance("HmacSHA256");
            }
            else if (HMAC_SHA2_512.equals(type)) {
                mac = Mac.getInstance("HmacSHA512");
            }
            else
                throw new IllegalArgumentException("Unkown algorithm " + type);
        }
        catch (NoSuchAlgorithmException e) {
            throw new IllegalArgumentException("Unknown algorithm " + type, e);
        }

        macSize = mac.getMacLength();

        if (type.endsWith("-96")) {
            outSize = 12;
            buffer = new byte[macSize];
        }
        else {
            outSize = macSize;
            buffer = null;
        }

        try {
            mac.init(new SecretKeySpec(key, type));
        }
        catch (InvalidKeyException e) {
            throw new IllegalArgumentException(e);
        }
    }

    public final void initMac(int seq) {
        mac.reset();
        mac.update((byte)(seq >> 24));
        mac.update((byte)(seq >> 16));
        mac.update((byte)(seq >> 8));
        mac.update((byte)(seq));
    }

    public final void update(byte[] packetdata, int off, int len) {
        mac.update(packetdata, off, len);
    }

    public final void getMac(byte[] out, int off) {
        try {
            if (buffer != null) {
                mac.doFinal(buffer, 0);
                System.arraycopy(buffer, 0, out, off, out.length - off);
            }
            else {
                mac.doFinal(out, off);
            }
        }
        catch (ShortBufferException e) {
            throw new IllegalStateException(e);
        }
    }

    public final int size() {
        return outSize;
    }
}