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

namespace TwlBackupBlock
{
    /// <summary>
    /// TWLバックアップデータの暗号化・復号化処理を提供します。
    /// </summary>
    public class Utility
    {
        private const int KEY_SIZE = 16;
        private const int ROUND_SCALE = 16;

        private const int AES_MAC_SIZE = 16;
        private const int AES_IV_SIZE = 16;
        private const int AES_SIGN_HEADER_SIZE = 32;

        private const int BANNER_SIZE = 16 * 1024;
        private const int HEADER_SIZE = 160;
        private const int SIGNATURE_SIZE = 1024;

        private const int ENCRYPT_BANNER_SIZE = BANNER_SIZE + AES_SIGN_HEADER_SIZE;
        private const int ENCRYPT_HEADER_SIZE = HEADER_SIZE + AES_SIGN_HEADER_SIZE;
        private const int ENCRYPT_SIGNATURE_SIZE = SIGNATURE_SIZE + AES_SIGN_HEADER_SIZE;

        /// <summary>
        /// バックアップデータをブロックに分割して復号化をおこないます。
        /// </summary>
        /// <param name="data">バックアップデータ。</param>
        /// <param name="key">データ本体の復号鍵。</param>
        /// <returns>分割、復号化されたブロック。</returns>
        public static Blocks DecryptBackupData(byte[] data, byte[] key)
        {
            if (data == null)
            {
                throw new ArgumentNullException("data");
            }
            if (key == null)
            {
                throw new ArgumentNullException("key");
            }
            {
                int minDataSize = ENCRYPT_BANNER_SIZE + ENCRYPT_HEADER_SIZE + ENCRYPT_SIGNATURE_SIZE;
                if (data.Length < minDataSize)
                {
                    string message = string.Format("data.Length < {0} (data.Length:{1})", data, minDataSize);
                    throw new ArgumentException("data", message);
                }
            }
            if (key.Length != KEY_SIZE)
            {
                string message = string.Format("key.Length != {0} (key.Length:{1})", KEY_SIZE, key.Length);
                throw new ArgumentException("key", message);
            }

            Blocks blocks = new Blocks();

            using (var ms = new MemoryStream(data))
            using (var br = new BinaryReader(ms))
            {
                blocks.banner.body = new Body(DecryptBlock(br, BANNER_SIZE, key, out blocks.banner.mac, out blocks.banner.iv));
                HeaderBody headerBody = new HeaderBody(DecryptBlock(br, HEADER_SIZE, key, out blocks.header.mac, out blocks.header.iv));
                blocks.header.body = headerBody;

                if (!CheckSignature(headerBody))
                {
                    throw new Exception("ヘッダブロックのシグネチャが正しくありません");
                }

                blocks.signature.body = new SignatureBody(DecryptBlock(br, SIGNATURE_SIZE, key, out blocks.signature.mac, out blocks.signature.iv));

                int TMD_SIZE = (int)headerBody.fileSizes[0];
                int CONTENT_SIZE = (int)headerBody.fileSizes[1];
                int SAVE_DATA_SIZE = (int)headerBody.fileSizes[2];
                int SUB_BANNER_SIZE = (int)headerBody.fileSizes[3];

                blocks.tmd.body = new Body(DecryptBlock(br, TMD_SIZE, key, out blocks.tmd.mac, out blocks.tmd.iv));
                blocks.content.body = new Body(DecryptBlock(br, CONTENT_SIZE, key, out blocks.content.mac, out blocks.content.iv));
                blocks.saveData.body = new Body(DecryptBlock(br, SAVE_DATA_SIZE, key, out blocks.saveData.mac, out blocks.saveData.iv));
                blocks.subBanner.body = new Body(DecryptBlock(br, SUB_BANNER_SIZE, key, out blocks.subBanner.mac, out blocks.subBanner.iv));

                Debug.Assert(ms.Position == data.Length);
            }

            return blocks;
        }

