chenx221/Textractor_test
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

refactor

Browse Source
This commit is contained in:
Akash Mozumdar 2018-10-08 00:26:43 -04:00
parent 9929c3fac1
commit 5a464fc083
6 changed files with 50 additions and 47 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
GUI/host/host.cc
View File

@ -26,7 +26,7 @@ namespace
std::recursive_mutex hostMutex; std::recursive_mutex hostMutex;
DWORD DUMMY[100]; DWORD DUMMY[1];
ThreadParam CONSOLE{ 0, -1ULL, -1ULL, -1ULL }; ThreadParam CONSOLE{ 0, -1ULL, -1ULL, -1ULL };
void DispatchText(ThreadParam tp, const BYTE* text, int len) void DispatchText(ThreadParam tp, const BYTE* text, int len)
@ -141,10 +141,9 @@ namespace
namespace Host 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; OnAttach = onAttach; OnDetach = onDetach; OnCreate = onCreate; OnRemove = onRemove; TextThread::Output = output;
TextThread::Prefilter = (int(*)(wchar_t*, const wchar_t*))GetProcAddress(LoadLibraryW(L"0_Prefilter.dll"), "OnNewData");
OnCreate(textThreadsByParams[CONSOLE] = new TextThread(CONSOLE, USING_UNICODE)); OnCreate(textThreadsByParams[CONSOLE] = new TextThread(CONSOLE, USING_UNICODE));
StartPipe(); StartPipe();
} }
@ -216,14 +215,14 @@ namespace Host
void InsertHook(DWORD pid, HookParam hp, std::string name) void InsertHook(DWORD pid, HookParam hp, std::string name)
{ {
auto info = InsertHookCmd(hp, name); auto command = InsertHookCmd(hp, name);
WriteFile(processRecordsByIds[pid].hostPipe, &info, sizeof(info), DUMMY, nullptr); WriteFile(processRecordsByIds[pid].hostPipe, &command, sizeof(command), DUMMY, nullptr);
} }
void RemoveHook(DWORD pid, uint64_t addr) void RemoveHook(DWORD pid, uint64_t addr)
{ {
auto info = RemoveHookCmd(addr); auto command = RemoveHookCmd(addr);
WriteFile(processRecordsByIds[pid].hostPipe, &info, sizeof(info), DUMMY, nullptr); WriteFile(processRecordsByIds[pid].hostPipe, &command, sizeof(command), DUMMY, nullptr);
} }
HookParam GetHookParam(DWORD pid, uint64_t addr) HookParam GetHookParam(DWORD pid, uint64_t addr)
@ -268,7 +267,7 @@ namespace Host
return textThreadsByParams[tp]; return textThreadsByParams[tp];
} }
void AddConsoleOutput(std::wstring text) { GetThread(CONSOLE)->AddSentence(std::wstring(text)); } void AddConsoleOutput(std::wstring text) { GetThread(CONSOLE)->AddSentence(text); }
} }
// EOF // EOF

2
GUI/host/host.h
View File

@ -12,7 +12,7 @@ typedef std::function<void(TextThread*)> ThreadEventCallback;
namespace Host 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(); void Close();
bool InjectProcess(DWORD pid, DWORD timeout = 5000); bool InjectProcess(DWORD pid, DWORD timeout = 5000);

17
GUI/host/textthread.cc
View File

@ -15,7 +15,7 @@ TextThread::~TextThread()
CloseHandle(deletionEvent); CloseHandle(deletionEvent);
} }
std::wstring TextThread::GetStore() std::wstring TextThread::GetStorage()
{ {
LOCK(ttMutex); LOCK(ttMutex);
return storage; return storage;
@ -36,19 +36,20 @@ void TextThread::Flush()
void TextThread::AddSentence(std::wstring sentence) void TextThread::AddSentence(std::wstring sentence)
{ {
// Dispatch to extensions occurs here. Don't hold mutex! Extensions might take a while! // Dispatch to extensions occurs here. Don't hold mutex! Extensions might take a while!
if (Output) sentence = Output(this, sentence); if (Output(this, sentence))
LOCK(ttMutex); {
storage.append(sentence); LOCK(ttMutex);
storage += sentence;
}
} }
void TextThread::AddText(const BYTE* data, int len) 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) ? std::wstring((wchar_t*)data, len / 2)
: StringToWideString(std::string((char*)data, len), status & USING_UTF8 ? CP_UTF8 : SHIFT_JIS); : StringToWideString(std::string((char*)data, len), status & USING_UTF8 ? CP_UTF8 : SHIFT_JIS);
LOCK(ttMutex); if (Filter) Filter(buffer.data(), storage.c_str());
if (Prefilter) wData.resize(Prefilter(wData.data(), storage.c_str()));
buffer.append(wData);
timestamp = GetTickCount(); timestamp = GetTickCount();
} }

