package applet;

import javacard.framework.APDU;
import javacard.framework.ISO7816;
import javacard.framework.ISOException;
import javacard.framework.Util;
import javacard.security.AESKey;
import javacard.security.KeyBuilder;
import javacardx.crypto.Cipher;

public class Milenage {
    private static final short FAIL_CODE = (short) 0x9862;

    private Cipher rijndael = null;
    private byte[] opc;
    private byte[] key;

    Milenage(byte[] opc, byte[] key) {
        this.rijndael = Cipher.getInstance(Cipher.ALG_AES_BLOCK_128_ECB_NOPAD, false);
        AESKey theKey = (AESKey) KeyBuilder.buildKey(KeyBuilder.TYPE_AES, KeyBuilder.LENGTH_AES_128, false);
        theKey.setKey(key, (short)0);
        this.rijndael.init(theKey, Cipher.MODE_ENCRYPT);
        this.opc = opc;
        this.key = key;
//        Util.arrayCopyNonAtomic(opc, (short) 0, this.opc, (short)0, (short)16);
//        Util.arrayCopyNonAtomic(key, (short) 0, this.key, (short)0, (short)16);
    }

//    Milenage() {
//        this.opc = new byte[] {
//                (byte) 0x00, (byte) 0x01, (byte) 0x02, (byte) 0x03,
//                (byte) 0x04, (byte) 0x05, (byte) 0x06, (byte) 0x07,
//                (byte) 0x08, (byte) 0x09, (byte) 0x0a, (byte) 0x0b,
//                (byte) 0x0c, (byte) 0x0d, (byte) 0x0e, (byte) 0x0f,
//        };
//        byte[] key = new byte[] {
//                (byte) 0x00, (byte) 0x01, (byte) 0x02, (byte) 0x03,
//                (byte) 0x04, (byte) 0x05, (byte) 0x06, (byte) 0x07,
//                (byte) 0x08, (byte) 0x09, (byte) 0x0a, (byte) 0x0b,
//                (byte) 0x0c, (byte) 0x0d, (byte) 0x0e, (byte) 0x0f,
//        };
//        this(key);
//    }

    public void resetKeys(byte[] opc, byte[] key){
        this.rijndael = Cipher.getInstance(Cipher.ALG_AES_BLOCK_128_ECB_NOPAD, false);
        AESKey theKey = (AESKey) KeyBuilder.buildKey(KeyBuilder.TYPE_AES, KeyBuilder.LENGTH_AES_128, false);
        theKey.setKey(key, (short)0);
        this.rijndael.init(theKey, Cipher.MODE_ENCRYPT);
        this.opc = opc;
        this.key = key;
//        Util.arrayCopyNonAtomic(opc, (short) 0, this.opc, (short)0, (short)16);
//        Util.arrayCopyNonAtomic(key, (short) 0, this.key, (short)0, (short)16);
    }

    private void f2345(byte[] op_c, byte[] rand, byte[] res, byte[] ck, byte[] ik, byte[] ak) {
        byte[] rijndael_input = new byte[16];
        byte[] tmp = new byte[16];
        byte[] out = new byte[16];

        for (byte i = 0; i < 16; i++) {
            rijndael_input[i] = (byte) (op_c[i] ^ rand[i]);
        }
        this.rijndael.update(rijndael_input, (short) 0, (short) 16, tmp, (short) 0);

        for (byte i = 0; i < 16; i++) {
            rijndael_input[i] = (byte) (op_c[i] ^ tmp[i]);
        }
        rijndael_input[15] ^= (byte) 1;
        this.rijndael.update(rijndael_input, (short) 0, (short) 16, out, (short) 0);
        for (byte i = 0; i < 16; i++) {
            out[i] ^= op_c[i];
        }

        Util.arrayCopyNonAtomic(out, (short) 8, res, (short) 0, (short) 8);
        Util.arrayCopyNonAtomic(out, (short) 0, ak, (short) 0, (short) 6);

        for (byte i = 0; i < 16; i++) {
            byte pos = (byte)((byte)(i + (byte) 12) % (byte) 16);
            rijndael_input[pos] = (byte)(tmp[i] ^ op_c[i]);
        }
        rijndael_input[15] ^= (byte) 2;
        this.rijndael.update(rijndael_input, (short) 0, (short) 16, out, (short) 0);
        for (byte i = 0; i < 16; i++) {
            ck[i] = (byte)(out[i] ^ op_c[i]);
        }

        for (byte i = 0; i < 16; i++) {
            byte pos = (byte)((byte)(i + (byte) 8) % (byte) 16);
            rijndael_input[pos] = (byte)(tmp[i] ^ op_c[i]);
        }
        rijndael_input[15] ^= (byte) 4;
        this.rijndael.update(rijndael_input, (short) 0, (short) 16, out, (short) 0);
        for (byte i = 0; i < 16; i++) {
            ik[i] = (byte)(out[i] ^ op_c[i]);
        }
    }

