using System;
using System.Runtime.InteropServices;
using System.Security.Cryptography;
using System.Linq;
using Tools;
using System.IO;

namespace Tools.Crypto
{
    static unsafe class MemoryHelper
    {
        public static byte[] UMToByteArray(void* ptr, int size)
        {
            byte[] output = new byte[size];
            Marshal.Copy((IntPtr)ptr, output, 0, size);
            return output;
        }

        public static void ByteArrayToUM(byte[] input, void* ptr, int size)
        {
            Marshal.Copy(input, 0, (IntPtr)ptr, size);
        }

        public static T[] UMToArray<T>(T* ptr, int size) where T : unmanaged
        {
            T[] output = new T[size];
            fixed (T* outputPtr = output)
                Buffer.MemoryCopy(ptr, outputPtr, size * sizeof(T), size * sizeof(T));
            return output;
        }
        public static void ArrayToUM<T>(T[] input, T* ptr, int size) where T : unmanaged
        {
            fixed (T* inputPtr = input)
                Buffer.MemoryCopy(inputPtr, ptr, sizeof(T) * size, sizeof(T) * size);
        }
    }
    public unsafe struct RSAHashSignature
    {
        public const int Length = 32;

        internal fixed byte sig[Length];

        public byte[] Signature
        {
            get
            {
                fixed (byte* ptr = sig)
                {
                    return MemoryHelper.UMToByteArray(ptr, Length);
                }
            }
            set
            {
                fixed (byte* ptr = sig)
                    MemoryHelper.ByteArrayToUM(value, ptr, Length);
            }
        }

        public static RSAHashSignature CreateSignature(RSACryptoServiceProvider rsa, SHA256 algo, byte[] data)
        {
            return rsa.SignData(data, algo).ToUnmanaged<RSAHashSignature>();
        }
        public static RSAHashSignature CreateSignature(RSACryptoServiceProvider rsa, byte[] data)
        {
            using (var algo = SHA256.Create())
                return CreateSignature(rsa, algo, data);
        }

        public bool Verify(RSACryptoServiceProvider rsa, SHA256 algo, byte[] data)
        {
            return rsa.VerifyData(data, algo, Signature);
        }
        public bool Verify(RSACryptoServiceProvider rsa, byte[] data)
        {
            using (var algo = SHA256.Create())
                return Verify(rsa, algo, data);
        }
    }
    public unsafe struct RSACiphertext
    {
        public const int Length = 128;

        internal fixed byte ciphertext[Length];

        public byte[] Ciphertext
        {
            get
            {
                fixed (byte* ptr = ciphertext)
                    return MemoryHelper.UMToByteArray(ptr, Length);
            }
            set
            {
                fixed (byte* ptr = ciphertext)
                    MemoryHelper.ByteArrayToUM(value, ptr, Length);
            }
        }

        public void EncryptSingleBlock(RSACryptoServiceProvider rsa, byte[] plain)
        {
            Ciphertext = rsa.Encrypt(plain, false);
        }
        public byte[] DecryptSingleBlock(RSACryptoServiceProvider rsa)
        {
            return rsa.Decrypt(Ciphertext, false);
        }