        /// <summary>
        /// ヘッダのデータ本体のシグネチャが正しいか判定します。
        /// </summary>
        /// <param name="header">ヘッダのデータ本体。</param>
        /// <returns>シグネチャが正しい場合はtrue。それ以外の場合はfalse。</returns>
        private static bool CheckSignature(HeaderBody header)
        {
            byte[] signature = { (byte)'3', (byte)'F', (byte)'D', (byte)'T' };
            byte[] bytes = BitConverter.GetBytes(header.signature);
            Debug.Assert(signature.Length == bytes.Length);

            for (int i = 0; i < bytes.Length; i++)
            {
                if (bytes[i] != signature[i])
                {
                    return false;
                }
            }

            return true;
        }

        /// <summary>
        /// binaryReaderから読み込んだブロックを復号化します。
        /// </summary>
        /// <param name="binaryReader">ブロックを読み出すBinaryReader。</param>
        /// <param name="bodySize">暗号化前のブロックのデータ本体のバイトサイズ。</param>
        /// <param name="key">データ本体の復号鍵。</param>
        /// <param name="mac">ブロックのMAC。</param>
        /// <param name="iv">ブロックのIV。</param>
        /// <returns>ブロックのデータ本体。</returns>
        private static byte[] DecryptBlock(BinaryReader binaryReader, int bodySize, byte[] key, out byte[] mac, out byte[] iv)
        {
            Debug.Assert(key.Length == KEY_SIZE);

            int encryptBodySize = RoundUp(bodySize, ROUND_SCALE);

            byte[] body = binaryReader.ReadBytes(encryptBodySize);
            mac = binaryReader.ReadBytes(AES_MAC_SIZE);
            iv = binaryReader.ReadBytes(AES_IV_SIZE);

            return DecryptBody(body, key, iv, bodySize);
        }

        /// <summary>
        /// ブロックのデータ本体を復号化します。
        /// </summary>
        /// <param name="body">暗号化されたデータ本体。</param>
        /// <param name="key">データ本体の復号鍵。</param>
        /// <param name="iv">暗号化に利用したIV。</param>
        /// <param name="size">暗号化前のブロックのデータ本体のバイトサイズ。</param>
        /// <returns>復号化されたデータ本体。</returns>
        public static byte[] DecryptBody(byte[] body, byte[] key, byte[] iv, int size)
        {
            if (body == null)
            {
                throw new ArgumentNullException("body");
            }
            if (key == null)
            {
                throw new ArgumentNullException("key");
            }
            if (iv == null)
            {
                throw new ArgumentNullException("iv");
            }
            if (body.Length % ROUND_SCALE != 0)
            {
                string message = string.Format("body.Length % {0} != 0 (body.Length:{1})", ROUND_SCALE, body.Length);
                throw new ArgumentException("body", message);
            }
            if (key.Length != KEY_SIZE)
            {
                string message = string.Format("key.Length != {0} (key.Length:{1})", KEY_SIZE, key.Length);
                throw new ArgumentException("key", message);
            }
            if (iv.Length != AES_IV_SIZE)
            {
                string message = string.Format("iv.Length != {0} (iv.Length:{1})", AES_IV_SIZE, iv.Length);
                throw new ArgumentException("iv", message);
            }
            if (size < 0)
            {
                string message = string.Format("size < 0 (size:{0})", size);
                throw new ArgumentException("size", message);
            }
            if (size > body.Length)
            {
                string message = string.Format("size > body.Length (size:{0}, body.Length:{1})", size, body.Length);
                throw new ArgumentException("size", message);
            }

            byte[] result = new byte[size];

            using (var ms = new MemoryStream(body))
            {
                RijndaelManaged aes = new RijndaelManaged();
                aes.Mode = CipherMode.CBC;
                aes.Padding = PaddingMode.Zeros;
                aes.Key = key;
                aes.IV = iv;
                using (CryptoStream cs = new CryptoStream(ms, aes.CreateDecryptor(), CryptoStreamMode.Read))
                {
                    cs.Read(result, 0, result.Length);
                }
            }

            return result;
        }