23
GUI/host/textthread.h
View File

@ -9,39 +9,38 @@
class TextThread class TextThread
{ {
typedef std::function<std::wstring(TextThread*, std::wstring)> ThreadOutputCallback;
public: public:
typedef std::function<bool(TextThread*, std::wstring&)> OutputCallback;
inline static OutputCallback Output;
inline static std::function<void(wchar_t* buffer, const wchar_t* storage)> 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(ThreadParam tp, DWORD status);
~TextThread(); ~TextThread();
std::wstring GetStore(); std::wstring GetStorage();
void AddText(const BYTE* data, int len); void AddText(const BYTE* data, int len);
void AddSentence(std::wstring sentence); void AddSentence(std::wstring sentence);
void RegisterOutputCallBack(ThreadOutputCallback cb) { Output = cb; }
const int64_t handle; const int64_t handle;
const std::wstring name; const std::wstring name;
const ThreadParam tp; 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: private:
void Flush(); void Flush();
std::wstring buffer; std::wstring buffer;
std::wstring storage; std::wstring storage;
std::recursive_mutex ttMutex; std::recursive_mutex ttMutex;
DWORD status;
HANDLE deletionEvent = CreateEventW(nullptr, FALSE, FALSE, NULL); HANDLE deletionEvent = CreateEventW(nullptr, FALSE, FALSE, NULL);
std::thread flushThread = std::thread([&] { while (WaitForSingleObject(deletionEvent, 10) == WAIT_TIMEOUT) Flush(); }); std::thread flushThread = std::thread([&] { while (WaitForSingleObject(deletionEvent, 10) == WAIT_TIMEOUT) Flush(); });
DWORD timestamp = GetTickCount(); DWORD timestamp = GetTickCount();
ThreadOutputCallback Output;
DWORD status;
}; };
// EOF // EOF

29
GUI/mainwindow.cpp
View File

@ -30,7 +30,8 @@ MainWindow::MainWindow(QWidget *parent) :
[&](DWORD processId) { emit SigAddProcess(processId); }, [&](DWORD processId) { emit SigAddProcess(processId); },
[&](DWORD processId) { emit SigRemoveProcess(processId); }, [&](DWORD processId) { emit SigRemoveProcess(processId); },
[&](TextThread* thread) { emit SigAddThread(thread); }, [&](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(); ReloadExtensions();
@ -47,7 +48,7 @@ MainWindow::~MainWindow()
delete ui; delete ui;
} }
void MainWindow::AddProcess(unsigned int processId) void MainWindow::AddProcess(unsigned processId)
{ {
processCombo->addItem(QString::number(processId, 16).toUpper() + ": " + GetModuleName(processId)); processCombo->addItem(QString::number(processId, 16).toUpper() + ": " + GetModuleName(processId));
QFile file("SavedHooks.txt"); 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)); 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) + 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) 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) QString MainWindow::TextThreadString(TextThread* thread)
{ {
ThreadParam tp = thread->tp; ThreadParam tp = thread->tp;
@ -234,7 +237,7 @@ void MainWindow::on_saveButton_clicked()
void MainWindow::on_ttCombo_activated(int index) 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); textOutput->moveCursor(QTextCursor::End);
} }

9
GUI/mainwindow.h
View File

@ -22,15 +22,15 @@ public:
~MainWindow(); ~MainWindow();
signals: signals:
void SigAddProcess(unsigned int processId); void SigAddProcess(unsigned processId);
void SigRemoveProcess(unsigned int processId); void SigRemoveProcess(unsigned processId);
void SigAddThread(TextThread* thread); void SigAddThread(TextThread* thread);
void SigRemoveThread(TextThread* thread); void SigRemoveThread(TextThread* thread);
void SigThreadOutput(TextThread* thread, QString output); void SigThreadOutput(TextThread* thread, QString output);
private slots: private slots:
void AddProcess(unsigned int processId); void AddProcess(unsigned processId);
void RemoveProcess(unsigned int processId); void RemoveProcess(unsigned processId);
void AddThread(TextThread* thread); void AddThread(TextThread* thread);
void RemoveThread(TextThread* thread); void RemoveThread(TextThread* thread);
void ThreadOutput(TextThread* thread, QString output); void ThreadOutput(TextThread* thread, QString output);
@ -44,6 +44,7 @@ private slots:
void on_rmvExtenButton_clicked(); void on_rmvExtenButton_clicked();
private: private:
bool ProcessThreadOutput(TextThread* thread, std::wstring& output);
QString TextThreadString(TextThread* thread); QString TextThreadString(TextThread* thread);
ThreadParam ParseTextThreadString(QString textThreadString); ThreadParam ParseTextThreadString(QString textThreadString);
DWORD GetSelectedProcessId(); DWORD GetSelectedProcessId();