From 5a464fc083b90a2a6175855ba35a95faa59377ee Mon Sep 17 00:00:00 2001 From: Akash Mozumdar Date: Mon, 8 Oct 2018 00:26:43 -0400 Subject: [PATCH] refactor --- GUI/host/host.cc | 17 ++++++++--------- GUI/host/host.h | 2 +- GUI/host/textthread.cc | 17 +++++++++-------- GUI/host/textthread.h | 23 +++++++++++------------ GUI/mainwindow.cpp | 29 ++++++++++++++++------------- GUI/mainwindow.h | 9 +++++---- 6 files changed, 50 insertions(+), 47 deletions(-) diff --git a/GUI/host/host.cc b/GUI/host/host.cc index 80410a3..7d566e5 100644 --- a/GUI/host/host.cc +++ b/GUI/host/host.cc @@ -26,7 +26,7 @@ namespace std::recursive_mutex hostMutex; - DWORD DUMMY[100]; + DWORD DUMMY[1]; ThreadParam CONSOLE{ 0, -1ULL, -1ULL, -1ULL }; void DispatchText(ThreadParam tp, const BYTE* text, int len) @@ -141,10 +141,9 @@ namespace namespace Host { - void Start(ProcessEventCallback onAttach, ProcessEventCallback onDetach, ThreadEventCallback onCreate, ThreadEventCallback onRemove) + void Start(ProcessEventCallback onAttach, ProcessEventCallback onDetach, ThreadEventCallback onCreate, ThreadEventCallback onRemove, TextThread::OutputCallback output) { - OnAttach = onAttach; OnDetach = onDetach; OnCreate = onCreate; OnRemove = onRemove; - TextThread::Prefilter = (int(*)(wchar_t*, const wchar_t*))GetProcAddress(LoadLibraryW(L"0_Prefilter.dll"), "OnNewData"); + OnAttach = onAttach; OnDetach = onDetach; OnCreate = onCreate; OnRemove = onRemove; TextThread::Output = output; OnCreate(textThreadsByParams[CONSOLE] = new TextThread(CONSOLE, USING_UNICODE)); StartPipe(); } @@ -216,14 +215,14 @@ namespace Host void InsertHook(DWORD pid, HookParam hp, std::string name) { - auto info = InsertHookCmd(hp, name); - WriteFile(processRecordsByIds[pid].hostPipe, &info, sizeof(info), DUMMY, nullptr); + auto command = InsertHookCmd(hp, name); + WriteFile(processRecordsByIds[pid].hostPipe, &command, sizeof(command), DUMMY, nullptr); } void RemoveHook(DWORD pid, uint64_t addr) { - auto info = RemoveHookCmd(addr); - WriteFile(processRecordsByIds[pid].hostPipe, &info, sizeof(info), DUMMY, nullptr); + auto command = RemoveHookCmd(addr); + WriteFile(processRecordsByIds[pid].hostPipe, &command, sizeof(command), DUMMY, nullptr); } HookParam GetHookParam(DWORD pid, uint64_t addr) @@ -268,7 +267,7 @@ namespace Host return textThreadsByParams[tp]; } - void AddConsoleOutput(std::wstring text) { GetThread(CONSOLE)->AddSentence(std::wstring(text)); } + void AddConsoleOutput(std::wstring text) { GetThread(CONSOLE)->AddSentence(text); } } // EOF diff --git a/GUI/host/host.h b/GUI/host/host.h index 79d08e0..b27d1a8 100644 --- a/GUI/host/host.h +++ b/GUI/host/host.h @@ -12,7 +12,7 @@ typedef std::function ThreadEventCallback; namespace Host { - void Start(ProcessEventCallback onAttach, ProcessEventCallback onDetach, ThreadEventCallback onCreate, ThreadEventCallback onRemove); + void Start(ProcessEventCallback onAttach, ProcessEventCallback onDetach, ThreadEventCallback onCreate, ThreadEventCallback onRemove, TextThread::OutputCallback output); void Close(); bool InjectProcess(DWORD pid, DWORD timeout = 5000); diff --git a/GUI/host/textthread.cc b/GUI/host/textthread.cc index b518eab..f1d7aab 100644 --- a/GUI/host/textthread.cc +++ b/GUI/host/textthread.cc @@ -15,7 +15,7 @@ TextThread::~TextThread() CloseHandle(deletionEvent); } -std::wstring TextThread::GetStore() +std::wstring TextThread::GetStorage() { LOCK(ttMutex); return storage; @@ -36,19 +36,20 @@ void TextThread::Flush() void TextThread::AddSentence(std::wstring sentence) { // Dispatch to extensions occurs here. Don't hold mutex! Extensions might take a while! - if (Output) sentence = Output(this, sentence); - LOCK(ttMutex); - storage.append(sentence); + if (Output(this, sentence)) + { + LOCK(ttMutex); + storage += sentence; + } } void TextThread::AddText(const BYTE* data, int len) { - std::wstring wData = status & USING_UNICODE + LOCK(ttMutex); + buffer += status & USING_UNICODE ? std::wstring((wchar_t*)data, len / 2) : StringToWideString(std::string((char*)data, len), status & USING_UTF8 ? CP_UTF8 : SHIFT_JIS); - LOCK(ttMutex); - if (Prefilter) wData.resize(Prefilter(wData.data(), storage.c_str())); - buffer.append(wData); + if (Filter) Filter(buffer.data(), storage.c_str()); timestamp = GetTickCount(); } diff --git a/GUI/host/textthread.h b/GUI/host/textthread.h index 868c81e..41fcf05 100644 --- a/GUI/host/textthread.h +++ b/GUI/host/textthread.h @@ -9,39 +9,38 @@ class TextThread { - typedef std::function ThreadOutputCallback; - public: + typedef std::function OutputCallback; + + inline static OutputCallback Output; + inline static std::function Filter = nullptr; + + inline static int FlushDelay = 250; // flush every 250ms by default + inline static int MaxBufferSize = 200; + inline static int ThreadCounter = 0; + TextThread(ThreadParam tp, DWORD status); ~TextThread(); - std::wstring GetStore(); + std::wstring GetStorage(); void AddText(const BYTE* data, int len); void AddSentence(std::wstring sentence); - void RegisterOutputCallBack(ThreadOutputCallback cb) { Output = cb; } const int64_t handle; const std::wstring name; const ThreadParam tp; - inline static int(*Prefilter)(wchar_t*, const wchar_t*) = nullptr; - inline static int FlushDelay = 250; // flush every 250ms by default - inline static int MaxBufferSize = 200; - inline static int ThreadCounter = 0; - private: void Flush(); std::wstring buffer; std::wstring storage; std::recursive_mutex ttMutex; + DWORD status; HANDLE deletionEvent = CreateEventW(nullptr, FALSE, FALSE, NULL); std::thread flushThread = std::thread([&] { while (WaitForSingleObject(deletionEvent, 10) == WAIT_TIMEOUT) Flush(); }); DWORD timestamp = GetTickCount(); - - ThreadOutputCallback Output; - DWORD status; }; // EOF diff --git a/GUI/mainwindow.cpp b/GUI/mainwindow.cpp index 8690ecf..407ed44 100644 --- a/GUI/mainwindow.cpp +++ b/GUI/mainwindow.cpp @@ -30,7 +30,8 @@ MainWindow::MainWindow(QWidget *parent) : [&](DWORD processId) { emit SigAddProcess(processId); }, [&](DWORD processId) { emit SigRemoveProcess(processId); }, [&](TextThread* thread) { emit SigAddThread(thread); }, - [&](TextThread* thread) { emit SigRemoveThread(thread); } + [&](TextThread* thread) { emit SigRemoveThread(thread); }, + [&](TextThread* thread, std::wstring& output) { return ProcessThreadOutput(thread, output); } ); ReloadExtensions(); @@ -47,7 +48,7 @@ MainWindow::~MainWindow() delete ui; } -void MainWindow::AddProcess(unsigned int processId) +void MainWindow::AddProcess(unsigned processId) { processCombo->addItem(QString::number(processId, 16).toUpper() + ": " + GetModuleName(processId)); QFile file("SavedHooks.txt"); @@ -65,7 +66,7 @@ void MainWindow::AddProcess(unsigned int processId) } } -void MainWindow::RemoveProcess(unsigned int processId) +void MainWindow::RemoveProcess(unsigned processId) { processCombo->removeItem(processCombo->findText(QString::number(processId, 16).toUpper() + ":", Qt::MatchStartsWith)); } @@ -79,15 +80,6 @@ void MainWindow::AddThread(TextThread* thread) GenerateCode(Host::GetHookParam(thread->tp), thread->tp.pid) + ")" ); - thread->RegisterOutputCallBack([&](TextThread* thread, std::wstring output) - { - if (DispatchSentenceToExtensions(output, GetInfoForExtensions(thread))) - { - output += L"\r\n"; - emit SigThreadOutput(thread, QString::fromStdWString(output)); - } - return output; - }); } void MainWindow::RemoveThread(TextThread* thread) @@ -112,6 +104,17 @@ void MainWindow::ThreadOutput(TextThread* thread, QString output) } } +bool MainWindow::ProcessThreadOutput(TextThread* thread, std::wstring& output) +{ + if (DispatchSentenceToExtensions(output, GetInfoForExtensions(thread))) + { + output += L"\r\n"; + emit SigThreadOutput(thread, QString::fromStdWString(output)); + return true; + } + return false; +} + QString MainWindow::TextThreadString(TextThread* thread) { ThreadParam tp = thread->tp; @@ -234,7 +237,7 @@ void MainWindow::on_saveButton_clicked() void MainWindow::on_ttCombo_activated(int index) { - textOutput->setPlainText(QString::fromStdWString(Host::GetThread(ParseTextThreadString(ttCombo->itemText(index)))->GetStore())); + textOutput->setPlainText(QString::fromStdWString(Host::GetThread(ParseTextThreadString(ttCombo->itemText(index)))->GetStorage())); textOutput->moveCursor(QTextCursor::End); } diff --git a/GUI/mainwindow.h b/GUI/mainwindow.h index ffa3313..5ca76cb 100644 --- a/GUI/mainwindow.h +++ b/GUI/mainwindow.h @@ -22,15 +22,15 @@ public: ~MainWindow(); signals: - void SigAddProcess(unsigned int processId); - void SigRemoveProcess(unsigned int processId); + void SigAddProcess(unsigned processId); + void SigRemoveProcess(unsigned processId); void SigAddThread(TextThread* thread); void SigRemoveThread(TextThread* thread); void SigThreadOutput(TextThread* thread, QString output); private slots: - void AddProcess(unsigned int processId); - void RemoveProcess(unsigned int processId); + void AddProcess(unsigned processId); + void RemoveProcess(unsigned processId); void AddThread(TextThread* thread); void RemoveThread(TextThread* thread); void ThreadOutput(TextThread* thread, QString output); @@ -44,6 +44,7 @@ private slots: void on_rmvExtenButton_clicked(); private: + bool ProcessThreadOutput(TextThread* thread, std::wstring& output); QString TextThreadString(TextThread* thread); ThreadParam ParseTextThreadString(QString textThreadString); DWORD GetSelectedProcessId();