        public unsafe static RSACiphertext[] Encrypt(RSACryptoServiceProvider rsa, byte[] plain)
        {
            byte[] ct = rsa.Encrypt(plain, false);

            if (ct.Length == Length) return new RSACiphertext[] { new RSACiphertext() { Ciphertext = ct } };
            else if (ct.Length % Length == 0)
            {
                int ciphers = ct.Length / Length;
                RSACiphertext* cipher = stackalloc RSACiphertext[ciphers];

                Marshal.Copy(ct, 0, (IntPtr)cipher, ciphers * sizeof(RSACiphertext));

                return MemoryHelper.UMToArray<RSACiphertext>(cipher, ciphers);
            }
            else
            {
                //Probably won't happen?
                int ciphers = (ct.Length / Length) + 1;
                RSACiphertext* cipher = stackalloc RSACiphertext[ciphers];

                Marshal.Copy(ct, 0, (IntPtr)cipher, ciphers * sizeof(RSACiphertext));

                return MemoryHelper.UMToArray<RSACiphertext>(cipher, ciphers);
            }
        }
        public static byte[] Decrypt(RSACryptoServiceProvider rsa, RSACiphertext[] cipher)
        {
            byte[] byts = new byte[sizeof(RSACiphertext) * cipher.Length];
            fixed (RSACiphertext* input = cipher)
            {
                fixed (byte* output = byts)
                {
                    Buffer.MemoryCopy(input, output, byts.Length, byts.Length);
                }
            }
            return rsa.Decrypt(byts, false);
        }
    }

    public unsafe struct AESKey
    {
        public const int KeySize = 32;
        public const int IVSize = 16;

        internal fixed byte key[KeySize];
        internal fixed byte iv[IVSize];

        /// <summary>
        /// The Key (256 bit)
        /// </summary>
        public byte[] Key
        {
            get
            {
                byte[] bytes = new byte[KeySize];
                fixed (byte* k = key)
                {
                    Marshal.Copy((IntPtr)k, bytes, 0, KeySize);
                }
                return bytes;
            }
            set
            {
                if (value.Length != KeySize) throw new ArgumentException(nameof(value) + " must be exaclty " + KeySize + " bytes (not " + value.Length + ")");
                fixed (byte* k = key)
                {
                    Marshal.Copy(value, 0, (IntPtr)k, KeySize);
                }
            }
        }

        /// <summary>
        /// The IV (128 bits)
        /// </summary>
        public byte[] IV
        {
            get
            {
                byte[] bytes = new byte[IVSize];
                fixed (byte* k = iv)
                {
                    Marshal.Copy((IntPtr)k, bytes, 0, IVSize);
                }
                return bytes;
            }
            set
            {
                if (value.Length != IVSize) throw new ArgumentException(nameof(value) + " must be exaclty " + IVSize + " bytes (not " + value.Length + ")");
                fixed (byte* k = iv)
                {
                    Marshal.Copy(value, 0, (IntPtr)k, IVSize);
                }
            }
        }


        /// <summary>
        /// Binary serialisation of this key
        /// </summary>
        public byte[] BinaryData
        {
            get
            {
                return this.ToByteArrayUnmanaged();
            }
            set
            {
                if (value.Length < sizeof(AESKey)) throw new ArgumentException(nameof(value) + " must be at least " + sizeof(AESKey) + " bytes (not " + value.Length + ")");
                fixed (AESKey* k = &this)
                {
                    Marshal.Copy(value, 0, (IntPtr)k, sizeof(AESKey));
                }
            }
        }

        /// <summary>
        /// Create a new key
        /// </summary>
        /// <returns>The new AES Key</returns>
        public static AESKey NewKey()
        {
            using (RNGCryptoServiceProvider rng = new RNGCryptoServiceProvider())
            {
                byte[] buffer = new byte[sizeof(AESKey)];
                rng.GetBytes(buffer);
                return buffer.ToUnmanaged<AESKey>();
            }
        }

        /// <summary>
        /// Set the key and iv to an AesCryptoServiceProvider
        /// </summary>
        /// <param name="r">The CSP</param>
        public void ToCSP(AesCryptoServiceProvider r)
        {
            r.KeySize = 256;
            r.BlockSize = 128;

            r.Key = Key;
            r.IV = IV;
        }

        public void FromAes(Aes aes)
        {
            Key = aes.Key;
            IV = aes.IV;
        }

        public void ToAes(Aes aes)
        {
            aes.KeySize = 256;
            aes.BlockSize = 128;

            aes.Key = Key;
            aes.IV = IV;
        }

