OperationRestoreSecurity.cpp

#include "OperationRestoreSecurity.h"
#include "InputOutput.h"
#include "Helpers.h"

#include <fstream>
#include <locale>
#include <codecvt>

ClassFactory<OperationRestoreSecurity> OperationRestoreSecurity::RegisteredFactory(GetCommand());

OperationRestoreSecurity::OperationRestoreSecurity(std::queue<std::wstring> & oArgList, const std::wstring & sCommand) : Operation(oArgList)
{
	// exit if there are not enough arguments to parse
	const std::vector<std::wstring> sSubArgs = ProcessAndCheckArgs(1, oArgList, L"\\0");

	// open the file
	std::wifstream fFile(sSubArgs.at(0).c_str());
	if (!fFile.is_open() || !fFile.good())
	{
		Print(L"ERROR: Could not open file '{}'.", sSubArgs.at(0));
		std::exit(-1);
	}

	// adapt the stream to read windows unicode files
	(void) fFile.imbue(std::locale(fFile.getloc(), new std::codecvt_utf8_utf16<wchar_t,
		0x10ffff, std::consume_header>));

	// read the file line-by-line
	std::wstring sLine;
	while (std::getline(fFile, sLine))
	{
		// parse the file name and descriptor which are separated by a '|' character
		// also, sometimes a carriage return appears in the input stream so adding 
		// it here ensures it is stripped from the very end
		std::vector<std::wstring> oLineItems = SplitArgs(sLine, L"\\||\r");

		// convert the long string descriptor its binary equivalent
		SmartPointer<PSECURITY_DESCRIPTOR> tDesc(LocalFree, nullptr);
		ULONG iDescSize = 0;
		if (oLineItems.size() != 2 ||
			ConvertStringSecurityDescriptorToSecurityDescriptor(oLineItems.at(1).c_str(),
			SDDL_REVISION_1, &tDesc, &iDescSize) == 0)
		{
			Print(L"ERROR: Unable to parse string security descriptor file for restoration.");
			std::exit(-1);
		}

		// retain an owned snapshot for repeated and concurrent restores
		const BYTE * pDesc = static_cast<const BYTE *>(*tDesc);
		oImportMap[oLineItems.at(0)].assign(pDesc, pDesc + iDescSize);
	}

	// reject incomplete input before starting any scan or permission changes
	if (!fFile.eof())
	{
		Print(L"ERROR: Unable to read UTF-8 input file '{}'.", sSubArgs.at(0));
		std::exit(-1);
	}

	// cleanup
	fFile.close();

	// flag this as being an ace-level action
	AppliesToSd = true;
	AppliesToDacl = true;
	AppliesToSacl = true;
	AppliesToOwner = true;
	AppliesToGroup = true;
}

bool OperationRestoreSecurity::ProcessSdAction(std::wstring & sFileName, ObjectEntry & tObjectEntry, PSECURITY_DESCRIPTOR & tDescriptor, bool & bDescReplacement)
{
	const auto oSecInfo = oImportMap.find(sFileName);
	if (oSecInfo == oImportMap.end())
	{
		InputOutput::AddError(L"Import file did not contain a descriptor for this path");
		return false;
	}

	// give each entry its own descriptor before releasing the current one
	SmartPointer<PSECURITY_DESCRIPTOR> tNewDescriptor(LocalFree,
		LocalAlloc(LMEM_FIXED, oSecInfo->second.size()));
	if (!tNewDescriptor)
	{
		InputOutput::AddError(L"Could not allocate security descriptor for restoration.");
		return false;
	}
	memcpy(tNewDescriptor, oSecInfo->second.data(), oSecInfo->second.size());

	if (bDescReplacement) LocalFree(tDescriptor);
	tDescriptor = tNewDescriptor;
	*tNewDescriptor = nullptr;
	bDescReplacement = true;
	return true;
}