        /// <summary>
        /// ブロックのデータ本体を暗号化します。
        /// </summary>
        /// <param name="body">復号化されたデータ本体。</param>
        /// <param name="key">データ本体の暗号鍵。</param>
        /// <param name="iv">暗号化に利用するIV。</param>
        /// <returns>暗号化されたデータ本体。</returns>
        public static byte[] EncryptBody(byte[] body, byte[] key, byte[] iv)
        {
            if (body == null)
            {
                throw new ArgumentNullException("body");
            }
            if (key == null)
            {
                throw new ArgumentNullException("key");
            }
            if (iv == null)
            {
                throw new ArgumentNullException("iv");
            }
            if (key.Length != KEY_SIZE)
            {
                string message = string.Format("key.Length != {0} (key.Length:{1})", KEY_SIZE, key.Length);
                throw new ArgumentException("key", message);
            }
            if (iv.Length != AES_IV_SIZE)
            {
                string message = string.Format("iv.Length != {0} (iv.Length:{1})", AES_IV_SIZE, iv.Length);
                throw new ArgumentException("iv", message);
            }

            byte[] result = new byte[RoundUp(body.Length, ROUND_SCALE)];

            using (var ms = new MemoryStream(body))
            {
                RijndaelManaged aes = new RijndaelManaged();
                aes.Mode = CipherMode.CBC;
                aes.Padding = PaddingMode.Zeros;
                aes.Key = key;
                aes.IV = iv;
                using (CryptoStream cs = new CryptoStream(ms, aes.CreateEncryptor(), CryptoStreamMode.Read))
                {
                    cs.Read(result, 0, result.Length);
                }
            }

            return result;
        }

        /// <summary>
        /// ブロックのデータ本体を暗号化します。内部的にはEncryptBody(body.GetBytes(), key, iv)を呼び出します。
        /// </summary>
        /// <param name="body">暗号化されたデータ本体。</param>
        /// <param name="key">データ本体の暗号鍵。</param>
        /// <param name="iv">暗号化に利用するIV。</param>
        /// <returns>暗号化されたデータ本体。</returns>
        public static byte[] EncryptBody(AbstractBody body, byte[] key, byte[] iv)
        {
            if (body == null)
            {
                throw new ArgumentNullException("body");
            }

            return EncryptBody(body.GetBytes(), key, iv);
        }

        /// <summary>
        /// ブロックのデータ本体からハッシュを取得します。
        /// </summary>
        /// <param name="data">暗号化前のデータ本体。</param>
        /// <returns>32バイトのハッシュ。</returns>
        public static byte[] GetSha256(byte[] data)
        {
            if (data == null)
            {
                throw new ArgumentNullException("data");
            }

            var sha = new SHA256Managed();

            byte[] buf = Enumerable.Repeat<byte>(0, RoundUp(data.Length, ROUND_SCALE)).ToArray();
            Array.Copy(data, 0, buf, 0, data.Length);

            return sha.ComputeHash(buf);
        }

        /// <summary>
        /// ハッシュからMACを取得します。
        /// </summary>
        /// <param name="data">データ本体のハッシュ。</param>
        /// <param name="key">AES-CMACの鍵。</param>
        /// <returns>16バイトのMAC。</returns>
        public static byte[] GetAesCmac(byte[] data, byte[] key)
        {
            return AesCmac.GetMac(data, key);
        }

        /// <summary>
        /// 基準値の倍数のうち、指定された整数以上で最小のものを返します。
        /// </summary>
        /// <param name="x">整数。</param>
        /// <param name="y">基準値。</param>
        /// <returns>yの倍数のうち、x以上で最小の整数。</returns>
        private static int RoundUp(int x, int y)
        {
            return (x + y - 1) / y * y;
        }

