/*
Copyright (C) 2006 - 2009 Evan Teran
                   eteran@alum.rit.edu

This program is free software; you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation; either version 2 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program; if not, write to the Free Software
Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
*/

#include "DialogAttach.h"
#include "Types.h"

#include <QMap>
#include <QHeaderView>
#include <QDir>
#include <QTextStream>

#include <tlhelp32.h>

struct ProcessInfo {
	edb::pid_t	pid;
	edb::uid_t	uid;
	QString		user;
	QString		name;
};

//------------------------------------------------------------------------------
// Name: DialogAttach(QWidget *parent) : QDialog(parent)
// Desc: constructor
//------------------------------------------------------------------------------
DialogAttach::DialogAttach(QWidget *parent) : QDialog(parent) {
	ui.setupUi(this);
	ui.tableWidget->verticalHeader()->hide();
	ui.tableWidget->horizontalHeader()->setStretchLastSection(true);
}

//------------------------------------------------------------------------------
// Name: isNumeric(const QString &s)
// Desc: returns true if the string only contains decimal digits
//------------------------------------------------------------------------------
bool DialogAttach::isNumeric(const QString &s) {
	Q_FOREACH(QChar ch, s) {
		if(!ch.isDigit()) {
			return false;
		}
	}

	return true;
}

//------------------------------------------------------------------------------
// Name: on_txtSearch_textChanged(const QString &text)
// Desc: 
//------------------------------------------------------------------------------
void DialogAttach::on_txtSearch_textChanged(const QString &text) {
	updateList(text);
}


QString dumpToken(HANDLE hProcess) {
	QString ret;
	HANDLE hToken;
	if (!OpenProcessToken(hProcess, TOKEN_QUERY, &hToken))
		return ret;

	DWORD needed;
	TOKEN_OWNER *tokenOwner = NULL;
	GetTokenInformation(hToken, TokenOwner, NULL, 0, &needed);
	tokenOwner = static_cast<TOKEN_OWNER *>(malloc(needed));

	if(tokenOwner) {
		if (GetTokenInformation(hToken, TokenOwner, tokenOwner, needed, &needed)) {
			WCHAR user[MAX_PATH];
			WCHAR domain[MAX_PATH];
			SID_NAME_USE snu;
			DWORD user_sz = MAX_PATH;
			DWORD domain_sz = MAX_PATH;
			if (LookupAccountSid(NULL, tokenOwner->Owner, user, &user_sz, domain, &domain_sz, &snu) && snu == SidTypeUser) {
				ret = QString::fromWCharArray(user);
			}
		}
		free(tokenOwner);
	}

	CloseHandle(hToken);
	return ret;
}

//------------------------------------------------------------------------------
// Name: updateList(const QString &filter)
// Desc: 
//------------------------------------------------------------------------------
void DialogAttach::updateList(const QString &filter) {

	QMap<edb::pid_t, ProcessInfo> procs;
	edb::uid_t myUID = 0;

	// FILL up procs
	HANDLE handle = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
	if(handle != INVALID_HANDLE_VALUE) {
		PROCESSENTRY32 lppe;
		
		std::memset(&lppe, 0, sizeof(lppe));
		lppe.dwSize = sizeof(lppe);

		if(Process32First(handle, &lppe)) {
			do {
				ProcessInfo pi;
				pi.pid = lppe.th32ProcessID;
				pi.uid = 0; // TODO
				pi.name = QString::fromWCharArray(lppe.szExeFile);

				HANDLE hProc = OpenProcess(PROCESS_QUERY_INFORMATION, FALSE, lppe.th32ProcessID);
				if(hProc != 0) {
					BOOL wow64 = FALSE;
					if(IsWow64Process(hProc, &wow64) && wow64) {
						pi.name += " *32";
					}

					pi.user = dumpToken(hProc);

					CloseHandle(hProc);
				}

				procs[pi.pid] = pi;

				std::memset(&lppe, 0, sizeof(lppe));
				lppe.dwSize = sizeof(lppe);
			} while(Process32Next(handle, &lppe));
		}
		
		CloseHandle(handle);
	}
	
	const bool filterUID = ui.chkFilterUID->isChecked();
	const QString lowerFilter = filter.toLower();
	ui.tableWidget->setSortingEnabled(false);
	ui.tableWidget->setRowCount(0);

	Q_FOREACH(const ProcessInfo &procInfo, procs) {

		const QString procName = procInfo.name;
		
		if(filter.isEmpty() || procName.toLower().contains(lowerFilter)) {	
			if(!filterUID || procInfo.uid == myUID) {
				const int row = ui.tableWidget->rowCount();
				ui.tableWidget->insertRow(row);
				
				QTableWidgetItem *const itemPID = new QTableWidgetItem;
				itemPID->setData(Qt::DisplayRole, static_cast<quint32>(procInfo.pid));

				QTableWidgetItem *itemUID;
				if(!procInfo.user.isEmpty()) {
					itemUID = new QTableWidgetItem(procInfo.user);
				} else {
					itemUID = new QTableWidgetItem;
					itemUID->setData(Qt::DisplayRole, static_cast<quint32>(procInfo.uid));
				}

				ui.tableWidget->setItem(row, 0, itemPID);
				ui.tableWidget->setItem(row, 1, itemUID);
		        ui.tableWidget->setItem(row, 2, new QTableWidgetItem(procInfo.name));
			}
		}
	}
	ui.tableWidget->setSortingEnabled(true);
}

//------------------------------------------------------------------------------
// Name: showEvent(QShowEvent *event)
// Desc: 
//------------------------------------------------------------------------------
void DialogAttach::showEvent(QShowEvent *event) {
	Q_UNUSED(event);
	updateList(ui.txtSearch->text());
}

//------------------------------------------------------------------------------
// Name: on_chkFilterUID_clicked(bool checked)
// Desc: 
//------------------------------------------------------------------------------
void DialogAttach::on_chkFilterUID_clicked(bool checked) {
	Q_UNUSED(checked);
	updateList(ui.txtSearch->text());
}

//------------------------------------------------------------------------------
// Name: 
// Desc: 
//------------------------------------------------------------------------------
edb::pid_t DialogAttach::selectedPID(bool &ok) const {
	const QList<QTableWidgetItem *> sel = ui.tableWidget->selectedItems();
	if(sel.size() != 0) {
		ok = true;
		return sel.first()->text().toUInt();
	} else {
		ok = false;
		return edb::pid_t();
	}
}