    private void f1(byte[] op_c, byte[] rand, byte[] sqn, byte[] amf, byte[] mac_a) {
        byte[] rijndael_input = new byte[16];
        byte[] in = new byte[16];
        byte[] out = new byte[16];

        for (byte i = 0; i < 16; i++) {
            rijndael_input[i] = (byte) (op_c[i] ^ rand[i]);
        }
        this.rijndael.update(rijndael_input, (short) 0, (short) 16, out, (short) 0);

        for (byte i = 0; i < 6; i++) {
            in[i] = sqn[i];
            in[(byte)(i + 8)] = sqn[i];
        }
        for (byte i = 0; i < 2; i++) {
            in[(byte)(i + 6)] = amf[i];
            in[(byte)(i + 14)] = amf[i];
        }

        for (byte i = 0; i < 16; i++) {
            byte pos = (byte)((byte)(i + (byte) 8) % (byte) 16);
            rijndael_input[pos] = (byte)(in[i] ^ op_c[i]);
        }
        for (byte i = 0; i < 16; i++) {
            rijndael_input[i] ^= out[i];
        }
        this.rijndael.update(rijndael_input, (short) 0, (short) 16, out, (short) 0);

        for (byte i = 0; i < 8; i++) {
            mac_a[i] = (byte)(out[i] ^ op_c[i]);
        }
    }

    public void process(APDU apdu) throws ISOException {
        byte[] rand = new byte[16];
        byte[] autn = new byte[16];
        byte[] res = new byte[8];
        byte[] ck = new byte[16];
        byte[] ak = new byte[6];
        byte[] ik = new byte[16];
        byte[] sqn = new byte[6];
        byte[] amf = new byte[2];
        byte[] mac = new byte[8];
        byte[] kc = new byte[8];
        //byte[] resp = new byte[53];

        short lc;

        lc = apdu.setIncomingAndReceive();
        if (lc != apdu.getIncomingLength() || lc != (short) 34) {
            ISOException.throwIt(ISO7816.SW_WRONG_LENGTH);
        }

        short offset_cdata = apdu.getOffsetCdata();
        byte[] buffer = apdu.getBuffer();

        // TODO: Check RAND and AUTN length
        Util.arrayCopyNonAtomic(buffer, (short) (offset_cdata + (short) 1), rand, (short) 0, (short) 16);
        Util.arrayCopyNonAtomic(buffer, (short) (offset_cdata + (short) 18), autn, (short) 0, (short) 16);

        f2345(this.opc, rand, res, ck, ik, ak);

        // TODO: check SQN for synchronization
        for (byte i = 0; i < 6; i++) {
            sqn[i] = (byte) (autn[i] ^ ak[i]);
        }
        amf[0] = autn[6];
        amf[1] = autn[7];

        f1(this.opc, rand, sqn, amf, mac);

        for (byte i = 0; i < 8; i++) {
            if (mac[i] != autn[(byte)(i + (byte) 8)]) {
                ISOException.throwIt((short)0x9862);
            }
        }

        for (byte i = 0; i < 8; i++) {
            byte a = (byte) (ck[i] ^ ck[(byte)(i + (byte) 8)]);
            byte b = (byte) (ik[i] ^ ik[(byte)(i + (byte) 8)]);
            kc[i] = (byte) (a ^ b);
        }

        buffer[0] = (byte) 0xdb;
        buffer[1] = (byte) 0x08;
        Util.arrayCopyNonAtomic(res, (short) 0, buffer, (short) 2, (short) 8);
        buffer[10] = (byte) 0x10;
        Util.arrayCopyNonAtomic(ck, (short) 0, buffer, (short) 11, (short) 16);
        buffer[27] = (byte) 0x10;
        Util.arrayCopyNonAtomic(ik, (short) 0, buffer, (short) 28, (short) 16);
        buffer[44] = (byte) 0x08;
        Util.arrayCopyNonAtomic(kc, (short) 0, buffer, (short) 45, (short) 8);

        apdu.setOutgoingAndSend((short)0, (short) 53);
        ISOException.throwIt(ISO7816.SW_NO_ERROR);
    }
}
