Skip to content

Commit

Permalink
Merge pull request #407 from p4p1/dev
Browse files Browse the repository at this point in the history 
Python API custom widget addition
  • Loading branch information
@S4ntiagoP
S4ntiagoP authored Oct 10, 2023
2 parents 0b0d276 + 59e1729 commit 2627c88
Show file tree
Hide file tree
Showing 4 changed files with 315 additions and 1 deletion.
2 changes: 2 additions & 0 deletions client/CMakeLists.txt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -68,6 +68,7 @@ set(
Include/Havoc/PythonApi/Event.h
Include/Havoc/PythonApi/HavocUi.h
Include/Havoc/PythonApi/PyAgentClass.hpp
Include/Havoc/PythonApi/PyWidgetClass.hpp

# Dialogs
Include/UserInterface/Dialogs/Payload.hpp
Expand Down Expand Up @@ -153,6 +154,7 @@ set(
Source/Havoc/PythonApi/PythonApi.cpp
Source/Havoc/PythonApi/Havoc.cpp
Source/Havoc/PythonApi/HavocUi.cpp
Source/Havoc/PythonApi/PyWidgetClass.cpp
Source/Havoc/PythonApi/PyDemonClass.cpp
Source/Havoc/PythonApi/Event.cpp
Source/Havoc/Havoc.cpp
Expand Down
50 changes: 50 additions & 0 deletions client/Include/Havoc/PythonApi/PyWidgetClass.hpp
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,50 @@
#ifndef HAVOC_PYWIDGETCLASS_H
#define HAVOC_PYWIDGETCLASS_H

#include <UserInterface/HavocUI.hpp>
#include <global.hpp>

#include <QVBoxLayout>
#include <QDialog>
#include <QLabel>
#include <QPushButton>
#include <QCheckBox>
#include <QLineEdit>

typedef struct
{
PyObject_HEAD

// Demon Info
char* title;

} PyWidgetClass, *PPyWidgetClass;

typedef struct
{

QDialog* window;
QVBoxLayout* layout;

} PyWidgetQWindow, *PPyWidgetQWindow;

extern PyTypeObject PyWidgetClass_Type;

void WidgetClass_dealloc( PPyWidgetClass self );
PyObject* WidgetClass_new( PyTypeObject *type, PyObject *args, PyObject *kwds );
int WidgetClass_init( PPyWidgetClass self, PyObject *args, PyObject *kwds );

// Methods

// PyObject* DemonClass_( PPyDemonClass self, PyObject *args );

PyObject* WidgetClass_exec( PPyWidgetClass self, PyObject *args );
PyObject* WidgetClass_close( PPyWidgetClass self, PyObject *args );
PyObject* WidgetClass_addLabel( PPyWidgetClass self, PyObject *args );
PyObject* WidgetClass_addButton( PPyWidgetClass self, PyObject *args );
PyObject* WidgetClass_addCheckbox( PPyWidgetClass self, PyObject *args );
PyObject* WidgetClass_addCombobox( PPyWidgetClass self, PyObject *args );
PyObject* WidgetClass_addLineedit( PPyWidgetClass self, PyObject *args );
PyObject* WidgetClass_ConsoleWrite( PPyWidgetClass self, PyObject *args );

#endif
11 changes: 10 additions & 1 deletion client/Source/Havoc/PythonApi/HavocUi.cpp
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,8 @@
#include <Havoc/PythonApi/PythonApi.h>
#include <UserInterface/HavocUI.hpp>

#include <Havoc/PythonApi/PyWidgetClass.hpp>

#include <QFile>
#include <QMessageBox>
#include <QFileDialog>
Expand Down Expand Up @@ -151,5 +154,11 @@ PyObject* PythonAPI::HavocUI::Core::SaveFileDialog(PyObject *self, PyObject *arg

PyMODINIT_FUNC PythonAPI::HavocUI::PyInit_HavocUI(void)
{
return PyModule_Create( &PythonAPI::HavocUI::PyModule::havocui );
PyObject* Module = PyModule_Create2( &PythonAPI::HavocUI::PyModule::havocui, PYTHON_API_VERSION );

if ( PyType_Ready( &PyWidgetClass_Type ) < 0 )
spdlog::error( "Couldn't check if WidgetClass is ready" );
else
PyModule_AddObject( Module, "Widget", (PyObject*) &PyWidgetClass_Type );
return Module;
}
253 changes: 253 additions & 0 deletions client/Source/Havoc/PythonApi/PyWidgetClass.cpp
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,253 @@

#define PY_SSIZE_T_CLEAN
#include <Python.h>
#include <structmember.h>

#include <Havoc/PythonApi/PythonApi.h>
#include <Havoc/PythonApi/PyWidgetClass.hpp>

PPyWidgetQWindow DialogWindow = NULL;

PyMemberDef PyWidgetClass_members[] = {

{ "title", T_STRING, offsetof( PyWidgetClass, title ), 0, "title" },

{ NULL },
};

PyMethodDef PyWidgetClass_methods[] = {

{ "exec", ( PyCFunction ) WidgetClass_exec, METH_VARARGS, "Display the window" },
{ "close", ( PyCFunction ) WidgetClass_close, METH_VARARGS, "Close the window" },
{ "addLabel", ( PyCFunction ) WidgetClass_addLabel, METH_VARARGS, "Insert a label in the window" },
{ "addButton", ( PyCFunction ) WidgetClass_addButton, METH_VARARGS, "Insert a button in the window" },
{ "addCheckbox", ( PyCFunction ) WidgetClass_addCheckbox, METH_VARARGS, "Insert a checkbox in the window" },
{ "addCombobox", ( PyCFunction ) WidgetClass_addCombobox, METH_VARARGS, "Insert a checkbox in the window" },
{ "addLineedit", ( PyCFunction ) WidgetClass_addLineedit, METH_VARARGS, "Insert a Line edit in the window" },
{ "ConsoleWrite", ( PyCFunction ) WidgetClass_ConsoleWrite, METH_VARARGS, "Print message console" },

{ NULL },
};

PyTypeObject PyWidgetClass_Type = {
PyVarObject_HEAD_INIT( &PyType_Type, 0 )

"havocui.widget", /* tp_name */
sizeof( PyWidgetClass ), /* tp_basicsize */
0, /* tp_itemsize */
( destructor ) WidgetClass_dealloc, /* tp_dealloc */
0, /* tp_print */
0, /* tp_getattr */
0, /* tp_setattr */
0, /* tp_reserved */
0, /* tp_repr */
0, /* tp_as_number */
0, /* tp_as_sequence */
0, /* tp_as_mapping */
0, /* tp_hash */
0, /* tp_call */
0, /* tp_str */
0, /* tp_getattro */
0, /* tp_setattro */
0, /* tp_as_buffer */
Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE, /* tp_flags */
"Widget Havoc Object", /* tp_doc */
0, /* tp_traverse */
0, /* tp_clear */
0, /* tp_richcompare */
0, /* tp_weaklistoffset */
0, /* tp_iter */
0, /* tp_iternext */
PyWidgetClass_methods, /* tp_methods */
PyWidgetClass_members, /* tp_members */
0, /* tp_getset */
0, /* tp_base */
0, /* tp_dict */
0, /* tp_descr_get */
0, /* tp_descr_set */
0, /* tp_dictoffset */
( initproc ) WidgetClass_init, /* tp_init */
0, /* tp_alloc */
WidgetClass_new, /* tp_new */
};

#define AllocMov( des, src, size ) \
if ( size > 0 ) \
{ \
des = ( char* ) malloc( size * sizeof( char ) ); \
memset( des, 0, size ); \
std::strcpy( des, src ); \
}

void WidgetClass_dealloc( PPyWidgetClass self )
{
Py_XDECREF( self->title );
delete DialogWindow->window;
free(DialogWindow);

Py_TYPE( self )->tp_free( ( PyObject* ) self );
}

PyObject* WidgetClass_new( PyTypeObject *type, PyObject *args, PyObject *kwds )
{
PPyWidgetClass self;

self = ( PPyWidgetClass ) PyType_Type.tp_alloc( type, 0 );

return ( PyObject* ) self;
}

int WidgetClass_init( PPyWidgetClass self, PyObject *args, PyObject *kwds )
{
if ( PyType_Type.tp_init( ( PyObject* ) self, args, kwds ) < 0 )
return -1;

char* title = NULL;
const char* kwdlist[] = { "title", NULL };

if ( ! PyArg_ParseTupleAndKeywords( args, kwds, "s", const_cast<char**>(kwdlist), &title ) )
return -1;
AllocMov( self->title, title, strlen(title) );
DialogWindow = (PPyWidgetQWindow)malloc(sizeof(PyWidgetQWindow));
if (DialogWindow == NULL)
return -1;
DialogWindow->window = new QDialog(HavocX::HavocUserInterface->HavocWindow);
DialogWindow->window->setWindowTitle(title);
DialogWindow->layout = new QVBoxLayout(DialogWindow->window);

return 0;
}

// Methods
PyObject* WidgetClass_exec( PPyWidgetClass self, PyObject *args )
{
DialogWindow->window->exec();

Py_RETURN_NONE;
}

PyObject* WidgetClass_addLabel( PPyWidgetClass self, PyObject *args )
{
char *text = nullptr;

if( !PyArg_ParseTuple( args, "s", &text) )
{
Py_RETURN_NONE;
}
QLabel* label = new QLabel(text, DialogWindow->window);
DialogWindow->layout->addWidget(label);

Py_RETURN_NONE;
}

PyObject* WidgetClass_addButton( PPyWidgetClass self, PyObject *args )
{
char *text = nullptr;
PyObject* button_callback = nullptr;

if( !PyArg_ParseTuple( args, "sO", &text, &button_callback) )
{
Py_RETURN_NONE;
}
if ( !PyCallable_Check(button_callback) )
{
PyErr_SetString(PyExc_TypeError, "parameter must be callable");
return NULL;
}
QPushButton* button = new QPushButton(text, DialogWindow->window);
DialogWindow->layout->addWidget(button);
QObject::connect(button, &QPushButton::clicked, DialogWindow->window, [button_callback]() {
PyObject_CallFunctionObjArgs(button_callback, nullptr);
});

Py_RETURN_NONE;
}

PyObject* WidgetClass_addCheckbox( PPyWidgetClass self, PyObject *args )
{
char *text = nullptr;
PyObject* checkbox_callback = nullptr;

if( !PyArg_ParseTuple( args, "sO", &text, &checkbox_callback) )
{
Py_RETURN_NONE;
}
if ( !PyCallable_Check(checkbox_callback) )
{
PyErr_SetString(PyExc_TypeError, "parameter must be callable");
return NULL;
}
QCheckBox* checkbox = new QCheckBox(text, DialogWindow->window);
DialogWindow->layout->addWidget(checkbox);
QObject::connect(checkbox, &QCheckBox::clicked, DialogWindow->window, [checkbox_callback]() {
PyObject_CallFunctionObjArgs(checkbox_callback, nullptr);
});

Py_RETURN_NONE;
}

PyObject* WidgetClass_addCombobox( PPyWidgetClass self, PyObject *args )
{
Py_ssize_t tuple_size = PyTuple_Size(args);
QComboBox* comboBox = new QComboBox(DialogWindow->window);

PyObject* callable_obj = PyTuple_GetItem(args, 0);
if ( !PyCallable_Check(callable_obj) )
{
PyErr_SetString(PyExc_TypeError, "parameter must be callable");
return NULL;
}
for (Py_ssize_t i = 1; i < tuple_size; i++) {
const char * string_obj = PyUnicode_AsUTF8(PyTuple_GetItem(args, i));

comboBox->addItem(string_obj);
}
DialogWindow->layout->addWidget(comboBox);
QObject::connect(comboBox, QOverload<int>::of(&QComboBox::activated), [callable_obj](int index) {
PyObject* pArg = PyLong_FromLong(index);
PyObject_CallFunctionObjArgs(callable_obj, pArg, nullptr);
});
Py_RETURN_NONE;
}

PyObject* WidgetClass_addLineedit( PPyWidgetClass self, PyObject *args )
{
char *text = nullptr;
PyObject* line_callback = nullptr;

if( !PyArg_ParseTuple( args, "sO", &text, &line_callback) )
{
Py_RETURN_NONE;
}
if ( !PyCallable_Check(line_callback) )
{
PyErr_SetString(PyExc_TypeError, "parameter must be callable");
return NULL;
}
QLineEdit* line = new QLineEdit(DialogWindow->window);
line->setPlaceholderText(text);
DialogWindow->layout->addWidget(line);
QObject::connect(line, &QLineEdit::editingFinished, DialogWindow->window, [line, line_callback]() {
QString text = line->text();
QByteArray byteArray = text.toUtf8();
char *charArray = byteArray.data();
PyObject* pyString = PyUnicode_DecodeFSDefault(charArray);
PyObject_CallFunctionObjArgs(line_callback, pyString, nullptr);
});

Py_RETURN_NONE;
}

PyObject* WidgetClass_close( PPyWidgetClass self, PyObject *args )
{
DialogWindow->window->accept();

Py_RETURN_NONE;
}

PyObject* WidgetClass_ConsoleWrite( PPyWidgetClass self, PyObject *args )
{
printf("%s\n", self->title);

Py_RETURN_NONE;
}

0 comments on commit 2627c88

Please sign in to comment.