using System;
using System.Diagnostics;
using System.IO;
using System.Linq;
using System.Security.Cryptography;

namespace TwlBackupBlock
{
    /// <summary>
    /// AES-CMACによるMACを生成します。
    /// 鍵は16バイト、24バイト、32バイトに対応しています。
    /// </summary>
    public class AesCmac
    {
        private const int BLOCK_SIZE = 16;
        private static readonly int[] LEGAL_KEY_SIZES = { 16, 24, 32 };

        /// <summary>
        /// 指定された鍵を利用してデータのMACを取得します。
        /// </summary>
        /// <param name="data">MACを発行したいデータ。</param>
        /// <param name="key">AES-CMACの鍵。</param>
        /// <returns><paramref name="data"/>に対するMAC。</returns>
        public static byte[] GetMac(byte[] data, byte[] key)
        {
            if (data == null)
            {
                throw new ArgumentNullException("data");
            }
            if (key == null)
            {
                throw new ArgumentNullException("key");
            }
            if (!CheckKeySize(key.Length))
            {
                string message = string.Format("key.Length is illegal (key.Length:{0})", key.Length);
                throw new ArgumentException("key", message);
            }

            var context = new AesCmacContext(block => Encrypt(block, key));
            
            context.Update(data);

            return context.GetMac();
        }

        /// <summary>
        /// 指定された鍵を利用して16バイトブロックをAESで暗号化します。
        /// </summary>
        /// <param name="block">暗号化するブロック。メソッド呼出し後は、暗号化されたブロック。</param>
        /// <param name="key">AESの鍵。</param>
        private static void Encrypt(byte[] block, byte[] key)
        {
            Debug.Assert(block.Length == BLOCK_SIZE);
            Debug.Assert(CheckKeySize(key.Length));

            using (var ms = new MemoryStream(block))
            {
                var aes = new RijndaelManaged();
                aes.BlockSize = BLOCK_SIZE * 8;
                aes.Mode = CipherMode.ECB;
                aes.Key = key;
                using (CryptoStream cs = new CryptoStream(ms, aes.CreateEncryptor(), CryptoStreamMode.Read))
                {
                    cs.Read(block, 0, block.Length);
                }
            }
        }

        /// <summary>
        /// 鍵のサイズが16バイト、24バイト、32バイトのいずれであるか判定します。
        /// </summary>
        /// <param name="keySize">鍵のバイトサイズ</param>
        /// <returns></returns>
        private static bool CheckKeySize(int keySize)
        {
            return LEGAL_KEY_SIZES.Any(e => (e == keySize));
        }

        /// <summary>
        /// AES-CMAC生成器です。データの入力履歴に応じたMAC生成器の内部状態を保持します。
        /// </summary>
        private class AesCmacContext
        {
            public delegate void Encrypt(byte[] block);

            private Encrypt encryptOperation;
            private byte[] block;
            private int filledIndex;

            /// <summary>
            /// 指定された暗号化処理をもとに、AesCmacContextクラスの新しいインスタンスを初期化します。
            /// </summary>
            /// <param name="op">16バイトブロックを暗号化するメソッド</param>
            public AesCmacContext(Encrypt op)
            {
                encryptOperation = op;
                block = Enumerable.Repeat<byte>(0, BLOCK_SIZE).ToArray();
                filledIndex = 0;
            }

            /// <summary>
            /// 指定された入力データによりMAC生成器を更新します。
            /// </summary>
            /// <param name="data">MACを発行したいデータ</param>
            public void Update(byte[] data)
            {
                for (int i = 0; i < data.Length; )
                {
                    if (filledIndex >= BLOCK_SIZE)
                    {
                        encryptOperation(block);
                        filledIndex = 0;
                    }
                    while (filledIndex < BLOCK_SIZE && i < data.Length)
                    {
                        block[filledIndex] ^= data[i];
                        ++filledIndex;
                        ++i;
                    }
                }
            }

            /// <summary>
            /// 現在までの入力履歴からMACを生成します。
            /// </summary>
            /// <returns>これまでの入力データに対するMAC</returns>
            public byte[] GetMac()
            {
                byte[] padding = Enumerable.Repeat<byte>(0, BLOCK_SIZE).ToArray();
                encryptOperation(padding);
                Multiply128(padding);

                if (filledIndex < BLOCK_SIZE)
                {
                    block[filledIndex] ^= 0x80;
                    Multiply128(padding);
                }

                for (int i = 0; i < BLOCK_SIZE; ++i)
                {
                    padding[i] ^= block[i];
                }

                encryptOperation(padding);

                return padding;
            }

            /// <summary>
            /// 16バイトの整数に128を乗じます。オーバーフローは無視します。
            /// </summary>
            /// <param name="block">16バイト整数</param>
            private void Multiply128(byte[] block)
            {
                bool isCarried = ((block[0] & 0x80) != 0);

                for (int i = 0; i < block.Length - 1; i++)
                {
                    block[i] = (byte)(((uint)block[i] << 1) | ((uint)block[i + 1] >> 7));
                }
                block[block.Length - 1] <<= 1;
                if (isCarried)
                {
                    block[block.Length - 1] = (byte)(block[block.Length - 1] ^ 0x87);
                }
            }
        }
    }
}