From 1a1a20895bfa9778df82b339bbd05b27ca1d7c10 Mon Sep 17 00:00:00 2001 From: Nikolaj Schlej Date: Sun, 23 Apr 2023 16:46:59 -0700 Subject: [PATCH] Refactor UEFIExtract a bit --- UEFIExtract/uefidump.cpp | 108 ++++++------ UEFIExtract/uefiextract_main.cpp | 282 +++++++++++++++++-------------- 2 files changed, 202 insertions(+), 188 deletions(-) diff --git a/UEFIExtract/uefidump.cpp b/UEFIExtract/uefidump.cpp index fd93961..1f10d93 100644 --- a/UEFIExtract/uefidump.cpp +++ b/UEFIExtract/uefidump.cpp @@ -85,67 +85,59 @@ USTATUS UEFIDumper::recursiveDump(const UModelIndex & index) if (!index.isValid()) return U_INVALID_PARAMETER; - //UByteArray itemHeader = model.header(index); - //UByteArray fileHeader = model.header(model.findParentOfType(index, Types::File)); - - //if (guid.length() == 0 || - // (itemHeader.size() >= sizeof (EFI_GUID) && guidToUString(*(const EFI_GUID*)itemHeader.constData()) == guid) || - // (fileHeader.size() >= sizeof(EFI_GUID) && guidToUString(*(const EFI_GUID*)fileHeader.constData()) == guid)) { - - // Construct file name - UString orgName = uniqueItemName(index); - UString name = orgName; - bool nameFound = false; - for (int i = 1; i < 1000; ++i) { - if (!isExistOnFs(name + UString("_info.txt"))) { - nameFound = true; - break; - } - name = orgName + UString("_") + usprintf("%03d", i); + // Construct file name + UString orgName = uniqueItemName(index); + UString name = orgName; + bool nameFound = false; + for (int i = 1; i < 1000; ++i) { + if (!isExistOnFs(name + UString("_info.txt"))) { + nameFound = true; + break; } - - if (!nameFound) { - printf("Cannot find unique name for \"%s\".\n", (const char*)orgName.toLocal8Bit()); - return U_INVALID_PARAMETER; //TODO: replace with proper errorCode + name = orgName + UString("_") + usprintf("%03d", i); + } + + if (!nameFound) { + printf("Cannot find unique name for \"%s\".\n", (const char*)orgName.toLocal8Bit()); + return U_INVALID_PARAMETER; //TODO: replace with proper errorCode + } + + // Add header and body only for leaf sections + if (model.rowCount(index) == 0) { + // Header + UByteArray data = model.header(index); + if (!data.isEmpty()) { + std::ofstream file; + UString str = name + UString("_header.bin"); + file.open(str.toLocal8Bit(), std::ios::out | std::ios::binary); + file.write(data.constData(), data.size()); + file.close(); } - - // Add header and body only for leaf sections - if (model.rowCount(index) == 0) { - // Header - UByteArray data = model.header(index); - if (!data.isEmpty()) { - std::ofstream file; - UString str = name + UString("_header.bin"); - file.open(str.toLocal8Bit(), std::ios::out | std::ios::binary); - file.write(data.constData(), data.size()); - file.close(); - } - - // Body - data = model.body(index); - if (!data.isEmpty()) { - std::ofstream file; - UString str = name + UString("_body.bin"); - file.open(str.toLocal8Bit(), std::ios::out | std::ios::binary); - file.write(data.constData(), data.size()); - file.close(); - } + + // Body + data = model.body(index); + if (!data.isEmpty()) { + std::ofstream file; + UString str = name + UString("_body.bin"); + file.open(str.toLocal8Bit(), std::ios::out | std::ios::binary); + file.write(data.constData(), data.size()); + file.close(); } - // Info - UString info = "Type: " + itemTypeToUString(model.type(index)) + "\n" + - "Subtype: " + itemSubtypeToUString(model.type(index), model.subtype(index)) + "\n"; - if (model.text(index).length() > 0) - info += "Text: " + model.text(index) + "\n"; - info += model.info(index) + "\n"; - - std::ofstream file; - UString str = name + UString("_info.txt"); - file.open(str.toLocal8Bit(), std::ios::out); - file.write(info.toLocal8Bit(), info.length()); - file.close(); - - dumped = true; - //} + } + // Info + UString info = "Type: " + itemTypeToUString(model.type(index)) + "\n" + + "Subtype: " + itemSubtypeToUString(model.type(index), model.subtype(index)) + "\n"; + if (model.text(index).length() > 0) + info += "Text: " + model.text(index) + "\n"; + info += model.info(index) + "\n"; + + std::ofstream file; + UString str = name + UString("_info.txt"); + file.open(str.toLocal8Bit(), std::ios::out); + file.write(info.toLocal8Bit(), info.length()); + file.close(); + + dumped = true; // Process child items USTATUS result; diff --git a/UEFIExtract/uefiextract_main.cpp b/UEFIExtract/uefiextract_main.cpp index 3bd2497..0a022ab 100644 --- a/UEFIExtract/uefiextract_main.cpp +++ b/UEFIExtract/uefiextract_main.cpp @@ -34,141 +34,94 @@ void print_usage() { std::cout << "UEFIExtract " PROGRAM_VERSION << std::endl << "Usage: UEFIExtract {-h | --help | -v | --version} - show help and/or version information." << std::endl - << " UEFIExtract imagefile - generate report and dump only leaf tree items into .dump folder." << std::endl - << " UEFIExtract imagefile all - generate report and dump all tree items." << std::endl - << " UEFIExtract imagefile unpack - generate report and dump all tree items into a single folder." << std::endl - << " UEFIExtract imagefile dump - only generate dump, no report needed." << std::endl - << " UEFIExtract imagefile report - only generate report, no dump needed." << std::endl - << " UEFIExtract imagefile guids - generate a CSV file with named GUIDs present in the image." << std::endl + << " UEFIExtract imagefile - generate report and GUID database, then dump only leaf tree items into .dump folder." << std::endl + << " UEFIExtract imagefile all - generate report and GUID database, then dump all tree items into .dump folder." << std::endl + << " UEFIExtract imagefile unpack - generate report, then dump all tree items into a single .dump folder (legacy UEFIDump compatibility mode)." << std::endl + << " UEFIExtract imagefile dump - only generate dump, no report or GUID database needed." << std::endl + << " UEFIExtract imagefile report - only generate report, no dump or GUID database needed." << std::endl + << " UEFIExtract imagefile guids - only generate GUID database, no dump or report needed." << std::endl << " UEFIExtract imagefile GUID_1 ... [ -o FILE_1 ... ] [ -m MODE_1 ... ] [ -t TYPE_1 ... ] -" << std::endl - << " Dump only FFS file(s) with specific GUID(s), without report." << std::endl + << " Dump only FFS file(s) with specific GUID(s), without report or GUID database." << std::endl << " Type is section type or FF to ignore. Mode is one of: all, body, header, info, file." << std::endl - << "Return value is a bit mask where 0 at position N means that file with GUID_N was found and unpacked, 1 otherwise." << std::endl; + << " Return value is a bit mask where 0 at position N means that file with GUID_N was found and unpacked, 1 otherwise." << std::endl; } int main(int argc, char *argv[]) { initGuidDatabase("guids.csv"); - if (argc > 1) { - if (argc == 2) { - UString arg = UString(argv[1]); - if (arg == UString("-h") || arg == UString("--help")) { - print_usage(); - return U_SUCCESS; - } - else if (arg == UString("-v") || arg == UString("--version")) { - std::cout << PROGRAM_VERSION << std::endl; - return U_SUCCESS; - } + if (argc <= 1) { + print_usage(); + return 1; + } + + // Help and version + if (argc == 2) { + UString arg = UString(argv[1]); + if (arg == UString("-h") || arg == UString("--help")) { + print_usage(); + return 0; } - - // Check that input file exists - USTATUS result; - UByteArray buffer; - UString path = getAbsPath(argv[1]); - result = readFileIntoBuffer(path, buffer); - if (result) - return result; - - // Hack to support legacy UEFIDump mode. - if (argc == 3 && !std::strcmp(argv[2], "unpack")) { - UEFIDumper uefidumper; - return (uefidumper.dump(buffer, UString(argv[1])) != U_SUCCESS); + else if (arg == UString("-v") || arg == UString("--version")) { + std::cout << PROGRAM_VERSION << std::endl; + return 0; } - - // Create model and ffsParser - TreeModel model; - FfsParser ffsParser(&model); - // Parse input buffer - result = ffsParser.parse(buffer); - if (result) - return result; - - ffsParser.outputInfo(); - - // Create ffsDumper - FfsDumper ffsDumper(&model); - - // Dump only leaf elements, no report - if (argc == 3 && !std::strcmp(argv[2], "dump")) { - return (ffsDumper.dump(model.index(0, 0), path + UString(".dump")) != U_SUCCESS); + } + + // Check that input file exists + USTATUS result; + UByteArray buffer; + UString path = getAbsPath(argv[1]); + result = readFileIntoBuffer(path, buffer); + if (result) + return result; + + // Hack to support legacy UEFIDump mode + if (argc == 3 && !std::strcmp(argv[2], "unpack")) { + UEFIDumper uefidumper; + return (uefidumper.dump(buffer, UString(argv[1])) != U_SUCCESS); + } + + // Create model and ffsParser + TreeModel model; + FfsParser ffsParser(&model); + // Parse input buffer + result = ffsParser.parse(buffer); + if (result) + return result; + + ffsParser.outputInfo(); + + // Create ffsDumper + FfsDumper ffsDumper(&model); + + // Dump only leaf elements, no report or GUID database + if (argc == 3 && !std::strcmp(argv[2], "dump")) { + return (ffsDumper.dump(model.index(0, 0), path + UString(".dump")) != U_SUCCESS); + } + // Dump named GUIDs found in the image, no dump or report + else if (argc == 3 && !std::strcmp(argv[2], "guids")) { + GuidDatabase db = guidDatabaseFromTreeRecursive(&model, model.index(0, 0)); + if (!db.empty()) { + return guidDatabaseExportToFile(path + UString(".guids.csv"), db); } - // Dump named GUIDs found in the image - else if (argc == 3 && !std::strcmp(argv[2], "guids")) { - GuidDatabase db = guidDatabaseFromTreeRecursive(&model, model.index(0, 0)); - if (!db.empty()) { - return guidDatabaseExportToFile(path + UString(".guids.csv"), db); - } + } + // Generate report, no dump or GUID database + else if (argc == 3 && !std::strcmp(argv[2], "report")) { + FfsReport ffsReport(&model); + std::vector report = ffsReport.generate(); + if (report.size()) { + std::ofstream file; + file.open((path + UString(".report.txt")).toLocal8Bit()); + for (size_t i = 0; i < report.size(); i++) + file << report[i].toLocal8Bit() << '\n'; + return 0; } - else if (argc > 3 || - (argc == 3 && std::strcmp(argv[2], "all") != 0 && std::strcmp(argv[2], "report") != 0)) { // Dump specific files, without report - std::vector inputs, outputs; - std::vector modes; - std::vector sectionTypes; - ReadType readType = READ_INPUT; - for (int i = 2; i < argc; i++) { - const char *arg = argv[i]; - if (!std::strcmp(arg, "-i")) { - readType = READ_INPUT; - continue; - } else if (!std::strcmp(arg, "-o")) { - readType = READ_OUTPUT; - continue; - } else if (!std::strcmp(arg, "-m")) { - readType = READ_MODE; - continue; - } else if (!std::strcmp(arg, "-t")) { - readType = READ_SECTION; - continue; - } - - if (readType == READ_INPUT) { - inputs.push_back(arg); - } else if (readType == READ_OUTPUT) { - outputs.push_back(getAbsPath(arg)); - } else if (readType == READ_MODE) { - if (!std::strcmp(arg, "all")) - modes.push_back(FfsDumper::DUMP_ALL); - else if (!std::strcmp(arg, "body")) - modes.push_back(FfsDumper::DUMP_BODY); - else if (!std::strcmp(arg, "header")) - modes.push_back(FfsDumper::DUMP_HEADER); - else if (!std::strcmp(arg, "info")) - modes.push_back(FfsDumper::DUMP_INFO); - else if (!std::strcmp(arg, "file")) - modes.push_back(FfsDumper::DUMP_FILE); - else - return U_INVALID_PARAMETER; - } else if (readType == READ_SECTION) { - char *converted = const_cast(arg); - UINT8 sectionType = (UINT8)std::strtol(arg, &converted, 16); - if (converted == arg) - return U_INVALID_PARAMETER; - sectionTypes.push_back(sectionType); - } - } - if (inputs.empty() || (!outputs.empty() && inputs.size() != outputs.size()) || - (!modes.empty() && inputs.size() != modes.size()) || - (!sectionTypes.empty() && inputs.size() != sectionTypes.size())) - return U_INVALID_PARAMETER; - - USTATUS lastError = U_SUCCESS; - for (size_t i = 0; i < inputs.size(); i++) { - UString outPath = outputs.empty() ? path + UString(".dump") : outputs[i]; - FfsDumper::DumpMode mode = modes.empty() ? FfsDumper::DUMP_ALL : modes[i]; - UINT8 type = sectionTypes.empty() ? FfsDumper::IgnoreSectionType : sectionTypes[i]; - result = ffsDumper.dump(model.index(0, 0), outPath, mode, type, inputs[i]); - if (result) { - std::cout << "Guid " << inputs[i].toLocal8Bit() << " failed with " << result << " code!" << std::endl; - lastError = result; - } - } - - return lastError; - } - - // Create ffsReport + return 1; + } + // Either default or all mode + else if (argc == 2 || (argc == 3 && !std::strcmp(argv[2], "all"))) { + // Generate report FfsReport ffsReport(&model); std::vector report = ffsReport.generate(); if (report.size()) { @@ -177,19 +130,88 @@ int main(int argc, char *argv[]) for (size_t i = 0; i < report.size(); i++) file << report[i].toLocal8Bit() << '\n'; } - - // Dump all non-leaf elements, with report, default + + // Create GUID database + GuidDatabase db = guidDatabaseFromTreeRecursive(&model, model.index(0, 0)); + if (!db.empty()) { + guidDatabaseExportToFile(path + UString(".guids.csv"), db); + } + + // Dump all non-leaf elements, with report and GUID database, default if (argc == 2) { return (ffsDumper.dump(model.index(0, 0), path + UString(".dump")) != U_SUCCESS); } - else if (argc == 3 && !std::strcmp(argv[2], "all")) { // Dump every element with report + else if (argc == 3 && !std::strcmp(argv[2], "all")) { // Dump every element with report and GUID database return (ffsDumper.dump(model.index(0, 0), path + UString(".dump"), FfsDumper::DUMP_ALL) != U_SUCCESS); } - else if (argc == 3 && !std::strcmp(argv[2], "report")) { // Skip dumping - return 0; - } } - + // Dump specific files, without report or GUID database + else { + std::vector inputs, outputs; + std::vector modes; + std::vector sectionTypes; + ReadType readType = READ_INPUT; + for (int i = 2; i < argc; i++) { + const char *arg = argv[i]; + if (!std::strcmp(arg, "-i")) { + readType = READ_INPUT; + continue; + } else if (!std::strcmp(arg, "-o")) { + readType = READ_OUTPUT; + continue; + } else if (!std::strcmp(arg, "-m")) { + readType = READ_MODE; + continue; + } else if (!std::strcmp(arg, "-t")) { + readType = READ_SECTION; + continue; + } + + if (readType == READ_INPUT) { + inputs.push_back(arg); + } else if (readType == READ_OUTPUT) { + outputs.push_back(getAbsPath(arg)); + } else if (readType == READ_MODE) { + if (!std::strcmp(arg, "all")) + modes.push_back(FfsDumper::DUMP_ALL); + else if (!std::strcmp(arg, "body")) + modes.push_back(FfsDumper::DUMP_BODY); + else if (!std::strcmp(arg, "header")) + modes.push_back(FfsDumper::DUMP_HEADER); + else if (!std::strcmp(arg, "info")) + modes.push_back(FfsDumper::DUMP_INFO); + else if (!std::strcmp(arg, "file")) + modes.push_back(FfsDumper::DUMP_FILE); + else + return U_INVALID_PARAMETER; + } else if (readType == READ_SECTION) { + char *converted = const_cast(arg); + UINT8 sectionType = (UINT8)std::strtol(arg, &converted, 16); + if (converted == arg) + return U_INVALID_PARAMETER; + sectionTypes.push_back(sectionType); + } + } + if (inputs.empty() || (!outputs.empty() && inputs.size() != outputs.size()) || + (!modes.empty() && inputs.size() != modes.size()) || + (!sectionTypes.empty() && inputs.size() != sectionTypes.size())) + return U_INVALID_PARAMETER; + + USTATUS lastError = U_SUCCESS; + for (size_t i = 0; i < inputs.size(); i++) { + UString outPath = outputs.empty() ? path + UString(".dump") : outputs[i]; + FfsDumper::DumpMode mode = modes.empty() ? FfsDumper::DUMP_ALL : modes[i]; + UINT8 type = sectionTypes.empty() ? FfsDumper::IgnoreSectionType : sectionTypes[i]; + result = ffsDumper.dump(model.index(0, 0), outPath, mode, type, inputs[i]); + if (result) { + std::cout << "Guid " << inputs[i].toLocal8Bit() << " failed with " << result << " code!" << std::endl; + lastError = result; + } + } + + return lastError; + } + // If parameters are different, show version and usage information print_usage(); return 1;