sftp/Test.cpp

#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::SSHConnection));
	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_pSFTP->LastStatusCode);	
	Log(_T("LastStatusMessage = \"%s\"."), (LPCTSTR)m_pSFTP->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->File = 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: %u"), (LPCTSTR)m_pConnection->Host, m_pConnection->Port);		
		m_pConnection->Connect();
		
		Log(_T("%s"), (LPCTSTR)m_pConnection->ServerState->RemoteId);

		SFTPTest();

		// Disconnect
		Log(_T("Disconnect"));		
		m_pConnection->Disconnect();

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

void CTest::SFTPTest()
{
	// create SFTPConnection
	m_pSFTP = m_pConnection->CreateSFTPConnection();

	m_pSFTP->Connect();

	Log(_T("SFTP channel successfully opened."));

	// get current folder
	_bstr_t bstrRealPath = L".";
	Log(_T("RealPath \"%s\""), (LPCTSTR)bstrRealPath);

	_bstr_t bstrCurrentFolder = m_pSFTP->RealPath(bstrRealPath);
	Log(_T("Home Folder = %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);		
	
	sfFTPLib::IFTPItemsPtr pItems = m_pSFTP->ReadDirectory(bstrCurrentFolder);
	ATLENSURE(pItems);

	int nCount = pItems->Count;
	Log(_T("Count = %d"), nCount);		

	// Enum
	if(nCount > 0)
	{
		IEnumVARIANTPtr pEnum = pItems->_NewEnum;
		ATLENSURE(pEnum);

		ULONG CeltFetched;
		CComVariant variant;
		while(pEnum->Next(1, &variant, &CeltFetched) == S_OK)
		{
			if(variant.vt == VT_DISPATCH
				|| variant.vt == VT_UNKNOWN)
			{
				sfFTPLib::IFTPItemPtr 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::ftpItemAttributeModifyTime) == VARIANT_TRUE)
				{
					CString strTime;
					if(FILETIMEToISO8601(pSFTPItem->ModifyTime, 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);							
	m_pSFTP->MakeDirectory(_bstr_t(MakeDirectory));
	Log(_T("Directory \"%s\" created."), (LPCTSTR)MakeDirectory);

	// Rename
	CUnixPath RenameFrom = MakeDirectory;		
	CUnixPath RenameTo = (LPCTSTR)bstrCurrentFolder;
	RenameTo.Append(_T("testfolder2"));
	Log(_T("Rename \"%s\" to \"%s\""), (LPCTSTR)RenameFrom, (LPCTSTR)RenameTo);		
	m_pSFTP->Rename(_bstr_t(RenameFrom), _bstr_t(RenameTo), 0);

	// RemoveDirectory
	CUnixPath RemoveDirectory = RenameTo;
	Log(_T("RemoveDirectory \"%s\""), (LPCTSTR)RemoveDirectory);							
	m_pSFTP->RemoveDirectory(_bstr_t(RemoveDirectory));
	Log(_T("Directory \"%s\" removed."), (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);		
		m_pSFTP->UploadFile(CComVariant(pMemFile.GetInterfacePtr()), bstrUploadFile, 0,0);
		Log(_T("File sucessfully uploaded."));

		// Stat. Stat doesn't follow symbolic links.
		_bstr_t bstrStat = bstrUploadFile;
		Log(_T("Stat \"%s\""), (LPCTSTR)bstrStat);		
		CComPtr<sfFTPLib::IFTPItem> pItem = m_pSFTP->Stat(bstrStat, sfFTPLib::ftpSFTPItemAttributeSize);
		if(pItem->IsValidAttribute(sfFTPLib::ftpItemAttributeSize))
			Log(_T("File Size = %I64u."), pItem->Size);

		// 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);		
			m_pSFTP->DownloadFileEx(bstrDownloadFile, CComVariant(pDownloadMemFile.GetInterfacePtr()), 0, 0);
			Log(_T("File sucessfully downloaded."));
		}

		// 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);		
		m_pSFTP->DownloadFile(bstrDownloadFile, CComVariant(bstrDownloadLocalFile.GetBSTR()), 0, 0);
		Log(_T("File sucessfully downloaded."));		
	}

	Log(_T("Closing channel."));
	m_pSFTP->Disconnect();
}

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;
}