using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;

namespace WinnyProtocolAnalyzer
{
    /// <summary>
    /// ストリーム暗号方式の RC4 を提供します。
    /// </summary>
    public class RC4
    {
        SBOX sbox = null;

        public byte[] RC4Key { get; private set; }

        /// <summary>
        /// 暗号化に使用する鍵を指定して RC4 クラスの新しいインスタンスを作成します。
        /// </summary>
        /// <param name="aobKey">暗号化に使用する鍵</param>
        public RC4(byte[] aobKey)
        {
            RC4Key = aobKey;
            sbox = new SBOX(aobKey);
        }

        /// <summary>
        /// 指定されたデータに対して RC4 を適用します。
        /// </summary>
        /// <param name="aobData">データ</param>
        /// <param name="iLength">データ長</param>
        public void Cipher(byte[] aobData, int iLength)
        {
            for (int i = 0; i < aobData.Length && i < iLength; i++)
            {
                aobData[i] = cipher(aobData[i]);
            }
        }

        private byte cipher(byte datum)
        {
            int q = (int)datum;

            if (q < 0)
            {
                q += 256;
            }

            q ^= sbox.getByte();

            return (byte)q;
        }
    }

    class SBOX
    {

        private int SIZE = 256;
        private byte[] sbox = null;
        private int idx = 0;
        private int num = 0;

        public SBOX(byte[] key)
        {
            initialize(key);
        }

        private void initialize(byte[] key)
        {
            sbox = new byte[SIZE];

            for (int i = 0; i < sbox.Length; i++)
            {
                sbox[i] = (byte)i;
            }

            byte[] sbox2 = new byte[SIZE];
            for (int i = 0; i < sbox2.Length; i++)
            {
                sbox2[i] = key[i % key.Length];
            }

            int j = 0;
            for (int i = 0; i < sbox.Length; i++)
            {
                int sboxi = (int)sbox[i];
                if (sboxi < 0)
                {
                    sboxi += 256;
                }

                int sbox2i = (int)sbox2[i];
                if (sbox2i < 0)
                {
                    sbox2i += 256;
                }

                j = (j + sboxi + sbox2i) % SIZE;

                byte tmp = sbox[i];
                sbox[i] = sbox[j];
                sbox[j] = tmp;
            }

            j = 0;

            for (int i = 0; i < sbox2.Length; i++)
            {
                sbox2[i] = 0;
            }
        }

        public byte getByte()
        {
            byte resultByte = 0;

            idx = (idx + 1) % SIZE;
            num = (num + sbox[idx]) % SIZE;

            byte tmp = sbox[idx];
            sbox[idx] = sbox[num];
            sbox[num] = tmp;

            resultByte = sbox[(sbox[idx] + sbox[num]) % SIZE];

            return resultByte;
        }
    }
}
