/******************************************************************************/
/* CodeSign, by LoRd_MuldeR <MuldeR2@GMX.de>                                  */
/* This work has been released under the CC0 1.0 Universal license!           */
/******************************************************************************/

using System;
using System.Runtime.InteropServices;
using System.Security;
using System.Text;

namespace PasswordDialog
{
    public class SecureWrapper : IDisposable
    {
        private static readonly char[] EMPTY_CHARS = new char[0];
        private static readonly byte[] EMPTY_BYTES = new byte[0];

        private readonly SecureString m_string;
        private readonly Object m_lock = new Object();

        private GCHandle m_charBuffer;
        private GCHandle m_byteBuffer;

        private volatile bool m_disposed = false;

        //-------------------------------------------------------------------
        // Constructor
        //-------------------------------------------------------------------

        public SecureWrapper(SecureString secureString)
        {
            if(ReferenceEquals(m_string = secureString, null))
            {
                throw new ArgumentNullException("SecureString must not be null!");
            }
        }

        ~SecureWrapper()
        {
            Dispose();
        }

        //-------------------------------------------------------------------
        // Public Interface
        //-------------------------------------------------------------------

        public char[] CharBuffer
        {
            get
            {
                lock(m_lock)
                {
                    if(m_disposed)
                    {
                        throw new ObjectDisposedException("Already disposed!");
                    }
                    if(m_string.Length > 0)
                    {
                        if (!m_charBuffer.IsAllocated)
                        {
                            m_charBuffer = AllocateBuffer<char>(m_string.Length);
                            try
                            {
                                CopyToBuffer((char[])m_charBuffer.Target, m_string);
                            }
                            catch
                            {
                                ClearBuffer(m_charBuffer);
                                throw; /*re-throw after clean-up!*/
                            }
                        }
                        return (char[])m_charBuffer.Target;
                    }
                    return EMPTY_CHARS;
                }
            }
        }

        public byte[] ByteBuffer
        {
            get
            {
                lock (m_lock)
                {
                    if(m_disposed)
                    {
                        throw new ObjectDisposedException("Already disposed!");
                    }
                    if (m_string.Length > 0)
                    {
                        if (!m_byteBuffer.IsAllocated)
                        {
                            UTF8Encoding encoding = new UTF8Encoding(false);
                            char[] chars = CharBuffer;
                            m_byteBuffer = AllocateBuffer<byte>(encoding.GetByteCount(chars));
                            try
                            {
                                encoding.GetBytes(chars, 0, chars.Length, (byte[])m_byteBuffer.Target, 0);
                            }
                            catch
                            {
                                ClearBuffer(m_byteBuffer);
                                throw; /*re-throw after clean-up!*/
                            }
                        }
                        return (byte[])m_byteBuffer.Target;
                    }
                    return EMPTY_BYTES;
                }
            }
        }
        
        public void Dispose()
        {
            lock (m_lock)
            {
                if (!m_disposed)
                {
                    m_disposed = true;
                    try
                    {
                        ClearBuffer(m_byteBuffer);
                    }
                    catch { }
                    try
                    {
                        ClearBuffer(m_charBuffer);
                    }
                    catch { }
                }
            }
        }

        //-------------------------------------------------------------------
        // Internal Methods
        //-------------------------------------------------------------------

        private static GCHandle AllocateBuffer<T>(int length)
        {
            T[] buffer = new T[length];
            return GCHandle.Alloc(buffer, GCHandleType.Pinned);
        }

        private static void CopyToBuffer(char[] buffer, SecureString secStr)
        {
            IntPtr valuePtr = IntPtr.Zero;
            try
            {
                valuePtr = Marshal.SecureStringToGlobalAllocUnicode(secStr);
                Marshal.Copy(valuePtr, buffer, 0, Math.Min(secStr.Length, buffer.Length));
            }
            finally
            {
                Marshal.ZeroFreeGlobalAllocUnicode(valuePtr);
            }
        }

        private static void ClearBuffer(GCHandle handle)
        {
            if(handle.IsAllocated)
            {
                try
                {
                    if(handle.Target is Array)
                    {
                        Array array = (Array)handle.Target;
                        Array.Clear(array, 0, array.Length);
                    }
                    else if(handle.Target is String)
                    {
                        char[] zeros = new char[((String)handle.Target).Length];
                        Marshal.Copy(zeros, 0, handle.AddrOfPinnedObject(), zeros.Length);
                    }
                    else
                    {
                        throw new ArgumentException("Unsupported buffer type!");
                    }
                }
                finally
                {
                    handle.Free();
                }
            }
        }
    }
}