sftp/Test.cpp

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

#include <format>

_Check_return_ HRESULT CTest::Initialize()
{
	const auto hr{ m_Connection.CoCreateInstance(__uuidof(sfFTPLib::SSHConnection)) };
	if (SUCCEEDED(hr))
	{
		return S_OK;
	}
	Log(L"Failed to create SFTPConnection instance. hr=0x%x.\nTry to register sfFTPLib.dll again.", hr);
	ATLASSERT(false);

	ATL::CComPtr<sfFTPLib::IGlobal> global;
	ATLENSURE_SUCCEEDED(global.CoCreateInstance(__uuidof(sfFTPLib::Global)));
	// If LoadLicense is not called a trial license is automatically obtained from the activation server. The FTP Library uses WinHTTP to access
	// the activation server at www.smartftp.com (TLS, port 443). Ensure that your application is not blocked by any firewall.
	// TODO: insert the provided serial after the purchase of a license
	//ATLENSURE_SUCCEEDED(global->LoadLicense(ATL::CComBSTR(L"XXXXXXXX-XXXX-XXXX-XXXX-XXXXXXXXXXXX")));

	return hr;
}

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

void CTest::ReportLastStatus()
{
	sfFTPLib::SFTPStatus statusCode;
	ATLENSURE_SUCCEEDED(m_SFTP->get_LastStatusCode(&statusCode));
	Log(L"LastStatusCode = %d.", statusCode);
	ATL::CComBSTR statusMessage;
	ATLENSURE_SUCCEEDED(m_SFTP->get_LastStatusMessage(&statusMessage));
	Log(L"LastStatusMessage = \"%s\".", (PCWSTR)statusMessage);
}

// Purpose: Converts FILETIME to ISO8601 string
_Check_return_ HRESULT CTest::FILETIMEToISO8601(const FILETIME& ft, std::wstring& retval)
{
	SYSTEMTIME st;
	if (::FileTimeToSystemTime(&ft, &st))
	{
		retval.resize(255);
		retval.resize(std::swprintf(retval.data(), retval.size() + 1, L"%04hu-%02hu-%02huT%02hu:%02hu:%02huZ", st.wYear, st.wMonth, st.wDay, st.wHour, st.wMinute, st.wSecond));
		return S_OK;
	}
	return E_FAIL;
}

