diff --git a/src/test/main.cpp b/src/test/main.cpp index df0c166..dee3658 100644 --- a/src/test/main.cpp +++ b/src/test/main.cpp @@ -15,6 +15,21 @@ bool running = true; TinyMITMProxy* proxy; ProxyConfigurator* conf; +void cleanup() +{ + static bool cleaned = false; + if (cleaned) return; + cleaned = true; + + if (conf) conf->clearProxy(); + if (proxy) + { + proxy->shutdown(); + delete proxy; + proxy = nullptr; + } +} + #ifdef _WIN32 #define _AMD64_ #include @@ -29,36 +44,35 @@ BOOL WINAPI consoleHandler(DWORD dwType) if (dwType == CTRL_C_EVENT || dwType == CTRL_CLOSE_EVENT || dwType == CTRL_LOGOFF_EVENT || dwType == CTRL_SHUTDOWN_EVENT) { - if (conf) conf->clearProxy(); + cleanup(); running = false; - exit(0); + return TRUE; } return FALSE; } +#endif -void createAndSetupConsole() +void reqListener(const std ::string& url, std::string& body, std::string& headers, bool& blockOutgoing) { - FILE* pstdout = stdout; - AllocConsole(); - - freopen_s(&pstdout, "CONOUT$", "w", stdout); - assert(pstdout != nullptr); - - DWORD conMode = 0; - GetConsoleMode(pstdout, &conMode); - conMode |= ENABLE_VIRTUAL_TERMINAL_PROCESSING; - SetConsoleMode(pstdout, conMode); - - SetConsoleCtrlHandler(consoleHandler, TRUE); + url; + body; + headers; + blockOutgoing; } -#endif +void resListener(const std::string& host, std::string& body, std::string& headers, bool wasBlocked) { + host; + body; + headers; + wasBlocked; +} int main() { #ifdef _WIN32 - createAndSetupConsole(); + SetConsoleCtrlHandler(consoleHandler, TRUE); #endif + std::atexit(cleanup); seallib::Logger mainLog("Main"); mainLog.addSink(std::make_shared()); @@ -68,6 +82,7 @@ int main() mainLog.info("creating proxy"); TinyMITMConfig proxyConfig; proxy = new TinyMITMProxy(proxyConfig); + proxy->addLogSink(std::make_shared()); mainLog.info("creating configurator"); conf = new ProxyConfigurator(); @@ -78,8 +93,13 @@ int main() mainLog.info("starting proxy"); proxy->init(); + proxy->onClientRequest.addListener(reqListener); + proxy->onServerResponse.addListener(resListener); + while (running) std::this_thread::sleep_for(std::chrono::milliseconds(100)); + cleanup(); + return 0; }