fix: add cleanup, remove unneeded fn, add event listeners
This commit is contained in:
+37
-17
@@ -15,6 +15,21 @@ bool running = true;
|
|||||||
TinyMITMProxy* proxy;
|
TinyMITMProxy* proxy;
|
||||||
ProxyConfigurator* conf;
|
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
|
#ifdef _WIN32
|
||||||
#define _AMD64_
|
#define _AMD64_
|
||||||
#include <consoleapi.h>
|
#include <consoleapi.h>
|
||||||
@@ -29,36 +44,35 @@ BOOL WINAPI consoleHandler(DWORD dwType)
|
|||||||
if (dwType == CTRL_C_EVENT || dwType == CTRL_CLOSE_EVENT || dwType == CTRL_LOGOFF_EVENT ||
|
if (dwType == CTRL_C_EVENT || dwType == CTRL_CLOSE_EVENT || dwType == CTRL_LOGOFF_EVENT ||
|
||||||
dwType == CTRL_SHUTDOWN_EVENT)
|
dwType == CTRL_SHUTDOWN_EVENT)
|
||||||
{
|
{
|
||||||
if (conf) conf->clearProxy();
|
cleanup();
|
||||||
running = false;
|
running = false;
|
||||||
exit(0);
|
return TRUE;
|
||||||
}
|
}
|
||||||
return FALSE;
|
return FALSE;
|
||||||
}
|
}
|
||||||
|
#endif
|
||||||
|
|
||||||
void createAndSetupConsole()
|
void reqListener(const std ::string& url, std::string& body, std::string& headers, bool& blockOutgoing)
|
||||||
{
|
{
|
||||||
FILE* pstdout = stdout;
|
url;
|
||||||
AllocConsole();
|
body;
|
||||||
|
headers;
|
||||||
freopen_s(&pstdout, "CONOUT$", "w", stdout);
|
blockOutgoing;
|
||||||
assert(pstdout != nullptr);
|
|
||||||
|
|
||||||
DWORD conMode = 0;
|
|
||||||
GetConsoleMode(pstdout, &conMode);
|
|
||||||
conMode |= ENABLE_VIRTUAL_TERMINAL_PROCESSING;
|
|
||||||
SetConsoleMode(pstdout, conMode);
|
|
||||||
|
|
||||||
SetConsoleCtrlHandler(consoleHandler, TRUE);
|
|
||||||
}
|
}
|
||||||
|
|
||||||
#endif
|
void resListener(const std::string& host, std::string& body, std::string& headers, bool wasBlocked) {
|
||||||
|
host;
|
||||||
|
body;
|
||||||
|
headers;
|
||||||
|
wasBlocked;
|
||||||
|
}
|
||||||
|
|
||||||
int main()
|
int main()
|
||||||
{
|
{
|
||||||
#ifdef _WIN32
|
#ifdef _WIN32
|
||||||
createAndSetupConsole();
|
SetConsoleCtrlHandler(consoleHandler, TRUE);
|
||||||
#endif
|
#endif
|
||||||
|
std::atexit(cleanup);
|
||||||
|
|
||||||
seallib::Logger mainLog("Main");
|
seallib::Logger mainLog("Main");
|
||||||
mainLog.addSink(std::make_shared<ConOutSink>());
|
mainLog.addSink(std::make_shared<ConOutSink>());
|
||||||
@@ -68,6 +82,7 @@ int main()
|
|||||||
mainLog.info("creating proxy");
|
mainLog.info("creating proxy");
|
||||||
TinyMITMConfig proxyConfig;
|
TinyMITMConfig proxyConfig;
|
||||||
proxy = new TinyMITMProxy(proxyConfig);
|
proxy = new TinyMITMProxy(proxyConfig);
|
||||||
|
proxy->addLogSink(std::make_shared<ConOutSink>());
|
||||||
|
|
||||||
mainLog.info("creating configurator");
|
mainLog.info("creating configurator");
|
||||||
conf = new ProxyConfigurator();
|
conf = new ProxyConfigurator();
|
||||||
@@ -78,8 +93,13 @@ int main()
|
|||||||
mainLog.info("starting proxy");
|
mainLog.info("starting proxy");
|
||||||
proxy->init();
|
proxy->init();
|
||||||
|
|
||||||
|
proxy->onClientRequest.addListener(reqListener);
|
||||||
|
proxy->onServerResponse.addListener(resListener);
|
||||||
|
|
||||||
while (running)
|
while (running)
|
||||||
std::this_thread::sleep_for(std::chrono::milliseconds(100));
|
std::this_thread::sleep_for(std::chrono::milliseconds(100));
|
||||||
|
|
||||||
|
cleanup();
|
||||||
|
|
||||||
return 0;
|
return 0;
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user