sftp/Test.cpp

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

_Check_return_ HRESULT CTest::Initialize()
{
	HRESULT hr = m_Connection.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;
}

_Check_return_ HRESULT CTest::Uninitialize()
{
	m_Connection.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
_Check_return_ HRESULT CTest::FILETIMEToISO8601(const FILETIME& ft, CString &retval)
{
	HRESULT hr = E_FAIL;
	SYSTEMTIME st;
	if(::FileTimeToSystemTime(&ft, &st))
	{
		retval.Format(_T("%04hu-%02hu-%02huT%02hu:%02hu:%02huZ"), st.wYear, st.wMonth, st.wDay, st.wHour, st.wMinute, st.wSecond);
		hr = S_OK;
	}
	return hr;
}

_Check_return_ HRESULT CTest::Run()
{
	HRESULT hr = E_FAIL;
	try
	{
		//m_Connection->Async = VARIANT_FALSE;
		m_Connection->Host = L"localhost";
		m_Connection->Port = 22;
		m_Connection->Username = L"user";
		m_Connection->Password = L"password";

		CComPtr<sfFTPLib::IFileLogger> fileLogger;
		ATLENSURE_SUCCEEDED(fileLogger.CoCreateInstance(__uuidof(sfFTPLib::FileLogger)));
		fileLogger->Putfile(L"ssh.log");

		CComQIPtr<sfFTPLib::ILogger> logger(fileLogger);
		m_Connection->Logger = logger;

		// 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_Connection->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_Connection->PutCompressions(&CComVariant(compressions));

		// Limit Encryptions
		CComSafeArray<VARIANT> encryptions(4);
		encryptions.SetAt(0, CComVariant(sfFTPLib::ftpEncryptionAES256CTR));
		encryptions.SetAt(1, CComVariant(sfFTPLib::ftpEncryptionAES192CTR));
		encryptions.SetAt(2, CComVariant(sfFTPLib::ftpEncryptionAES128CTR));
		encryptions.SetAt(3, CComVariant(sfFTPLib::ftpEncryption3DES));
		//m_Connection->PutEncryptions(&CComVariant(encryptions));

		// Limit KeyExchange Algorithms
		CComSafeArray<VARIANT> keyexchanges(2);
		keyexchanges.SetAt(0, CComVariant(sfFTPLib::ftpKeyExchangeDiffieHellmanGroup14SHA1));
		keyexchanges.SetAt(1, CComVariant(sfFTPLib::ftpKeyExchangeDiffieHellmanGroup1SHA1));
		//m_Connection->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);
			CComPtr<sfFTPLib::IKey> pKey;
			if(pKeyManager->raw_LoadFile(bstrFilePrivate, bstrPassword, &pKey) == S_OK)
			{
				if(pKey->Type == sfFTPLib::ftpKeyTypePrivateKey)
				{
					m_Connection->PrivateKey = pKey;
					Log(_T("Private key successfully loaded from \"%s\"."), (LPCTSTR)bstrFilePrivate);
				}
			}
			else
			{
				Log(_T("Failed to load key."));
			}
		}
		else
		{
			ATLASSERT(0);
		}

		Log(_T("Connecting to %s Port: %u"), (LPCTSTR)m_Connection->Host, m_Connection->Port);
		m_Connection->Connect();

		Log(_T("%s"), (LPCTSTR)m_Connection->ServerState->RemoteId);

		SFTPTest();

		// Disconnect
		Log(_T("Disconnect"));
		m_Connection->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()
{
	m_pSFTP = m_Connection->CreateSFTPConnection();

	CComPtr<sfFTPLib::IFileLogger> fileLogger;
	ATLENSURE_SUCCEEDED(fileLogger.CoCreateInstance(__uuidof(sfFTPLib::FileLogger)));
	fileLogger->Putfile(L"sftp.log");

	CComQIPtr<sfFTPLib::ILogger> logger(fileLogger);
	m_pSFTP->Logger = logger;

	m_pSFTP->Connect();

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

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

	_bstr_t currentFolder = m_pSFTP->RealPath(realPath);
	Log(_T("Home Folder = %s"), (LPCTSTR)currentFolder);

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

	Log(_T("Reading Directory \"%s\""), (LPCTSTR)currentFolder);

	sfFTPLib::IFTPItemsPtr items = m_pSFTP->ReadDirectory(currentFolder);

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

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

		ULONG CeltFetched;
		CComVariant variant;
		while(pEnum->Next(1, &variant, &CeltFetched) == S_OK)
		{
			if(variant.vt == VT_DISPATCH || variant.vt == VT_UNKNOWN)
			{
				CComQIPtr<sfFTPLib::IFTPItem> pSFTPItem = variant.pdispVal;

				// 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::ftpFTPItemAttributeModifyTime) == 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)currentFolder);
	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)currentFolder;
	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
	CComPtr<IStream> pMemFile;
	static const DWORD dwSize = 1000 * 1024; // 1000 KiB
	if(CreateMemFile(dwSize, 0, &pMemFile) == S_OK)
	{
		// Upload File
		CUnixPath UploadFile((LPCTSTR)currentFolder);
		UploadFile.Append(_T("memfile"));
		_bstr_t bstrUploadFile = (LPCTSTR)UploadFile;

		Log(_T("UploadFile to \"%s\""), (LPCTSTR)bstrUploadFile);
		m_pSFTP->UploadFile(CComVariant(pMemFile.p), bstrUploadFile, sfFTPLib::ftpDataTransferTypeImage, 0,0);
		Log(_T("File successfully 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::ftpFTPItemAttributeSize);
		if (pItem->IsValidAttribute(sfFTPLib::ftpFTPItemAttributeSize))
		{
			Log(_T("File Size = %I64u."), pItem->Size);
		}
		// DownloadFile to memory file
		CComPtr<IStream> 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.p), sfFTPLib::ftpDataTransferTypeImage, 0, 0, sfFTPLib::ftpDownloadFlagReadBeyondEnd, nullptr);
			Log(_T("File successfully downloaded."));
		}

		// DownloadFile to physical file
		_bstr_t bstrDownloadFile = bstrUploadFile;
		TCHAR szCurrentDirectory[MAX_PATH] = {};
		::GetCurrentDirectory(ARRAYSIZE(szCurrentDirectory), szCurrentDirectory);
		::PathAppend(szCurrentDirectory, _T("Download"));
		::SHCreateDirectoryEx(NULL, szCurrentDirectory, NULL);
		::PathAppend(szCurrentDirectory, _T("memfile"));
		_bstr_t downloadLocalFile = szCurrentDirectory;
		Log(_T("DownloadFile \"%s\" to \"%s\""), (LPCTSTR)bstrDownloadFile, (LPCTSTR)downloadLocalFile);
		m_pSFTP->DownloadFile(bstrDownloadFile, CComVariant(downloadLocalFile.GetBSTR()), sfFTPLib::ftpDataTransferTypeImage, 0, 0);
		Log(_T("File successfully downloaded."));
	}

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

void CTest::Log(_In_z_ PCTSTR 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
_Check_return_ HRESULT CTest::CreateMemFile(DWORD nSize, int fillMethod, _COM_Outptr_ IStream **retval)
{
	*retval = nullptr;

	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(fillMethod == 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;
}