void CTest::Run()
{
	const wchar_t host []{ L"localhost" };
	const long port{ 22 };
	const wchar_t username []{ L"user" };
	const wchar_t password []{ L"pass" };

	//m_Connection->Async = VARIANT_FALSE;
	ATLENSURE_SUCCEEDED(m_Connection->put_Host(ATL::CComBSTR(host)));
	ATLENSURE_SUCCEEDED(m_Connection->put_Port(port));
	ATLENSURE_SUCCEEDED(m_Connection->put_Username(ATL::CComBSTR(username)));
	ATLENSURE_SUCCEEDED(m_Connection->put_Password(ATL::CComBSTR(password)));

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

	ATL::CComQIPtr<sfFTPLib::ILogger> logger{ fileLogger };
	ATLENSURE_SUCCEEDED(m_Connection->put_Logger(logger));

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

	// Disable Compression
	// Uncomment to disable compression
	ATL::CComSafeArray<VARIANT> compressions(2);
	compressions.SetAt(0, ATL::CComVariant(sfFTPLib::ftpSSHCompressionzlibopenssh));
	compressions.SetAt(0, ATL::CComVariant(sfFTPLib::ftpSSHCompressionNone));
	//m_Connection->put_Compressions(&ATL::CComVariant(compressions));

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

	// Limit KeyExchange Algorithms
	ATL::CComSafeArray<VARIANT> keyexchanges(2);
	keyexchanges.SetAt(0, ATL::CComVariant(sfFTPLib::ftpKeyExchangeDiffieHellmanGroup14SHA1));
	//m_Connection->put_KeyExchanges(&ATL::CComVariant(keyexchanges));

	ATL::CComPtr<sfFTPLib::IKeyManager> keyManager;
	ATLENSURE_SUCCEEDED(keyManager.CoCreateInstance(__uuidof(sfFTPLib::KeyManager)));
	
	// Uncomment to generate new key
	ATL::CComBSTR bstrFilePrivate(L"Identity");
	ATL::CComBSTR bstrFilePublic(L"Identity.pub");
	ATL::CComBSTR bstrPassword(L"");

#if 0
	// Uncomment to create private key
	// -> broken
	// For VShell copy public key (Identity.pub) to user's folder: C:\Program Files\VShell\PublicKey\<user>
	ATL::CComPtr<sfFTPLib::OpenSSLKey> pRSA;
	ATLENSURE_SUCCEEDED(pRSA.CoCreateInstance(__uuidof(sfFTPLib::OpenSSLKey)));

	// Generate 2048-bit RSA key
	pRSA->Generate(2048);
	// Save private key in PKCS12 format (.p12)
	ATLENSURE_SUCCEEDED(keyManager->SaveFile(sfFTPLib::ftpKeyFileFormatPKCS12, pRSA, sfFTPLib::ftpKeyTypePrivateKey, bstrFilePrivate, bstrPassword));
	// Save public key. Password is ignored.
	ATLENSURE_SUCCEEDED(keyManager->SaveFile(sfFTPLib::ftpKeyFileFormatSSH, pRSA, sfFTPLib::ftpKeyTypePublicKey, bstrFilePublic, bstrPassword));
	// Save public key (for OpenSSH only). Password is ignored.
	//ATLENSURE_SUCCEEDED(keyManager->SaveFile(sfFTPLib::ftpKeyFileFormatOpenSSH, pRSA, sfFTPLib::ftpKeyTypePublicKey, bstrFilePublic, bstrPassword);
#endif

	Log(L"Loading private key \"%s\".", (PCWSTR) bstrFilePrivate);
	ATL::CComPtr<sfFTPLib::IKey> pKey;
	if (keyManager->LoadFile(bstrFilePrivate, bstrPassword, &pKey) == S_OK)
	{
		sfFTPLib::KeyType keyType;
		ATLENSURE_SUCCEEDED(pKey->get_Type(&keyType));
		if (keyType == sfFTPLib::ftpKeyTypePrivateKey)
		{
			m_Connection->put_PrivateKey(pKey);
			Log(L"Private key successfully loaded from \"%s\".", (PCWSTR) bstrFilePrivate);
		}
	}
	else
	{
		Log(L"Failed to load key.");
	}

	Log(L"Connecting to %s Port: %u", (PCWSTR)host, port);
	ATLENSURE_SUCCEEDED(m_Connection->Connect());

	ATL::CComPtr<sfFTPLib::ISSHServerState> serverState;
	ATLENSURE_SUCCEEDED(m_Connection->get_ServerState(&serverState));
	ATL::CComBSTR remoteId;
	ATLENSURE_SUCCEEDED(serverState->get_RemoteId(&remoteId));
	Log(L"%s", (PCWSTR) remoteId);

	SFTPTest();

	// Disconnect
	Log(L"Disconnect");
	ATLENSURE_SUCCEEDED(m_Connection->Disconnect());
}

