text: /**
   * @file src/ar-extractor/archive_wrapper.cpp
   * @brief Definition of ArchiveWrapper class.
   * @copyright (c) 2017 Avast Software, licensed under the MIT license
   */
  
  #include <algorithm>
  #include <cctype>
  #include <cstring>
  #include <fstream>
  #include <ostream>
  
  #include <rapidjson/document.h>
  #include <rapidjson/prettywriter.h>
  
  #include "retdec/utils/filesystem.h"
  #include "retdec/utils/string.h"
  #include "retdec/ar-extractor/archive_wrapper.h"
  
  using namespace retdec::utils;
  using namespace llvm;
  using namespace llvm::object;
  using namespace rapidjson;
  
  namespace
  {
  
  /**
   * Fix name.
   *
   * @param name input name
   *
   * @return new fixed name
   */
  std::string fixName(
  	std::string name)
  {
  	std::transform(name.begin(), name.end(), name.begin(),
  		[](const unsigned char c) {
  			return (!std::isalnum(c) && !strchr("-. \\", c)) ? '_' : c;
  		});
  	return name;
  }
  
  /**
   * Write single file.
   *
   * @param outputPath output file path
   * @param inputBuffer input buffer with data
   * @param errorMessage error message if @c false is returned
   *
   * @return @c true if no errors occurred, @c false otherwise
   */
  bool writeFile(
  	const std::string &outputPath,
  	const llvm::StringRef &inputBuffer,
  	std::string &errorMessage)
  {
  	std::ofstream outStream(outputPath, std::ofstream::binary);
  	if (outStream) {
  		outStream.write(inputBuffer.data(), inputBuffer.size());
  		if (!outStream) {
  			errorMessage = "Could not write content to output file";
  			return false;
  		}
  
  		// Everything OK.
  		return true;
  	}
  
  	errorMessage = "Could not open output file for writing";
  	return false;
  }
  
  /**
   * Check for errors.
   *
   * @param error input Error object
   * @param errorMessage error message if @c true is returned
   *
   * @return @c true if error was found, @c false otherwise
   */
  bool checkError(
  	Error &error,
  	std::string &errorMessage)
  {
  	if (error) {
  		errorMessage = llvm::toString(std::move(error));
  		return true;
  	}
  
  	return false;
  }
  
  } // anonymous namespace
  
  namespace retdec {
  namespace ar_extractor {
  
  /**
   * Constructor.
   *
   * @param archivePath path to input archive
   * @param succes result of object construction
   * @param errorMessage possible error message if @p success is set to false
   */
  ArchiveWrapper::ArchiveWrapper(
  	const std::string &archivePath,
  	bool &succes,
  	std::string &errorMessage)
  	: buffer(MemoryBuffer::getFile(llvm::Twine(archivePath)))
  {
  	succes = false;
  	if (!buffer) {
  		errorMessage = "Could not create file buffer";
  		return;
  	}
  
  	Error error = Error::success();
  	archive = std::make_unique<Archive>(buffer.get()->getMemBufferRef(), error);
  	if (error) {
  		errorMessage = llvm::toString(std::move(error));
  		return;
  	}
  
  	// Get object count - this iterates over all objects.
  	succes = getCount(objectCount, errorMessage);
  }
  
  /**
   * Get number of objects.
   *
   * @return number of objects in file.
   */
  std::size_t ArchiveWrapper::getNumberOfObjects() const
  {
  	return objectCount;
  }
  
  /**
   * Check whether archive is thin archive.
   *
   * @return @c true if archive is thin, @c false otherwise
   */
  bool ArchiveWrapper::isThinArchive() const
  {
  	return archive->isThin();
  }
  
  /**
   * Check whether archive is empty.
   *
   * @return @c true if archive is empty, @c false otherwise
   */
  bool ArchiveWrapper::isEmptyArchive() const
  {
  	return objectCount == 0;
  }
  
  /**
   * Get list of object file names in plain-text.
   *
   * @param result string with complete list in plain-text
   * @param errorMessage possible error message if @c false is returned
   * @param niceNames names will be made nicer if @c true
   * @param numbers list will be numbered if @c true
   *
   * @return @c true if no errors occurred, @c false otherwise
   */
  bool ArchiveWrapper::getPlainTextList(
  	std::string &result,
  	std::string &errorMessage,
  	bool niceNames,
  	bool numbers) const
  {
  	result.clear();
  	std::size_t counter = 0;
  	std::vector<std::string> names;
  
  	if (getNames(names, errorMessage)) {
  		for (const auto &name : names) {
  			const auto outName = niceNames ? fixName(name) : name;
  			result += numbers ? std::to_string(counter++) + "\t" : "";
  			result += outName + "\n";
  		}
  		return true;
  	}
  
  	return false;
  }
  
  /**
   * Get list of object file names in JSON format.
   *
   * @param result string with complete list in JSON format
   * @param errorMessage possible error message if @c false is returned
   * @param niceNames names will be made nicer if @c true
   * @param numbers list will be numbered if @c true
   *
   * @return @c true if no errors occurred, @c false otherwise
   */
  bool ArchiveWrapper::getJsonList(
  	std::string &result,
  	std::string &errorMessage,
  	bool niceNames,
  	bool numbers) const
  {
  	result.clear();
  	std::size_t counter = 0;
  	std::vector<std::string> names;
  
  	if (getNames(names, errorMessage)) {
  		Value objects(kArrayType);
  		Document outFile(kObjectType);
  		auto& allocator = outFile.GetAllocator();
  
  		for (const auto &name : names) {
  			const std::string outName = niceNames ? fixName(name) : name;
  
  			Value object(kObjectType);
  			object.AddMember("name", Value(outName.c_str(), allocator).Move(), allocator);
  			if (numbers) {
  				object.AddMember("index", static_cast<uint64_t>(counter++), allocator);
  			}
  			objects.PushBack(object, allocator);
  		}
  		outFile.AddMember("objects", objects, allocator);
  
  		StringBuffer buffer;
  		PrettyWriter<StringBuffer> writer(buffer);
  		outFile.Accept(writer);
  		result = buffer.GetString();
  		return true;
  	}
  
  	return false;
  }
  
  /**
   * Extract all object files.
   *
   * If directory is not specified, current directory is used. If multiple files
   * have same name, they are decorated with their index suffix.
   *
   * @param errorMessage possible error message if @c false is returned
   * @param directory optional target directory
   *
   * @return @c true if no errors occurred, @c false otherwise
   */
  bool ArchiveWrapper::extract(
  	std::string &errorMessage,
  	const std::string &directory) const
  {
  	// Check if target directory exists if string not empty.
  	if (!directory.empty() && !fs::is_directory(directory)) {
  		errorMessage = "Invalid target directory";
  		return false;
  	}
  
  	/**
  	 * @todo In writeFile function call, separator is added to path, this may
  	 * or may not work on Windows OS.
  	 */
  
  	// Map for non-unique names - counts number of name occurrences.
  	std::map<std::string, std::size_t> nameMap;
  
  	Error error = Error::success();
  	for (const auto &child : archive->children(error)) {
  		if (checkError(error, errorMessage)) {
  			return false;
  		}
  
  		// Try to get name.
  		auto nameOrErr = child.getName();
  		std::string name = nameOrErr ? fixName(nameOrErr->str()) : "invalid_name";
  
  		// Increment name count and fix name if it is not unique.
  		if (++nameMap[name] != 1) {
  			name += "." + std::to_string(nameMap[name]);
  		}
  
  		auto bufferOrErr = child.getBuffer();
  		if (!bufferOrErr) {
  			errorMessage = "Could not get file buffer";
  			return false;
  		}
  		else {
  			auto dir = directory.empty() ? directory : directory + '/';
  			if (!writeFile(dir + name, *bufferOrErr, errorMessage)) {
  				return false;
  			}
  		}
  	}
  
  	return !checkError(error, errorMessage);
  }
  
  /**
   * Extract object file by its name.
   *
   * If output path is not given, object name and current directory is used. If
   * multiple files with the same name are present, only first one is extracted.
   *
   * @param name target name
   * @param errorMessage possible error message if @c false is returned
   * @param outputPath optional output path
   *
   * @return @c true if no errors occurred, @c false otherwise
   */
  bool ArchiveWrapper::extractByName(
  	const std::string &name,
  	std::string &errorMessage,
  	const std::string &outputPath) const
  {
  	Error error = Error::success();
  	for (const auto &child : archive->children(error)) {
  		if (checkError(error, errorMessage)) {
  			return false;
  		}
  
  		auto nameOrErr = child.getName();
  		if (!nameOrErr) {
  			// Could not get name.
  			continue;
  		}
  
  		if (name != fixName(nameOrErr->str())) {
  			// Name does not match.
  			continue;
  		}
  
  		// Get buffer and try to write to a file.
  		auto bufferOrErr = child.getBuffer();
  		if (!bufferOrErr) {
  			errorMessage = "Could not get file buffer";
  			return false;
  		}
  		else {
  			auto path = outputPath.empty() ? name : outputPath;
  			return writeFile(path, *bufferOrErr, errorMessage);
  		}
  	}
  
  	if (checkError(error, errorMessage)) {
  		return false;
  	}
  
  	errorMessage = "Could not find desired file";
  	return false;
  }
  
  /**
   * Extract object file by its index.
   *
   * If output path is not given, object name and current directory is used. If
   * name cannot be retrieved, name 'invalid_name' is used.
   *
   * @param index target index
   * @param errorMessage possible error message if @c false is returned
   * @param outputPath optional output path
   *
   * @return @c true if no errors occurred, @c false otherwise
   */
  bool ArchiveWrapper::extractByIndex(
  	const std::size_t index,
  	std::string &errorMessage,
  	const std::string &outputPath) const
  {
  	Error error = Error::success();
  	std::size_t counter = 0;
  	for (const auto &child : archive->children(error)) {
  		if (checkError(error, errorMessage)) {
  			return false;
  		}
  
  		// No random access available.
  		if (index != counter++) {
  			continue;
  		}
  
  		// Get buffer and try to write to a file.
  		auto bufferOrErr = child.getBuffer();
  		if (!bufferOrErr) {
  			errorMessage = "Could not get file buffer";
  			return false;
  		}
  		else {
  			std::string path;
  			if (outputPath.empty()) {
  				// No path given - use object name.
  				auto nameOrErr = child.getName();
  				path = nameOrErr ? fixName(nameOrErr->str()) : "invalid_name";
  			}
  			else {
  				path = outputPath;
  			}
  			return writeFile(path, *bufferOrErr, errorMessage);
  		}
  	}
  
  	if (checkError(error, errorMessage)) {
  		return false;
  	}
  
  	errorMessage = "Could not find desired file";
  	return false;
  }
  
  /**
   * Get names of all object files in archive.
   *
   * If name of object could not be read from input archive, name 'invalid_name'
   * is used.
   *
   * @param result container where names will be added
   * @param errorMessage possible error message if @c false is returned
   *
   * @return @c true if no errors occurred, @c false otherwise
   */
  bool ArchiveWrapper::getNames(
  	std::vector<std::string> &result,
  	std::string &errorMessage) const
  {
  	Error error = Error::success();
  	for (const auto &child : archive->children(error)) {
  		if (checkError(error, errorMessage)) {
  			return false;
  		}
  
  		auto nameOrErr = child.getName();
  		if (!nameOrErr) {
  			result.emplace_back("invalid_name");
  		}
  		else {
  			result.push_back(*nameOrErr);
  		}
  	}
  
  	return !checkError(error, errorMessage);
  }
  
  /**
   * Get object count.
   *
   * @param count result
   * @param errorMessage possible error message if @c false is returned
   *
   * @return @c true if no errors occurred, @c false otherwise
   */
  bool ArchiveWrapper::getCount(
  	std::size_t &count,
  	std::string &errorMessage) const
  {
  	Error error = Error::success();
  	count = 0; // Reset counter.
  	const auto &ar = archive;
  	for (auto i = ar->child_begin(error), e = ar->child_end(); i != e; ++i) {
  		if (checkError(error, errorMessage)) {
  			return false;
  		}
  		count++;
  	}
  
  	return !checkError(error, errorMessage);
  }
  
  } // namespace ar_extractor
  } // namespace retdec
  