export_payload_dialog.py

import asyncio
import qtinter
import ipaddress

from PyQt5.QtWidgets import *
from PyQt5.QtCore import *
from PyQt5.QtGui import *

class ExportPayloadDialog( QObject ):
    """
    Creates a dialog for exporting a payload.
    """
    resetup_dialog = pyqtSignal();

    def __init__( self, parent, rpc_object ):
        """
        Sets up the user interface and signals.
        """
        super( QObject, self ).__init__( parent );

        # set the rpc object
        self.rpccli = rpc_object

        # Force all labels to be vertically placed downwards
        layout_labels = QVBoxLayout();
        layout_labels.addWidget( QLabel( 'IPv4 Address' ) )
        layout_labels.addWidget( QLabel( 'x64' ) )
        layout_labels.addWidget( QLabel( 'Sleep Time' ) )
        layout_labels.addWidget( QLabel( 'Chunk Size' ) )
        layout_labels.addWidget( QLabel( 'Kill Date' ) )
        layout_labels.setSpacing( 6 );

        # Force all inputs to be vertically placed downwards
        layout_inputs = QVBoxLayout();
        
        # address for filling in the IPv4 address
        self.ipv4_address = QLineEdit();
        self.ipv4_address.setPlaceholderText( 'e.g. 192.168.66.11' )

        # is this an x64 arch or is it x86?
        self.is_arch_64   = QCheckBox();
        
        # Sleep time. Max 24 hours or less
        self.sleep_time   = QSpinBox();
        self.sleep_time.setRange( 1000, 86400000 );

        # set allows us to configure the chunk size
        self.chunk_size   = QSpinBox();
        self.chunk_size.setRange( 100, 1400 );
        
        # set the kill date to kill the agent
        self.kill_date    = QDateEdit();

        layout_inputs.addWidget( self.ipv4_address );
        layout_inputs.addWidget( self.is_arch_64 );
        layout_inputs.addWidget( self.sleep_time );
        layout_inputs.addWidget( self.chunk_size );
        layout_inputs.addWidget( self.kill_date );
        layout_inputs.setSpacing( 6 );

        # create the labels next to the inputs!
        layout_labels_inputs = QHBoxLayout();
        layout_labels_inputs.addLayout( layout_labels );
        layout_labels_inputs.addLayout( layout_inputs );
        layout_labels_inputs.setSpacing( 6 );

        # create the button and the spacer
        layout_spacer_button = QHBoxLayout();
        spacer = QSpacerItem( 40, 20, QSizePolicy.Expanding, QSizePolicy.Minimum );
        self.generate_button = QPushButton( "Generate" );
        self.generate_button.clicked.connect( self._generate );
        self.generate_button.setEnabled( False );

        layout_spacer_button.addItem( spacer );
        layout_spacer_button.addWidget( self.generate_button );
        layout_spacer_button.setSpacing( 6 );

        # create a vertical layout to hold the labels/inputs + line + spacer/button 
        layout_form = QVBoxLayout();
        layout_form.addLayout( layout_labels_inputs );
        layout_form.addLayout( layout_spacer_button );
        layout_form.setSpacing( 6 );

        # create the output dialog
        layout_dialog = QGridLayout();
        layout_dialog.addLayout( layout_form, 0, 0, 1, 1 );
        layout_dialog.setSpacing( 6 );
        
        # set the layout
        parent.setLayout( layout_dialog );

        # set the title
        parent.setWindowTitle( 'Payload Configuration' );

        # resize the window
        parent.resize( 480, 167 );

        # set the dialog
        self.dialog = parent

        # connect the resetup signal to resetup
        self.resetup_dialog.connect( self._resetup );

        # is this lock set?
        self.lock = asyncio.Lock();

        # a timer for checking if the input is filled in
        self.timer = QTimer();
        self.timer.setInterval( 100 );
        self.timer.timeout.connect( self._validate_input );
        self.timer.start();

    def reset( self ):
        """
        Emits the signal to recreate the user interface.
        """
        self.resetup_dialog.emit();

    def _validate_ipv4_address( self ):
        """
        Validates the IPv4 address and returns true on whether it succeeded.
        """
        # the buffer was not filled at all!
        if self.ipv4_address.text() == '':
            return False
        else:
            # has a text valud within the widget
            try:
                # is this an IPv4 address that was passed?
                ipaddress.ip_address( self.ipv4_address.text() );

                # return true
                return True
            except ValueError:
                # No address was passed
                return False

    @qtinter.asyncslot
    async def _validate_input( self ):
        """
        Checks the input of each required variable and enables the button
        if all the fields are properly validated.
        """
        # Is this lock set currently?
        if not self.lock.locked():
            # Attempt to get the lock
            async with self.lock:
                # Is the IPv4 address filled in?
                if self._validate_ipv4_address():
                    # set the button as enabled
                    self.generate_button.setEnabled( True );
                else:
                    # set the button as disabled
                    self.generate_button.setEnabled( False );

    @qtinter.asyncslot
    async def _resetup( self ):
        """
        Regenerages a already generated user interface and resets it up!
        """
        # attempt to get the lock
        async with self.lock:
            # close and refrence the ui!
            self.dialog.hide();
            self.dialog.resize( 480, 167 );

            # clear all the values
            self.ipv4_address.clear();
            self.is_arch_64.setChecked( True ); 
            self.sleep_time.clear();
            self.sleep_time.setValue( 1000 );
            self.chunk_size.clear();
            self.chunk_size.setValue( 100 );
            self.kill_date.clear();

            # disable the button
            self.generate_button.setEnabled( False );

            # reshow the user interface
            self.dialog.show();

    @qtinter.asyncslot
    async def _generate( self ):
        """
        A user has requested to generate a payload. We ask the teamserver to send us an
        exported payload
        """
        # attempt to get the lock
        async with self.lock:
            # close the dialog
            self.dialog.close();

            # export a payload and save it to the disk if we can
            shell_code = await self.rpccli.export_payload( self.ipv4_address.text(), self.is_arch_64.isChecked(), self.sleep_time.value() );

            # do we have a buffer available
            if shell_code:
                # save it to the disk so that we can do shit with it.
                path_to_file, _ = QFileDialog.getSaveFileName( self.dialog, "Save Payload",  "", "Binary Files (*.bin)", options = QFileDialog.Options() | QFileDialog.DontUseNativeDialog );

                # Did we get a file path?
                if path_to_file:
                    # save the path to the file!
                    with open( path_to_file, 'wb+' ) as open_file:
                        # write the shellcode!
                        open_file.write( shell_code );