#include "StdAfx.h"
#include "Test.h"
#include "UnixPath.h"

CTest::CTest(void)
{
}

CTest::~CTest(void)
{
}


HRESULT CTest::Initialize()
{
    HRESULT hr = m_pConnection.CreateInstance(__uuidof(sfFTPLib::SFTPConnection));
    if(SUCCEEDED(hr))
    {
        hr = S_OK;
    }
    else
    {
        Log(_T("Failed to create SFTPConnection instance. hr=0x%x.\nTry to register sfFTPLib.dll again."), hr);
        ATLASSERT(0);
    }
    return hr;
}

HRESULT CTest::Uninitialize()
{
    m_pConnection.Release();
    return S_OK;
}

void CTest::ReportLastStatus()
{
    Log(_T("LastStatusCode = %d."), m_pConnection->LastStatusCode);
    Log(_T("LastStatusMessage = \"%s\"."), (LPCTSTR)m_pConnection->LastStatusMessage); 
}

// Purpose: Converts FILETIME to ISO8601 string
HRESULT CTest::FILETIMEToISO8601(const FILETIME& ft, CString &retval)
{
    HRESULT hr = E_FAIL;
    SYSTEMTIME st;
    if(::FileTimeToSystemTime(&ft, &st))
    {
        retval.Format(_T("%04d-%02d-%02dT%02d:%02d:%02dZ"), st.wYear, st.wMonth, st.wDay, st.wHour, st.wMinute, st.wSecond);
        hr = S_OK;
    }
    return hr;
}