        /// <summary>
        /// インポートシーケンスと同様の手順でバックアップデータを検証します。
        /// ただし、このメソッドでは署名データの検証はおこないません。
        /// 署名以外の整合性がとれており、署名の整合性がとれていない場合、このメソッドはtrueを返します。
        /// </summary>
        /// <param name="data">検証にかけるバックアップデータ。</param>
        /// <param name="aesCbcKey">データ本体の暗号鍵。</param>
        /// <param name="aesCmacKey">MACの鍵。</param>
        /// <returns>検証に成功した場合はtrue。それ以外の場合はfalse。</returns>
        public static bool Verify(byte[] data, byte[] aesCbcKey, byte[] aesCmacKey)
        {
            if (data == null)
            {
                throw new ArgumentNullException("data");
            }
            if (aesCbcKey == null)
            {
                throw new ArgumentNullException("key");
            }
            if (aesCbcKey.Length != KEY_SIZE)
            {
                string message = string.Format("key.Length != {0} (key.Length:{1})", KEY_SIZE, aesCbcKey.Length);
                throw new ArgumentException("key", message);
            }

            {
                int minDataSize = ENCRYPT_BANNER_SIZE + ENCRYPT_HEADER_SIZE + ENCRYPT_SIGNATURE_SIZE;
                if (data.Length < minDataSize)
                {
                    Console.WriteLine("バックアップファイルサイズがバナー＋ヘッダ＋署名ブロック未満です");
                    return false;
                }
            }

            Blocks blocks = new Blocks();

            using (var ms = new MemoryStream(data))
            using (var br = new BinaryReader(ms))
            {
                blocks.banner.body = new Body(DecryptBlock(br, BANNER_SIZE, aesCbcKey, out blocks.banner.mac, out blocks.banner.iv));
                HeaderBody headerBody = new HeaderBody(DecryptBlock(br, HEADER_SIZE, aesCbcKey, out blocks.header.mac, out blocks.header.iv));
                blocks.header.body = headerBody;
                SignatureBody signatureBody = new SignatureBody(DecryptBlock(br, SIGNATURE_SIZE, aesCbcKey, out blocks.signature.mac, out blocks.signature.iv));
                blocks.signature.body = signatureBody;

                if (!CheckMac(blocks.header, aesCmacKey, blocks.header.mac))
                {
                    Console.WriteLine("ヘッダブロックのMACが一致しません");
                    return false;
                }
                if (!CheckSignature(headerBody))
                {
                    Console.WriteLine("ヘッダブロックのシグネチャが正しくありません");
                    return false;
                }
                if (!CheckMac(blocks.signature, aesCmacKey, blocks.signature.mac))
                {
                    Console.WriteLine("署名ブロックのMACが一致しません");
                    return false;
                }
                /* 署名データの検証 */
                if (!CheckHash(blocks.header, signatureBody.digest.header))
                {
                    Console.WriteLine("ヘッダブロックのハッシュが一致しません");
                    return false;
                }
                if (!CheckMac(blocks.banner, aesCmacKey, blocks.banner.mac))
                {
                    Console.WriteLine("バナーブロックのMACが一致しません");
                    return false;
                }
                if (!CheckHash(blocks.banner, signatureBody.digest.banner))
                {
                    Console.WriteLine("バナーブロックのハッシュが一致しません");
                    return false;
                }

                int TMD_SIZE = (int)headerBody.fileSizes[0];
                int CONTENT_SIZE = (int)headerBody.fileSizes[1];
                int SAVE_DATA_SIZE = (int)headerBody.fileSizes[2];
                int SUB_BANNER_SIZE = (int)headerBody.fileSizes[3];

                int restSize = RoundUp(TMD_SIZE, ROUND_SCALE) + RoundUp(CONTENT_SIZE, ROUND_SCALE) +
                    RoundUp(SAVE_DATA_SIZE, ROUND_SCALE) + RoundUp(SUB_BANNER_SIZE, ROUND_SCALE) +
                    AES_SIGN_HEADER_SIZE * 4;
                if (data.Length - ms.Position != restSize)
                {
                    Console.WriteLine("TMD、コンテンツ、publicセーブデータ、サブバナーブロックサイズが実際のサイズと一致しません");
                    return false;
                }
                blocks.tmd.body = new Body(DecryptBlock(br, TMD_SIZE, aesCbcKey, out blocks.tmd.mac, out blocks.tmd.iv));
                if (!CheckMac(blocks.tmd, aesCmacKey, blocks.tmd.mac))
                {
                    Console.WriteLine("TMDブロックのMACが一致しません");
                    return false;
                }
                if (!CheckHash(blocks.tmd, signatureBody.digest.section[0]))
                {
                    Console.WriteLine("TMDブロックのハッシュが一致しません");
                    return false;
                }
                blocks.content.body = new Body(DecryptBlock(br, CONTENT_SIZE, aesCbcKey, out blocks.content.mac, out blocks.content.iv));
                if (!CheckMac(blocks.content, aesCmacKey, blocks.content.mac))
                {
                    Console.WriteLine("コンテンツブロックのMACが一致しません");
                    return false;
                }
                if (!CheckHash(blocks.content, signatureBody.digest.section[1]))
                {
                    Console.WriteLine("コンテンツブロックのハッシュが一致しません");
                    return false;
                }
                blocks.saveData.body = new Body(DecryptBlock(br, SAVE_DATA_SIZE, aesCbcKey, out blocks.saveData.mac, out blocks.saveData.iv));
                if (!CheckMac(blocks.saveData, aesCmacKey, blocks.saveData.mac))
                {
                    Console.WriteLine("publicセーブデータブロックのMACが一致しません");
                    return false;
                }
                if (!CheckHash(blocks.saveData, signatureBody.digest.section[2]))
                {
                    Console.WriteLine("publicセーブデータブロックのハッシュが一致しません");
                    return false;
                }
                blocks.subBanner.body = new Body(DecryptBlock(br, SUB_BANNER_SIZE, aesCbcKey, out blocks.subBanner.mac, out blocks.subBanner.iv));
                if (!CheckMac(blocks.subBanner, aesCmacKey, blocks.subBanner.mac))
                {
                    Console.WriteLine("サブバナーブロックの内容とMACが一致しません");
                    return false;
                }
                if (!CheckHash(blocks.subBanner, signatureBody.digest.section[3]))
                {
                    Console.WriteLine("サブバナーブロックのハッシュが一致しません");
                    return false;
                }

                Debug.Assert(ms.Position == data.Length);
            }

            return true;
        }

