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;
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

2
GUI/host/host.h
View File

@ -12,7 +12,7 @@ typedef std::function<void(TextThread*)> 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);

17
GUI/host/textthread.cc
View File

@ -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();
}

23
GUI/host/textthread.h
View File

@ -9,39 +9,38 @@
class TextThread
{
typedef std::function<std::wstring(TextThread*, std::wstring)> ThreadOutputCallback;
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();
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

29
GUI/mainwindow.cpp
View File

@ -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);
}

9
GUI/mainwindow.h
View File

@ -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();