#ifndef CPP_PROTECTED_STRING_H_
#define CPP_PROTECTED_STRING_H_

#include <WinCryptEx.h>

class CProtectedPin
{
public:
    CProtectedPin()
    {
        m_isSetted = false;
    }

    bool IsPresent()
    {
        return m_isSetted;
    }

#ifdef UNIX
    HRESULT ProtectPin(CAtlStringA& strToProtect)
    {
        if (m_isSetted)
        {
            m_isSetted = false;
            m_strProtected.Empty();
        }
        m_strProtected = strToProtect;
        m_isSetted = true;
        SecureZeroString(strToProtect);
        return S_OK;
    }

    HRESULT SetPinToCryptProv(CCryptProv& cryptProv, DWORD dwPinFlag)
    {
        if (!m_isSetted)
        {
            return OLE_E_BLANK;
        }

        return cryptProv.SetParam(dwPinFlag,
            reinterpret_cast<BYTE*>(const_cast<char*>(m_strProtected.GetString())), 0);
    }
#else // !UNIX
    HRESULT ProtectPin(CAtlStringA& strToProtect)
    {
        if (m_isSetted)
        {
            m_isSetted = false;
            m_strProtected.Empty();
        }
        DATA_BLOB dataBlob;
        DATA_BLOB dataOutBlob;
        dataBlob.pbData = reinterpret_cast<BYTE*>(const_cast<char*>(strToProtect.GetString()));
        //Завершающий ноль нужен обязательно
        dataBlob.cbData = (strToProtect.GetLength() + 1) * sizeof(char);

        if (!CryptProtectData(&dataBlob, L"PIN-code", NULL, NULL, NULL,
            0, &dataOutBlob))
        {
            return ATL::AtlHresultFromLastError();
        }

        if (!dataOutBlob.pbData)
        {
	    return E_FAIL;
        }

        // TODO: логичнее сделать dataOutBlob приватным членом и хранить его, а не копировать в CAtlStringA
        // и когда на UNIX появится CryptProtectData так и стоит сделать.
        // пока что копируем в строку, чтобы была уверенность в освобождении памяти
        m_strProtected = CAtlStringA(reinterpret_cast<char*>(dataOutBlob.pbData),
            dataOutBlob.cbData);
        LocalFree(dataOutBlob.pbData);

        SecureZeroString(strToProtect);
        m_isSetted = true;
        return S_OK;
    }

    HRESULT SetPinToCryptProv(CCryptProv& cryptProv, DWORD dwPinFlag)
    {
        if (!m_isSetted)
        {
            return OLE_E_BLANK;
        }
        LPWSTR pDescrOut =  NULL;
        DATA_BLOB dataOutBlob;
        DATA_BLOB dataBlob;
        dataBlob.pbData = reinterpret_cast<BYTE*>(const_cast<char*>(m_strProtected.GetString()));
        dataBlob.cbData = m_strProtected.GetLength()*sizeof(char);

        if (!CryptUnprotectData(&dataBlob, &pDescrOut, NULL, NULL, NULL, 0,
            &dataOutBlob))
        {
            return ATL::AtlHresultFromLastError();
        }

        HRESULT hr = cryptProv.SetParam(dwPinFlag, dataOutBlob.pbData, 0);
            
        LocalFree(dataOutBlob.pbData);
        LocalFree(pDescrOut);

        return hr;
    }
#endif // !UNIX
private:
    DISALLOW_COPY_AND_ASSIGN(CProtectedPin);
    bool m_isSetted;
    CAtlStringA m_strProtected;
};

#endif // CPP_PROTECTED_STRING_H_