        /// <summary>
        /// Get the key and iv from and AESCryptoServiceProvider
        /// </summary>
        /// <param name="r">The CSP</param>
        public void FromCSP(AesCryptoServiceProvider r)
        {
            Key = r.Key;
            IV = r.IV;
        }
        /// <summary>
        /// Initialise a new AESKey from an AESCryptoServiceProvider
        /// </summary>
        /// <param name="aes">The AES key</param>
        public AESKey(AesCryptoServiceProvider aes)
        {
            FromCSP(aes);
        }

        public AESKey(Aes aes)
        {
            FromAes(aes);
        }
    }

    public unsafe struct RSAPublicKey
    {
        public const int ModulusSize = 128;
        public const int ExponentSize = 3;

        internal fixed byte mod[ModulusSize];
        internal fixed byte exp[ExponentSize];

        /// <summary>
        /// The modulus of this key
        /// </summary>
        public byte[] Modulus
        {
            get
            {
                byte[] bytes = new byte[ModulusSize];
                fixed (byte* m = mod)
                {
                    Marshal.Copy((IntPtr)m, bytes, 0, ModulusSize);
                }
                return bytes;
            }
            set
            {
                if (value.Length != ModulusSize) throw new ArgumentException(nameof(value) + " must be exaclty " + ModulusSize + " bytes (not " + value.Length + ")");
                fixed (byte* m = mod)
                {
                    Marshal.Copy(value, 0, (IntPtr)m, ModulusSize);
                }
            }
        }
        /// <summary>
        /// The public exponent of this key
        /// </summary>
        public byte[] Exponent
        {
            get
            {
                byte[] bytes = new byte[ExponentSize];
                fixed (byte* m = exp)
                {
                    Marshal.Copy((IntPtr)m, bytes, 0, ExponentSize);
                }
                return bytes;
            }
            set
            {
                if (value.Length != ExponentSize) throw new ArgumentException(nameof(value) + " must be exaclty " + ExponentSize + " bytes (not " + value.Length + ")");
                fixed (byte* m = exp)
                {
                    Marshal.Copy(value, 0, (IntPtr)m, ExponentSize);
                }
            }
        }

        /// <summary>
        /// Binary serialisation of this key
        /// </summary>
        public byte[] BinaryData
        {
            get
            {
                return this.ToByteArrayUnmanaged();
            }
            set
            {
                if (value.Length < sizeof(RSAPublicKey)) throw new ArgumentException(nameof(value) + " must be at least " + sizeof(RSAPublicKey) + " bytes (not " + value.Length + ")");
                fixed (RSAPublicKey* k = &this)
                {
                    Marshal.Copy(value, 0, (IntPtr)k, sizeof(RSAPublicKey));
                }
            }
        }

        /// <summary>
        /// Set the public key to a RSACryptoServiceProvider
        /// </summary>
        /// <param name="csp">The CSP to set the key to</param>
        public void ToCSP(RSACryptoServiceProvider csp)
        {
            var p = csp.ExportParameters(false);
            p.Modulus = Modulus;
            p.Exponent = Exponent;
            csp.ImportParameters(p);
        }
        /// <summary>
        /// Get the public key information from an RSACryptoServiceProvider and return it in an RSAPublicKey struct
        /// </summary>
        /// <param name="csp">The CSP</param>
        /// <returns>A new RSAPublicKey struct</returns>
        public static RSAPublicKey FromCSP(RSACryptoServiceProvider csp)
        {
            RSAPublicKey rp = new RSAPublicKey();
            var p = csp.ExportParameters(false);
            rp.Modulus = p.Modulus;
            rp.Exponent = p.Exponent;
            return rp;
        }

        public static implicit operator byte[](RSAPublicKey rp)
        {
            return rp.ToByteArrayUnmanaged();
        }
        public static explicit operator RSAPublicKey(byte[] byt)
        {
            return byt.ToUnmanaged<RSAPublicKey>();
        }
    }