        /// <summary>
        /// ハッシュによりブロックが正しいか判定します。
        /// </summary>
        /// <param name="block">ブロック。</param>
        /// <param name="hash">ハッシュ。</param>
        /// <returns>ハッシュがブロックから生成したものと一致した場合はtrue。それ以外の場合はfalse。</returns>
        private static bool CheckHash(Block block, byte[] hash)
        {
            byte[] tryHash = GetSha256(block.body.GetBytes());
            Debug.Assert(tryHash.Length == hash.Length);

            for (int i = 0; i < tryHash.Length; i++)
            {
                if (tryHash[i] != hash[i])
                {
                    return false;
                }
            }
            return true;
        }

        /// <summary>
        /// MACによりブロックが正しいか判定します。
        /// </summary>
        /// <param name="block">ブロック。</param>
        /// <param name="key">MACの鍵。</param>
        /// <param name="mac">MAC。</param>
        /// <returns>MACがブロックから生成したものと一致した場合はtrue。それ以外の場合はfalse。</returns>
        private static bool CheckMac(Block block, byte[] key, byte[] mac)
        {
            byte[] tryMac = GetAesCmac(GetSha256(block.body.GetBytes()), key);
            Debug.Assert(tryMac.Length == mac.Length);

            for (int i = 0; i < tryMac.Length; ++i)
            {
                if (tryMac[i] != mac[i])
                {
                    return false;
                }
            }
            return true;
        }
    }
}