From 80bcc0268d61ade51cfb372b3d452277c8cbf29b Mon Sep 17 00:00:00 2001 From: YZBruh Date: Mon, 25 Aug 2025 13:12:46 +0300 Subject: [PATCH] pmt: Improve, etc. - Added custom file descriptors for better print status handling. - Added macro to be used for adding functions to make things easier. - Duplicate function protection added. --- include/PartitionManager/PartitionManager.hpp | 3 +- src/FunctionManager.cpp | 8 +- src/PartitionManager.cpp | 75 +++++++++++-------- 3 files changed, 54 insertions(+), 32 deletions(-) diff --git a/include/PartitionManager/PartitionManager.hpp b/include/PartitionManager/PartitionManager.hpp index 4c5792d..d9c3c3c 100644 --- a/include/PartitionManager/PartitionManager.hpp +++ b/include/PartitionManager/PartitionManager.hpp @@ -51,7 +51,7 @@ private: public: void registerFunction(std::unique_ptr _func, CLI::App &_app); - [[nodiscard]] bool isUsed(std::string name) const; + [[nodiscard]] bool isUsed(const std::string &name) const; [[nodiscard]] bool handleAll() const; }; @@ -75,6 +75,7 @@ using VariableTable = basic_variables; using Error = Helper::Error; extern std::unique_ptr Variables; +extern FILE *pstdout, *pstderr; int Main(int argc, char **argv); diff --git a/src/FunctionManager.cpp b/src/FunctionManager.cpp index 9da3329..5efffe1 100644 --- a/src/FunctionManager.cpp +++ b/src/FunctionManager.cpp @@ -70,6 +70,12 @@ void processCommandLine(std::vector &vec1, void basic_function_manager::registerFunction( std::unique_ptr _func, CLI::App &_app) { LOGN(PMTF, INFO) << "registering function: " << _func->name() << std::endl; + for (const auto& f : _functions) { + if (strcmp(f->name(), _func->name()) != 0) { + LOGN(PMTF, INFO) << "Function is already registered: " << _func->name() << ". Skipping." << std::endl; + return; + } + } if (!_func->init(_app)) throw Error("Cannot init function: %s\n", _func->name()); _functions.push_back(std::move(_func)); @@ -77,7 +83,7 @@ void basic_function_manager::registerFunction( << std::endl; } -bool basic_function_manager::isUsed(const std::string name) const { +bool basic_function_manager::isUsed(const std::string &name) const { if (_functions.empty()) return false; for (const auto &func : _functions) { if (func->name() == name) return func->isUsed(); diff --git a/src/PartitionManager.cpp b/src/PartitionManager.cpp index b04de6e..e0d6826 100644 --- a/src/PartitionManager.cpp +++ b/src/PartitionManager.cpp @@ -26,18 +26,9 @@ namespace PartitionManager { -__attribute__((constructor)) -void init() { - Helper::LoggingProperties::setLogFile("/sdcard/Documents/last_pmt_logs.log"); -} - -static void sigHandler(const int sig) { - // Even if only SIGINT is to be captured for now, this is still a more appropriate code - if (sig == SIGINT) println("\n%sInterrupted.%s", YELLOW, STYLE_RESET); - exit(sig); -} - -auto Variables = std::make_unique(); +// Usage: REGISTER_FUNCTION(FUNCTION_CLASS); +#define REGISTER_FUNCTION(cls) \ + FuncManager.registerFunction(std::make_unique(), AppMain) basic_variables::basic_variables() : logFile(Helper::LoggingProperties::FILE), onLogical(false), @@ -49,9 +40,38 @@ basic_variables::basic_variables() } } +__attribute__((constructor)) +void init() { + Helper::LoggingProperties::setLogFile("/sdcard/Documents/last_pmt_logs.log"); +} + +static void sigHandler(const int sig) { + // Even if only SIGINT is to be captured for now, this is still a more appropriate code + if (sig == SIGINT) println("\n%sInterrupted.%s", YELLOW, STYLE_RESET); + exit(sig); +} + +static int write(void *cookie, const char *buf, const int size) { + auto *real = static_cast(cookie); + if (!Variables->quietProcess) return fwrite(buf, 1, static_cast(size), real); + else return size; +} + +static FILE* make_fp(FILE* real) { + return funopen(real, nullptr, write, nullptr, nullptr); +} + +auto Variables = std::make_unique(); +FILE* pstdout = make_fp(stdout); +FILE* pstderr = make_fp(stderr); + int Main(int argc, char **argv) { try { // try-catch start + Helper::garbageCollector collector; + collector.closeAfterProgress(pstdout); + collector.closeAfterProgress(pstderr); + if (argc < 2) { println( "Usage: %s [OPTIONS] [SUBCOMMAND]\nUse --help for more information.", @@ -93,17 +113,15 @@ int Main(int argc, char **argv) { AppMain.add_flag("-v,--version", Variables->viewVersion, "Print version and exit"); - FuncManager.registerFunction(std::make_unique(), AppMain); - FuncManager.registerFunction(std::make_unique(), AppMain); - FuncManager.registerFunction(std::make_unique(), AppMain); - FuncManager.registerFunction(std::make_unique(), - AppMain); - FuncManager.registerFunction(std::make_unique(), AppMain); - FuncManager.registerFunction(std::make_unique(), AppMain); - FuncManager.registerFunction(std::make_unique(), AppMain); - FuncManager.registerFunction(std::make_unique(), AppMain); - FuncManager.registerFunction(std::make_unique(), - AppMain); + REGISTER_FUNCTION(backupFunction); + REGISTER_FUNCTION(flashFunction); + REGISTER_FUNCTION(eraseFunction); + REGISTER_FUNCTION(partitionSizeFunction); + REGISTER_FUNCTION(infoFunction); + REGISTER_FUNCTION(realPathFunction); + REGISTER_FUNCTION(typeFunction); + REGISTER_FUNCTION(rebootFunction); + REGISTER_FUNCTION(memoryTestFunction); CLI11_PARSE(AppMain, argc, argv); @@ -132,8 +150,7 @@ int Main(int argc, char **argv) { } catch (Helper::Error &error) { // catch Helper::Error - if (!Variables->quietProcess) - fprintf(stderr, "%s%sERROR(S) OCCURRED:%s\n%s", RED, BOLD, STYLE_RESET, + fprintf(pstderr, "%s%sERROR(S) OCCURRED:%s\n%s", RED, BOLD, STYLE_RESET, error.what()); return EXIT_FAILURE; } catch (CLI::Error &error) { @@ -148,17 +165,15 @@ int Main(int argc, char **argv) { void print(const char *format, ...) { va_list args; va_start(args, format); - if (!Variables->quietProcess) vfprintf(stdout, format, args); + vfprintf(pstdout, format, args); va_end(args); } void println(const char *format, ...) { va_list args; va_start(args, format); - if (!Variables->quietProcess) { - vfprintf(stdout, format, args); - print("\n"); - } + vfprintf(pstdout, format, args); + print("\n"); va_end(args); }