    public unsafe struct RSAPrivateKey
    {

        public override string ToString()
        {
            return BitConverter.ToString(this.ToByteArrayUnmanaged()).Replace("-", "");
        }

        public override bool Equals(object obj)
        {
            return obj is RSAPrivateKey priv && priv.ToByteArrayUnmanaged().SequenceEqual(this.ToByteArrayUnmanaged());
        }

        public override int GetHashCode()
        {
            unchecked
            {
                return (int)DamienG.Security.Cryptography.Crc32.Compute(this.ToByteArrayUnmanaged());
            }
        }

        public RSAPublicKey PublicPart => new RSAPublicKey() { Modulus = Modulus, Exponent = Exponent };

        public const int ModulusLength = 128;
        public const int ExponentLength = 3;
        public const int DLength = 128;
        public const int PLength = 64;
        public const int QLength = 64;
        public const int DPLength = 64;
        public const int DQLength = 64;
        public const int InverseQLength = 64;

        internal fixed byte modulus[ModulusLength];
        internal fixed byte exponent[ExponentLength];
        internal fixed byte d[DLength];
        internal fixed byte p[PLength];
        internal fixed byte q[QLength];
        internal fixed byte dp[DPLength];
        internal fixed byte dq[DQLength];
        internal fixed byte inverseQ[InverseQLength];

        public byte[] Modulus
        {
            get
            {
                byte[] bytes = new byte[ModulusLength];
                fixed (byte* m = modulus)
                {
                    Marshal.Copy((IntPtr)m, bytes, 0, ModulusLength);
                }
                return bytes;
            }
            set
            {
                if (value.Length != ModulusLength) throw new ArgumentException(nameof(value) + " must be exaclty " + ModulusLength + " bytes (not " + value.Length + ")");
                fixed (byte* m = modulus)
                {
                    Marshal.Copy(value, 0, (IntPtr)m, ModulusLength);
                }
            }
        }

        public byte[] Exponent
        {
            get
            {
                byte[] bytes = new byte[ExponentLength];
                fixed (byte* m = exponent)
                {
                    Marshal.Copy((IntPtr)m, bytes, 0, ExponentLength);
                }
                return bytes;
            }
            set
            {
                if (value.Length != ExponentLength) throw new ArgumentException(nameof(value) + " must be exaclty " + ExponentLength + " bytes (not " + value.Length + ")");
                fixed (byte* m = exponent)
                {
                    Marshal.Copy(value, 0, (IntPtr)m, ExponentLength);
                }
            }
        }

        public byte[] D
        {
            get
            {
                byte[] bytes = new byte[DLength];
                fixed (byte* m = d)
                {
                    Marshal.Copy((IntPtr)m, bytes, 0, DLength);
                }
                return bytes;
            }
            set
            {
                if (value.Length != DLength) throw new ArgumentException(nameof(value) + " must be exaclty " + DLength + " bytes (not " + value.Length + ")");
                fixed (byte* m = d)
                {
                    Marshal.Copy(value, 0, (IntPtr)m, DLength);
                }
            }
        }

        public byte[] P
        {
            get
            {
                byte[] bytes = new byte[PLength];
                fixed (byte* m = p)
                {
                    Marshal.Copy((IntPtr)m, bytes, 0, PLength);
                }
                return bytes;
            }
            set
            {
                if (value.Length != PLength) throw new ArgumentException(nameof(value) + " must be exaclty " + PLength + " bytes (not " + value.Length + ")");
                fixed (byte* m = p)
                {
                    Marshal.Copy(value, 0, (IntPtr)m, PLength);
                }
            }
        }