void CTest::SFTPTest()
{
	ATLENSURE_SUCCEEDED(m_Connection->CreateSFTPConnection(&m_SFTP));

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

	ATL::CComQIPtr<sfFTPLib::ILogger> logger{ fileLogger };
	ATLENSURE_SUCCEEDED(m_SFTP->put_Logger(logger));

	ATLENSURE_SUCCEEDED(m_SFTP->Connect());

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

	// get current folder
	ATL::CComBSTR realPath = L".";
	Log(L"RealPath \"%s\"", (PCWSTR)realPath);

	ATL::CComBSTR currentFolder;
	ATLENSURE_SUCCEEDED(m_SFTP->RealPath(realPath, &currentFolder));
	Log(L"Home Folder = %s", (PCWSTR)currentFolder);

	// overriding currentFolder for debug purpose
	//currentFolder = L"/c/archive";
	//Log(L"Overriding current folder. \"%s\"", (PCWSTR)currentFolder);

	Log(L"Reading Directory \"%s\"", (PCWSTR)currentFolder);

	ATL::CComPtr<sfFTPLib::IFTPItems> items;
	ATLENSURE_SUCCEEDED(m_SFTP->ReadDirectory(currentFolder, &items));

	long count;
	ATLENSURE_SUCCEEDED(items->get_Count(&count));
	Log(L"Count = %d", count);

	// Enum
	if (count > 0)
	{
		ATL::CComPtr<IUnknown> unkEnum;
		ATLENSURE_SUCCEEDED(items->get__NewEnum(&unkEnum));
		ATL::CComQIPtr<IEnumVARIANT> pEnum(unkEnum);

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

				// TODO: Check for valid attributes (IsValidAttribute())
				sfFTPLib::ItemType itemType;
				ATLENSURE_SUCCEEDED(ftpItem->get_Type(&itemType));
				ATL::CComBSTR itemName;
				ATLENSURE_SUCCEEDED(ftpItem->get_Name(&itemName));
				ULONGLONG itemSize;
				ATLENSURE_SUCCEEDED(ftpItem->get_Size(&itemSize));

				auto str{ std::format(L"Type={:#x}; Name={}; Size={}", (unsigned int)itemType, itemName.operator LPWSTR(), itemSize)};

				VARIANT_BOOL isValidAttribute;
				ATLENSURE_SUCCEEDED(ftpItem->IsValidAttribute(sfFTPLib::ftpFTPItemAttributeModifyTime, &isValidAttribute));
				if (isValidAttribute)
				{
					FILETIME modifyTime;
					ATLENSURE_SUCCEEDED(ftpItem->get_ModifyTime(&modifyTime));
					std::wstring modifyTimeAsString;
					if (FILETIMEToISO8601(modifyTime, modifyTimeAsString) == S_OK)
					{
						str += L"; ModifyTime=";
						str += modifyTimeAsString;
					}
				}

				Log(str.c_str());
			}

			variant.Clear();
		}
	}

	// MakeDirectory
	CUnixPath makeDirectory{ (PCWSTR) currentFolder };
	makeDirectory.Append(L"testfolder");
	Log(L"MakeDirectory \"%s\"", (PCWSTR)makeDirectory);
	ATLENSURE_SUCCEEDED(m_SFTP->MakeDirectory(ATL::CComBSTR(makeDirectory)));
	Log(L"Directory \"%s\" created.", (PCWSTR)makeDirectory);

	// Rename
	CUnixPath renameFrom{ makeDirectory };
	CUnixPath renameTo = (PCWSTR)currentFolder;
	renameTo.Append(L"testfolder2");
	Log(L"Rename \"%s\" to \"%s\"", (PCWSTR)renameFrom, (PCWSTR)renameTo);
	ATLENSURE_SUCCEEDED(m_SFTP->Rename(ATL::CComBSTR(renameFrom), ATL::CComBSTR(renameTo), 0));

	// RemoveDirectory
	CUnixPath directoryToRemove{ renameTo };
	Log(L"RemoveDirectory \"%s\"", (PCWSTR) directoryToRemove);
	ATLENSURE_SUCCEEDED(m_SFTP->RemoveDirectory(ATL::CComBSTR(directoryToRemove)));
	Log(L"Directory \"%s\" removed.", (PCWSTR) directoryToRemove);

	// Creating temporary memory file
	ATL::CComPtr<IStream> memFile;
	static const DWORD dwSize{ 1000 * 1024 }; // 1000 KiB
	ATLENSURE_SUCCEEDED(CreateMemFile(dwSize, 0, &memFile));

	// Upload File
	CUnixPath UploadFile((PCWSTR) currentFolder);
	UploadFile.Append(L"memfile");
	ATL::CComBSTR bstrUploadFile{ (PCWSTR) UploadFile };

	Log(L"UploadFile to \"%s\"", (PCWSTR) bstrUploadFile);
	ATLENSURE_SUCCEEDED(m_SFTP->UploadFileEx(ATL::CComVariant(memFile.p), bstrUploadFile, sfFTPLib::ftpDataTransferTypeImage, 0, nullptr));
	Log(L"File successfully uploaded.");

	// Stat. Stat doesn't follow symbolic links.
	ATL::CComBSTR bstrStat = bstrUploadFile;
	Log(L"Stat \"%s\"", (PCWSTR) bstrStat);
	ATL::CComPtr<sfFTPLib::IFTPItem> pItem;
	ATLENSURE_SUCCEEDED(m_SFTP->Stat(bstrStat, sfFTPLib::ftpFTPItemAttributeSize, &pItem));
	VARIANT_BOOL isValidAttribute;
	ATLENSURE_SUCCEEDED(pItem->IsValidAttribute(sfFTPLib::ftpFTPItemAttributeSize, &isValidAttribute));
	if (isValidAttribute)
	{
		ULONGLONG size;
		ATLENSURE_SUCCEEDED(pItem->get_Size(&size));
		Log(L"File Size = %I64u.", size);
	}
	// DownloadFile to memory file
	ATL::CComPtr<IStream> pDownloadMemFile;
	if (CreateMemFile(dwSize, 0, &pDownloadMemFile) == S_OK)
	{
		ATL::CComBSTR bstrDownloadFile = bstrUploadFile;
		Log(L"DownloadFile \"%s\" to memfile", (PCWSTR) bstrDownloadFile);
		ATLENSURE_SUCCEEDED(m_SFTP->DownloadFileEx(bstrDownloadFile, ATL::CComVariant(pDownloadMemFile.p), sfFTPLib::ftpDataTransferTypeImage, 0, 0, sfFTPLib::ftpDownloadFlagReadBeyondEnd, nullptr));
		Log(L"File successfully downloaded.");
	}

	// DownloadFile to physical file
	ATL::CComBSTR bstrDownloadFile{ bstrUploadFile };
	wchar_t szCurrentDirectory[MAX_PATH]{};
	::GetCurrentDirectoryW(ARRAYSIZE(szCurrentDirectory), szCurrentDirectory);
	::PathAppendW(szCurrentDirectory, L"Download");
	::SHCreateDirectoryExW(nullptr, szCurrentDirectory, nullptr);
	::PathAppendW(szCurrentDirectory, L"memfile");
	ATL::CComBSTR downloadLocalFile{ szCurrentDirectory };
	Log(L"DownloadFile \"%s\" to \"%s\"", (PCWSTR) bstrDownloadFile, (PCWSTR) downloadLocalFile);
	ATLENSURE_SUCCEEDED(m_SFTP->DownloadFileEx(bstrDownloadFile, ATL::CComVariant(szCurrentDirectory), sfFTPLib::ftpDataTransferTypeImage, 0, MAXULONGLONG, sfFTPLib::ftpDownloadFlagReadBeyondEnd, nullptr));
	Log(L"File successfully downloaded.");

	Log(L"Closing channel.");
	ATLENSURE_SUCCEEDED(m_SFTP->Disconnect());
}

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

	wchar_t chMsg[LOG_EVENT_MSG_SIZE];
	va_list pArg;

	va_start(pArg, pszFormat);
	_vsntprintf_s(chMsg, LOG_EVENT_MSG_SIZE, LOG_EVENT_MSG_SIZE-1, pszFormat, pArg);

	std::wstring message{ chMsg };
	message += L"\n";
	ATLTRACE(message.c_str());
	_tprintf(message.c_str());
}

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

	const auto hMem{ ::GlobalAlloc(GMEM_MOVEABLE | GMEM_ZEROINIT, static_cast<SIZE_T>(nSize)) };
	if (!hMem)
	{
		return E_OUTOFMEMORY;
	}

	const auto 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
		return ::CreateStreamOnHGlobal(hMem, TRUE, retval);
	}
	return E_FAIL;
}