HRESULT CTest::Run()
{
    ATLENSURE_RETURN_HR(m_pConnection, E_UNEXPECTED);

    HRESULT hr = E_FAIL;
    try
    {
        //m_pConnection->Async = VARIANT_FALSE;
        m_pConnection->Host = L"localhost";
        m_pConnection->Port = 22;
        m_pConnection->Username = L"user";
        m_pConnection->Password = L"password";
        m_pConnection->LogFile = L"Win32.log";

        // Authentication
        // Notes:
        // - Titan FTP Server: If "publickey" authentication fails the server disconnects without accepting any further methods. e.g. password
        CComSafeArray<VARIANT> authentications(2);
        //authentications.SetAt(0, CComVariant(sfFTPLib::ftpSSHAuthenticationNone));
        authentications.SetAt(0, CComVariant(sfFTPLib::ftpSSHAuthenticationPassword));
        authentications.SetAt(1, CComVariant(sfFTPLib::ftpSSHAuthenticationPublicKey));
        //m_pConnection->PutAuthentications(&CComVariant(authentications));

        // Disable Compression
        // Uncomment to disable compression
        CComSafeArray<VARIANT> compressions(2);
        compressions.SetAt(0, CComVariant(sfFTPLib::ftpSSHCompressionzlib));
        compressions.SetAt(0, CComVariant(sfFTPLib::ftpSSHCompressionNone));
        //m_pConnection->PutCompressions(&CComVariant(compressions));

        // Limit Encryptions
        CComSafeArray<VARIANT> encryptions(4);
        encryptions.SetAt(0, CComVariant(sfFTPLib::ftpEncryptionAES256));
        encryptions.SetAt(1, CComVariant(sfFTPLib::ftpEncryptionAES192));
        encryptions.SetAt(2, CComVariant(sfFTPLib::ftpEncryption3DES));
        encryptions.SetAt(3, CComVariant(sfFTPLib::ftpEncryptionAES128));
        //m_pConnection->PutEncryptions(&CComVariant(encryptions));

        // Limit KeyExchange Algorithms
        CComSafeArray<VARIANT> keyexchanges(2);
        keyexchanges.SetAt(0, CComVariant(sfFTPLib::ftpKeyExchangeDiffieHellmanGroup14SHA1));
        keyexchanges.SetAt(1, CComVariant(sfFTPLib::ftpKeyExchangeDiffieHellmanGroup1SHA1));
        //m_pConnection->PutKeyExchanges(&CComVariant(keyexchanges));

        sfFTPLib::IKeyManagerPtr pKeyManager;
        if(SUCCEEDED(pKeyManager.CreateInstance(__uuidof(sfFTPLib::KeyManager))))
        {
            // Uncomment to generate new key
            _bstr_t bstrFilePrivate(L"Identity");
            _bstr_t bstrFilePublic(L"Identity.pub");
            _bstr_t bstrPassword(L"");

#if 0
            // Uncomment to create private key
            // For VShell copy public key (Identity.pub) to user's folder: C:\Program Files\VShell\PublicKey\<user>
            sfFTPLib::IRSAKeyPtr pRSA;
            if(SUCCEEDED(pRSA.CreateInstance(__uuidof(sfFTPLib::RSAKey))))
            {
                // Generate 1024-bit RSA key
                pRSA->Generate(1024);
                // Save private key in PKCS12 format (.p12)
                pKeyManager->SaveFile(sfFTPLib::ftpKeyFileFormatPKCS12, pRSA, sfFTPLib::ftpKeyTypePrivateKey, bstrFilePrivate, bstrPassword);
                // Save public key. Password is ignored.
                pKeyManager->SaveFile(sfFTPLib::ftpKeyFileFormatSSH, pRSA, sfFTPLib::ftpKeyTypePublicKey, bstrFilePublic, bstrPassword);
                // Save public key (for OpenSSH only). Password is ignored.
                //pKeyManager->SaveFile(sfFTPLib::ftpKeyFileFormatOpenSSH, pRSA, sfFTPLib::ftpKeyTypePublicKey, bstrFilePublic, bstrPassword);
            }
#endif
            Log(_T("Loading private key \"%s\"."), (LPCTSTR)bstrFilePrivate);      
            sfFTPLib::IKeyPtr pKey;
            if(pKeyManager->raw_LoadFile(bstrFilePrivate, bstrPassword, &pKey) == S_OK)
            {
                if(pKey->Type == sfFTPLib::ftpKeyTypePrivateKey)
                {
                    m_pConnection->PrivateKey = pKey;
                    Log(_T("Private key sucessfully loaded from \"%s\"."), (LPCTSTR)bstrFilePrivate);  
                }
            }
            else
            {
                Log(_T("Failed to load key."));    
            }
        }
        else
        {
            ATLASSERT(0);
        }

        Log(_T("Connecting to %s Port: %d"), (LPCTSTR)m_pConnection->Host, m_pConnection->Port);       
        if(m_pConnection->Connect() == sfFTPLib::ftpErrorSuccess)
        {
            Log(_T("%s"), (LPCTSTR)m_pConnection->RemoteId);

            // get current folder
            _bstr_t bstrRealPath = L".";
            Log(_T("RealPath \"%s\""), (LPCTSTR)bstrRealPath);     
            if(m_pConnection->RealPath(bstrRealPath) == sfFTPLib::ftpErrorSuccess)
            {
                _bstr_t bstrCurrentFolder = m_pConnection->LastPath;
                Log(_T("LastPath = %s"), (LPCTSTR)bstrCurrentFolder);

                // overriding CurrentFolder for debug purpose
                //bstrCurrentFolder = L"/c/archive";
                //Log(_T("Overriding current folder. \"%s\""), (LPCTSTR)bstrCurrentFolder);

                Log(_T("Reading Directory \"%s\""), (LPCTSTR)bstrCurrentFolder);       
                if(m_pConnection->ReadDirectory(bstrCurrentFolder) == sfFTPLib::ftpErrorSuccess)
                {
                    sfFTPLib::ISFTPItemsPtr pItems = m_pConnection->Items;
                    if(pItems)
                    {
                        int nCount = pItems->Count;
                        Log(_T("Count = %d"), nCount);     

                        // Enum
                        if(nCount > 0)
                        {
                            IEnumVARIANTPtr pEnum = pItems->_NewEnum;
                            if(pEnum)
                            {
                                ULONG CeltFetched;
                                CComVariant variant;
                                while(pEnum->Next(1, &variant, &CeltFetched) == S_OK)
                                {
                                    if(variant.vt == VT_DISPATCH
                                        || variant.vt == VT_UNKNOWN)
                                    {
                                        sfFTPLib::ISFTPItemPtr pSFTPItem = variant.pdispVal;
                                        ATLENSURE(pSFTPItem);
                                       
                                        // TODO: Check for valid attributes (IsValidAttribute())
                                        CString str;
                                        str.Format(_T("Type=0x%x; Name=%s; Size=%d"), pSFTPItem->Type, (LPCTSTR)pSFTPItem->Name, pSFTPItem->Size);

                                        if(pSFTPItem->IsValidAttribute(sfFTPLib::ftpSFTPItemAttributeModifyTime) == VARIANT_TRUE)
                                        {
                                            CString strTime;
                                            if(FILETIMEToISO8601(pSFTPItem->ModifyTimeAsFileTime, strTime) == S_OK)
                                            {
                                                str += _T("; ModifyTime=") + strTime;
                                            }
                                        }

                                        Log(str);                                      
                                    }
                                    // need to manually clear variant
                                    variant.Clear();
                                }
                            }
                        }
                    }

                    // MakeDirectory
                    CUnixPath MakeDirectory((LPCTSTR)bstrCurrentFolder);
                    MakeDirectory.Append(_T("testfolder"));
                    Log(_T("MakeDirectory \"%s\""), (LPCTSTR)MakeDirectory);                           
                    if(m_pConnection->MakeDirectory(_bstr_t(MakeDirectory)) == sfFTPLib::ftpErrorSuccess)
                    {
                        Log(_T("Directory \"%s\" created."), (LPCTSTR)MakeDirectory);
                    }
                    else
                    {

                        Log(_T("Failed to create directory \"%s\"."), (LPCTSTR)MakeDirectory);
                    }

                    // Rename
                    CUnixPath RenameFrom = MakeDirectory;      
                    CUnixPath RenameTo = (LPCTSTR)bstrCurrentFolder;
                    RenameTo.Append(_T("testfolder2"));
                    Log(_T("Rename \"%s\" to \"%s\""), (LPCTSTR)RenameFrom, (LPCTSTR)RenameTo);    
                    if(m_pConnection->Rename(_bstr_t(RenameFrom), _bstr_t(RenameTo), 0) == sfFTPLib::ftpErrorSuccess)
                    {
                        Log(_T("File sucessfully renamed."));

                    }
                    else
                    {
                        Log(_T("Rename failed."));
                        ReportLastStatus();
                    }

                    // RemoveDirectory
                    CUnixPath RemoveDirectory = RenameTo;
                    Log(_T("RemoveDirectory \"%s\""), (LPCTSTR)RemoveDirectory);                           
                    if(m_pConnection->RemoveDirectory(_bstr_t(RemoveDirectory)) == sfFTPLib::ftpErrorSuccess)
                    {
                        Log(_T("Directory \"%s\" removed."), (LPCTSTR)RemoveDirectory);
                    }
                    else
                    {

                        Log(_T("Failed to removedirectory \"%s\"."), (LPCTSTR)RemoveDirectory);
                    }

                    // Creating temporary memory file
                    IStreamPtr pMemFile;
                    DWORD dwSize = 1000 * 1024; // 1000 KiB
                    if(CreateMemFile(dwSize, 0, &pMemFile) == S_OK)
                    {
                        // Upload File
                        CUnixPath UploadFile((LPCTSTR)bstrCurrentFolder);
                        UploadFile.Append(_T("memfile"));
                        _bstr_t bstrUploadFile = (LPCTSTR)UploadFile;      

                        Log(_T("UploadFile to \"%s\""), (LPCTSTR)bstrUploadFile);      
                        if(m_pConnection->UploadFile(CComVariant(pMemFile.GetInterfacePtr()), bstrUploadFile, 0,0, 0,0) == sfFTPLib::ftpErrorSuccess)
                        {
                            Log(_T("File sucessfully uploaded."));

                        }
                        else
                        {
                            Log(_T("UploadFile failed."));
                            ReportLastStatus();
                        }

                        // Stat. Stat doesn't follow symbolic links.
                        _bstr_t bstrStat = bstrUploadFile;
                        Log(_T("Stat \"%s\""), (LPCTSTR)bstrStat);     
                        if(m_pConnection->Stat(bstrStat, sfFTPLib::ftpSFTPItemAttributeSize) == sfFTPLib::ftpErrorSuccess)
                        {
                            if(m_pConnection->LastItem->ValidAttributes & sfFTPLib::ftpSFTPItemAttributeSize)
                                Log(_T("File Size = %d."), m_pConnection->LastItem->Size);

                        }
                        else
                        {
                            Log(_T("Stat failed."));   
                            ReportLastStatus();
                        }

                        // DownloadFile to memory file
                        IStreamPtr pDownloadMemFile;
                        if(CreateMemFile(dwSize, 0, &pDownloadMemFile) == S_OK)
                        {
                            _bstr_t bstrDownloadFile = bstrUploadFile;

                                Log(_T("DownloadFile \"%s\" to memfile"), (LPCTSTR)bstrDownloadFile);      
                                if(m_pConnection->DownloadFile(bstrDownloadFile, CComVariant(pDownloadMemFile.GetInterfacePtr()), 0,0, 0,0) == sfFTPLib::ftpErrorSuccess)
                                {
                                    Log(_T("File sucessfully downloaded."));       
                                }
                                else
                                {
                                    Log(_T("DownloadFile failed."));
                                    ReportLastStatus();
                                }
                        }

                        // DownloadFile to physical file
                        _bstr_t bstrDownloadFile = bstrUploadFile;
                        TCHAR szCurrentDirectory[MAX_PATH] = {0};
                        ::GetCurrentDirectory(ARRAYSIZE(szCurrentDirectory), szCurrentDirectory);
                        ::PathAppend(szCurrentDirectory, _T("Download"));
                        ::SHCreateDirectoryEx(NULL, szCurrentDirectory, NULL);
                        ::PathAppend(szCurrentDirectory, _T("memfile"));
                        _bstr_t bstrDownloadLocalFile = szCurrentDirectory;
                        Log(_T("DownloadFile \"%s\" to \"%s\""), (LPCTSTR)bstrDownloadFile, (LPCTSTR)bstrDownloadLocalFile);       
                        if(m_pConnection->DownloadFile(bstrDownloadFile, CComVariant(bstrDownloadLocalFile.GetBSTR()), 0,0, 0,0) == sfFTPLib::ftpErrorSuccess)
                        {
                            Log(_T("File sucessfully downloaded."));       
                        }
                        else
                        {
                            Log(_T("DownloadFile failed."));
                            ReportLastStatus();
                        }
                    }
                }
                else
                {
                    Log(_T("ReadDirectory failed."));
                    ReportLastStatus();
                }
            }
            else
            {
                Log(_T("RealPath failed."));
                ReportLastStatus();
            }

            // Disconnect
            Log(_T("Disconnect"));     
            if(m_pConnection->Disconnect() == sfFTPLib::ftpErrorSuccess)
            {
                Log(_T("Sucessfully disconnected."));      
            }
            else
            {
                Log(_T("Disconnect failed."));
                ReportLastStatus();
            }
        }
        else
        {
            Log(_T("Connect failed."));
        }

        hr = S_OK;
    }
    catch (_com_error &e)
    {
        e; 
        Log(_T("_com_error hr=0x%x"), e.Error());
        ATLASSERT(0);
    }
    return hr;
}