        public byte[] Q
        {
            get
            {
                byte[] bytes = new byte[QLength];
                fixed (byte* m = q)
                {
                    Marshal.Copy((IntPtr)m, bytes, 0, QLength);
                }
                return bytes;
            }
            set
            {
                if (value.Length != QLength) throw new ArgumentException(nameof(value) + " must be exaclty " + QLength + " bytes (not " + value.Length + ")");
                fixed (byte* m = q)
                {
                    Marshal.Copy(value, 0, (IntPtr)m, QLength);
                }
            }
        }

        public byte[] DP
        {
            get
            {
                byte[] bytes = new byte[DPLength];
                fixed (byte* m = dp)
                {
                    Marshal.Copy((IntPtr)m, bytes, 0, DPLength);
                }
                return bytes;
            }
            set
            {
                if (value.Length != DPLength) throw new ArgumentException(nameof(value) + " must be exaclty " + DPLength + " bytes (not " + value.Length + ")");
                fixed (byte* m = dp)
                {
                    Marshal.Copy(value, 0, (IntPtr)m, DPLength);
                }
            }
        }

        public byte[] DQ
        {
            get
            {
                byte[] bytes = new byte[DQLength];
                fixed (byte* m = dq)
                {
                    Marshal.Copy((IntPtr)m, bytes, 0, DQLength);
                }
                return bytes;
            }
            set
            {
                if (value.Length != DQLength) throw new ArgumentException(nameof(value) + " must be exaclty " + DQLength + " bytes (not " + value.Length + ")");
                fixed (byte* m = dq)
                {
                    Marshal.Copy(value, 0, (IntPtr)m, DQLength);
                }
            }
        }

        public byte[] InverseQ
        {
            get
            {
                byte[] bytes = new byte[InverseQLength];
                fixed (byte* m = inverseQ)
                {
                    Marshal.Copy((IntPtr)m, bytes, 0, InverseQLength);
                }
                return bytes;
            }
            set
            {
                if (value.Length != InverseQLength) throw new ArgumentException(nameof(value) + " must be exaclty " + InverseQLength + " bytes (not " + value.Length + ")");
                fixed (byte* m = inverseQ)
                {
                    Marshal.Copy(value, 0, (IntPtr)m, InverseQLength);
                }
            }
        }


        public void FromCSP(RSACryptoServiceProvider rsa)
        {
            var param = rsa.ExportParameters(true);

            Modulus = param.Modulus;
            Exponent = param.Exponent;
            D = param.D;
            P = param.P;
            Q = param.Q;
            DP = param.DP;
            DQ = param.DQ;
            InverseQ = param.InverseQ;
        }

        public void ToCSP(RSACryptoServiceProvider rsa)
        {
            var param = new RSAParameters();

            param.Modulus = Modulus;
            param.Exponent = Exponent;
            param.D = D;
            param.P = P;
            param.Q = Q;
            param.DP = DP;
            param.DQ = DQ;
            param.InverseQ = InverseQ;

            rsa.ImportParameters(param);
        }

        public static RSAPrivateKey Generate()
        {
            using (var rsa = new RSACryptoServiceProvider())
            {
                RSAPrivateKey ret = new RSAPrivateKey();
                ret.FromCSP(rsa);
                return ret;
            }
        }

        public static bool operator ==(RSAPrivateKey left, RSAPrivateKey right)
        {
            return left.Equals(right);
        }

        public static bool operator !=(RSAPrivateKey left, RSAPrivateKey right)
        {
            return !(left == right);
        }
    }

    public unsafe struct SHA256Hash
    {
        public const int Size = 32;
        internal fixed byte hash[Size];

        public override bool Equals(object obj)
        {
            return obj is SHA256Hash hash && hash.ToByteArrayUnmanaged().SequenceEqual(this.ToByteArrayUnmanaged());
        }

        public override int GetHashCode()
        {
            unchecked
            {
                return (int)DamienG.Security.Cryptography.Crc32.Compute(this.ToByteArrayUnmanaged());
            }
        }

        public byte[] Hash
        {
            get
            {
                byte[] bytes = new byte[Size];
                fixed (byte* m = hash)
                {
                    Marshal.Copy((IntPtr)m, bytes, 0, Size);
                }
                return bytes;
            }
            set
            {
                if (value.Length != Size) throw new ArgumentException(nameof(value) + " must be exaclty " + Size + " bytes (not " + value.Length + ")");
                fixed (byte* m = hash)
                {
                    Marshal.Copy(value, 0, (IntPtr)m, Size);
                }
            }
        }

        public SHA256Hash(byte[] hash)
        {
            Hash = hash;
        }

        public static SHA256Hash Generate(byte[] data, SHA256 hasher = null)
        {
            if (hasher == null)
                using (var hash = SHA256.Create())
                    return Generate(data, hash);

            return (SHA256Hash)hasher.ComputeHash(data);
        }

        public static implicit operator byte[](SHA256Hash hash) => hash.ToByteArrayUnmanaged();
        public static explicit operator SHA256Hash(byte[] hash) => hash.Length == Size ?  hash.ToUnmanaged<SHA256Hash>() : throw new InvalidCastException("Hash must be exactly "+Size+" bytes");

        public static bool operator ==(SHA256Hash left, SHA256Hash right)
        {
            return left.Equals(right);
        }

        public static bool operator !=(SHA256Hash left, SHA256Hash right)
        {
            return !(left == right);
        }

        public override string ToString()
        {
            return BitConverter.ToString(this.ToByteArrayUnmanaged()).Replace("-", "").ToLower();
        }
    }
}

namespace Tools
{
    using Tools.Crypto;

    public class IncrementalHashStream : Stream
    {
        private long written = 0;
        private IncrementalHash hasher;
        public IncrementalHashStream(HashAlgorithmName algo)
        {
            hasher = IncrementalHash.CreateHash(algo);
        }

        private byte[] hash;
        public SHA256Hash Hash => new SHA256Hash() { Hash = hash ?? throw new InvalidOperationException("Hash has not been computed. (Call Flush())") };
        public SHA256Hash HashAndReset()
        {
            Flush();
            return Hash;
        }

        public override bool CanRead => false;

        public override bool CanSeek => false;

        public override bool CanWrite => true;

        public override long Length => written;

        public override long Position { get => throw new NotSupportedException(); set => throw new NotSupportedException(); }

        /// <summary>
        /// Compute the hash of all written data, store it in <see cref="Hash"/>, and reset the algorithm.
        /// </summary>
        public override void Flush()
        {
            lock (mutex)
            {
                hash = hasher.GetHashAndReset();
            }
            written = 0;
        }

        public override int Read(byte[] buffer, int offset, int count)
        {
            throw new NotSupportedException();
        }

        public override long Seek(long offset, SeekOrigin origin)
        {
            throw new NotSupportedException();
        }

        public override void SetLength(long value)
        {
            throw new NotSupportedException();
        }

        private readonly object mutex = new object();
        public override void Write(byte[] buffer, int offset, int count)
        {
            var dt = new ArraySegment<byte>(buffer, offset, count).Array;
            lock (mutex)
            {
                hasher.AppendData(dt);
                written += dt.Length;
            }
        }

        public SHA256Hash GetFinalHash()
        {
            if (hasher == null) throw new ObjectDisposedException(this.ToString());
            Flush();
            Dispose();
            return Hash;
        }

        protected override void Dispose(bool disposing)
        {
            base.Dispose(disposing);
            if (disposing) hasher?.Dispose();
            hasher = null;
        }

    }
    public static class ExCryptoTools
    {
        public static SHA256Hash SHA256Hash(this byte[] data)
        {
            using (var sha = SHA256.Create())
                return (SHA256Hash)sha.ComputeHash(data);
        }

        public static SHA256Hash SHA256Hash(this Stream str)
        {
            using (var sha = SHA256.Create())
                return (SHA256Hash)sha.ComputeHash(str);
        }
    }
}