void CTest::Log(LPCTSTR pszFormat, ...)
{
    // max limit of log message set to 4096. Increase if message gets cut.
    const int LOG_EVENT_MSG_SIZE = 4096;

    TCHAR chMsg[LOG_EVENT_MSG_SIZE];
    va_list pArg;

    va_start(pArg, pszFormat);
#if _SECURE_ATL
    _vsntprintf_s(chMsg, LOG_EVENT_MSG_SIZE, LOG_EVENT_MSG_SIZE-1, pszFormat, pArg);
#else
    _vsntprintf(chMsg, LOG_EVENT_MSG_SIZE, pszFormat, pArg);
#endif

    CString strMsg(chMsg);
    strMsg += _T("\n");
    ATLTRACE(strMsg);
    _tprintf(strMsg);
}

// Purpose: Creates memory file with Global Memory (GlobalAlloc)
// nFillMethod: 0: zero data, 1: fill with 0-255
HRESULT CTest::CreateMemFile(DWORD nSize, int nFillMethod, IStream **retval)
{
    ATLENSURE_RETURN_HR(retval, E_INVALIDARG);
    *retval = NULL;

    HRESULT hr = E_FAIL;
    HGLOBAL hMem = ::GlobalAlloc(GMEM_MOVEABLE | GMEM_ZEROINIT, static_cast<SIZE_T>(nSize));
    if (!hMem)
        return E_OUTOFMEMORY;

    BYTE *pImage = reinterpret_cast<BYTE*>(::GlobalLock(hMem));
    if(pImage)
    {
        if(nFillMethod == 1)
        {
            // fill with 0-255
            for(DWORD i=0; i<nSize; i++)
                pImage[i] = static_cast<BYTE>(i);
        }
        ::GlobalUnlock(hMem);

        // Create Stream from hMem. Automatically release hMem
        hr = ::CreateStreamOnHGlobal(hMem, TRUE, retval);
    }
    return hr;
}