Skip to content

Instantly share code, notes, and snippets.

@realoriginal

realoriginal/agent_proc_list_tab.py

Created September 15, 2023 14:32
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save realoriginal/e3d726e06902764d87183fc7ce4edf78 to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save realoriginal/e3d726e06902764d87183fc7ce4edf78 to your computer and use it in GitHub Desktop.
Download ZIP

Revisions

  1. realoriginal created this gist Sep 15, 2023.
    106 changes: 106 additions & 0 deletions agent_proc_list_tab.py
    View file
    Original file line number Diff line number Diff line change
    @@ -0,0 +1,106 @@
    #!/usr/bin/env python3
    # -*- coding:utf-8 -*-
    import PyQt5
    import qtinter
    import asyncio

    class AgentProcListTab( PyQt5.QtWidgets.QWidget ):
    """
    Tasks the specified agent with requesting a process listing against the
    specified agent and executes the rendered results.
    """
    COLUMN_NAMES = [ "Process", "Session ID", "PID", "PPID", "Username" ]
    COLUMN_COUNT = len( COLUMN_NAMES );

    def __init__( self, ghost, agent_id ):
    super( PyQt5.QtWidgets.QWidget, self ).__init__( ghost );

    # set the ghost object
    self.ghost = ghost

    # set the agent ID
    self.agent_id = agent_id

    # Set the primary layout
    self.layout = PyQt5.QtWidgets.QHBoxLayout();

    # Process table for displaying the process information
    self.process_table = PyQt5.QtWidgets.QTableWidget();
    self.process_table.setShowGrid( False );
    self.process_table.setFocusPolicy( PyQt5.QtCore.Qt.NoFocus );
    self.process_table.setSelectionBehavior( PyQt5.QtWidgets.QTableView.SelectRows );
    self.process_table.setSelectionMode( PyQt5.QtWidgets.QTableView.SingleSelection );
    self.process_table.setRowCount( 0 );
    self.process_table.setColumnCount( self.COLUMN_COUNT );
    self.process_table.setHorizontalHeaderLabels( self.COLUMN_NAMES );
    self.process_table.verticalHeader().setVisible( False );
    self.process_table.horizontalHeader().setHighlightSections( False );

    # Loop through each column
    for i in range( 0, self.COLUMN_COUNT ):
    # Request that the column be stretched to fit the table view
    self.process_table.horizontalHeader().setSectionResizeMode( i, PyQt5.QtWidgets.QHeaderView.Stretch )

    # Setup a one-time startup timer to request the results!
    self.startup_event = PyQt5.QtCore.QTimer()
    self.startup_event.setSingleShot( True );
    self.startup_event.setInterval( 0 );
    self.startup_event.timeout.connect( self._start_event_cb );
    self.startup_event.start()

    # Ad the table to the primary layout
    self.layout.addWidget( self.process_table );

    # Set the primary layout
    self.setLayout( self.layout );

    async def parse_proc_list( self, proc_list ):
    """
    Parses a list of process information dict and renders it within a table.
    """
    # Reset the table
    self.process_table.setRowCount( 0 );

    # Loops through each process
    for proc in proc_list:
    # Updates the row count with one more row
    self.process_table.setRowCount( self.process_table.rowCount() + 1 );

    # Column: "Process"
    item = PyQt5.QtWidgets.QTableWidgetItem( f'{proc[ "process" ]}' );
    item.setTextAlignment( PyQt5.QtCore.Qt.AlignCenter );
    item.setFlags( PyQt5.QtCore.Qt.ItemIsSelectable | PyQt5.QtCore.Qt.ItemIsEnabled );
    self.process_table.setItem( self.process_table.rowCount() - 1, 0, item );

    # Column: "Session ID"
    item = PyQt5.QtWidgets.QTableWidgetItem( f'{proc[ "session_id" ]}' );
    item.setTextAlignment( PyQt5.QtCore.Qt.AlignCenter );
    item.setFlags( PyQt5.QtCore.Qt.ItemIsSelectable | PyQt5.QtCore.Qt.ItemIsEnabled );
    self.process_table.setItem( self.process_table.rowCount() - 1, 1, item );

    # Column: "PID"
    item = PyQt5.QtWidgets.QTableWidgetItem( f'{proc[ "pid" ]}' );
    item.setTextAlignment( PyQt5.QtCore.Qt.AlignCenter );
    item.setFlags( PyQt5.QtCore.Qt.ItemIsSelectable | PyQt5.QtCore.Qt.ItemIsEnabled );
    self.process_table.setItem( self.process_table.rowCount() - 1, 2, item );

    # Column: "PPID"
    item = PyQt5.QtWidgets.QTableWidgetItem( f'{proc[ "ppid" ]}' );
    item.setTextAlignment( PyQt5.QtCore.Qt.AlignCenter );
    item.setFlags( PyQt5.QtCore.Qt.ItemIsSelectable | PyQt5.QtCore.Qt.ItemIsEnabled );
    self.process_table.setItem( self.process_table.rowCount() - 1, 3, item );

    # Column: "Username"
    item = PyQt5.QtWidgets.QTableWidgetItem( f'{proc[ "username" ]}' );
    item.setTextAlignment( PyQt5.QtCore.Qt.AlignCenter );
    item.setFlags( PyQt5.QtCore.Qt.ItemIsSelectable | PyQt5.QtCore.Qt.ItemIsEnabled );
    self.process_table.setItem( self.process_table.rowCount() - 1, 4, item );

    @qtinter.asyncslot
    async def _start_event_cb( self ):
    """
    Requests a process listing from the specified agent on tab creation. Further
    queries are handled via a "refresh" of the results.
    """
    # Request a process listing callback to this specified ID
    await self.ghost.rpc.agent_proc_list( self.agent_id, ( await self.ghost.tab_widget.tab_get_id_from_object( self ) ) );
    171 changes: 171 additions & 0 deletions agents_widget.py
    View file
    Original file line number Diff line number Diff line change
    @@ -0,0 +1,171 @@
    #!/usr/bin/env python3
    # -*- coding:utf-8 -*-
    import PyQt5
    import asyncio
    import qtinter

    # Option Functionality
    from lib.ui.tabs import agent_proc_list_tab

    class AgentsWidget( PyQt5.QtWidgets.QWidget ):
    """
    A 'core' display within the ghost application view. Intended
    to display all the active agents and their status, as well
    as the information about them.
    Furthermore, its a core widget for performing actions on the
    widgets with the mouse.
    Provides functionality to interact with a pseudo-console,
    perform LDAP queries, and bulk interaction with multiple
    agents.
    """
    COLUMN_NAMES = [ "ID", "OS", "Process", "PID", "PPID", "Username" ]
    COLUMN_COUNT = len( COLUMN_NAMES );

    def __init__( self, ghost ):
    # init the widget from the ghost object
    super( PyQt5.QtWidgets.QWidget, self ).__init__( ghost );

    # set the ghost 'object' for performing most base operations
    self.ghost = ghost

    # create the primary layout for this widget
    self.layout = PyQt5.QtWidgets.QHBoxLayout();

    # List of agents in the database
    self.agents = []

    # create the table to display the agents
    self.agent_table = PyQt5.QtWidgets.QTableWidget();
    self.agent_table.setShowGrid( False );
    self.agent_table.setFocusPolicy( PyQt5.QtCore.Qt.NoFocus );
    self.agent_table.setSelectionBehavior( PyQt5.QtWidgets.QTableView.SelectRows );
    self.agent_table.setSelectionMode( PyQt5.QtWidgets.QTableView.SingleSelection );
    self.agent_table.setRowCount( 0 );
    self.agent_table.setColumnCount( self.COLUMN_COUNT );
    self.agent_table.setHorizontalHeaderLabels( self.COLUMN_NAMES );
    self.agent_table.verticalHeader().setVisible( False );
    self.agent_table.horizontalHeader().setHighlightSections( False );
    self.agent_table.setContextMenuPolicy( PyQt5.QtCore.Qt.CustomContextMenu );
    self.agent_table.customContextMenuRequested.connect( self._context_menu_requested );

    # Menu: "Options". A collection of options that can be used to task or request information
    # about the specific agent
    self.agent_option_menu = PyQt5.QtWidgets.QMenu()

    # Menu: "Options". Action: "Process List"
    self.agent_option_menu_action_process_list = PyQt5.QtWidgets.QAction( 'Process List' );
    self.agent_option_menu_action_process_list.triggered.connect( lambda: self._menu_agent_option_action_process_list( self.agent_table.selectionModel().selectedRows()[0].row() ) );
    self.agent_option_menu.addAction( self.agent_option_menu_action_process_list );

    # Loop through each column
    for i in range( 0, self.COLUMN_COUNT ):
    # Request that the column be stretched to fit the table view
    self.agent_table.horizontalHeader().setSectionResizeMode( i, PyQt5.QtWidgets.QHeaderView.Stretch );

    # add the table to the primary layout
    self.layout.addWidget( self.agent_table );

    # Lock for monitoring agents callback
    self.monitor_agent_lock = asyncio.Lock();

    # create the async queue for monitoring agents
    self.monitor_agent = PyQt5.QtCore.QTimer();
    self.monitor_agent.setInterval( 100 );
    self.monitor_agent.timeout.connect( self._monitor_agents );
    self.monitor_agent.start()

    # set the layout for this layout
    self.setLayout( self.layout );

    async def agent_table_get_id_by_row( self, agent_row ):
    """
    Returns the agent ID based on its row number.
    """
    # Returns an integer representing the ID of the agent
    return int( self.agent_table.item( agent_row, self.COLUMN_NAMES.index( "ID" ) ).text() )

    @qtinter.asyncslot
    async def _menu_agent_option_action_process_list( self, agent_row ):
    """
    Requests that the agent perform a process listing.
    """
    # Extract the agent ID from the row
    agent_id = await self.agent_table_get_id_by_row( agent_row );

    # Create the new tab
    await self.ghost.tab_widget.tab_add( agent_proc_list_tab.AgentProcListTab( self.ghost, agent_id ), f'[{agent_id}] Process Listing', True );

    @qtinter.asyncslot
    async def _context_menu_requested( self, pos ):
    """
    A custom menu called when the user right clicks on a row.
    """
    # Was one row selected? Then we can open a few options for that specific agent
    if len( self.agent_table.selectionModel().selectedRows() ) == 1:
    # Trigger the menu @ the position we were passed
    self.agent_option_menu.popup( PyQt5.QtGui.QCursor.pos() );

    async def _monitor_agents_add( self, agent_info : dict ):
    """
    Adds an agent to the table.
    """
    # Updates the row count with one more row
    self.agent_table.setRowCount( self.agent_table.rowCount() + 1 );

    # Column: "ID"
    item = PyQt5.QtWidgets.QTableWidgetItem( f'{agent_info[ "id" ]}' );
    item.setTextAlignment( PyQt5.QtCore.Qt.AlignCenter );
    item.setFlags( PyQt5.QtCore.Qt.ItemIsSelectable | PyQt5.QtCore.Qt.ItemIsEnabled );
    self.agent_table.setItem( self.agent_table.rowCount() - 1, 0, item );

    # Column: "OS"
    item = PyQt5.QtWidgets.QTableWidgetItem( f'{agent_info[ "os_major" ]}.{agent_info[ "os_minor" ]}.{agent_info[ "os_build" ]}' );
    item.setTextAlignment( PyQt5.QtCore.Qt.AlignCenter );
    item.setFlags( PyQt5.QtCore.Qt.ItemIsSelectable | PyQt5.QtCore.Qt.ItemIsEnabled );
    self.agent_table.setItem( self.agent_table.rowCount() - 1, 1, item );

    # Column: "Process"
    item = PyQt5.QtWidgets.QTableWidgetItem( f'{agent_info[ "process" ]}' );
    item.setTextAlignment( PyQt5.QtCore.Qt.AlignCenter );
    item.setFlags( PyQt5.QtCore.Qt.ItemIsSelectable | PyQt5.QtCore.Qt.ItemIsEnabled );
    self.agent_table.setItem( self.agent_table.rowCount() - 1, 2, item );

    # Column: "PID"
    item = PyQt5.QtWidgets.QTableWidgetItem( f'{agent_info[ "pid" ]}' );
    item.setTextAlignment( PyQt5.QtCore.Qt.AlignCenter );
    item.setFlags( PyQt5.QtCore.Qt.ItemIsSelectable | PyQt5.QtCore.Qt.ItemIsEnabled );
    self.agent_table.setItem( self.agent_table.rowCount() - 1, 3, item );

    # Column: "PPID"
    item = PyQt5.QtWidgets.QTableWidgetItem( f'{agent_info[ "ppid" ]}' );
    item.setTextAlignment( PyQt5.QtCore.Qt.AlignCenter );
    item.setFlags( PyQt5.QtCore.Qt.ItemIsSelectable | PyQt5.QtCore.Qt.ItemIsEnabled );
    self.agent_table.setItem( self.agent_table.rowCount() - 1, 4, item );

    # Column: "Username"
    item = PyQt5.QtWidgets.QTableWidgetItem( f'{agent_info[ "domain" ]}\\{agent_info[ "username" ]}' );
    item.setTextAlignment( PyQt5.QtCore.Qt.AlignCenter );
    item.setFlags( PyQt5.QtCore.Qt.ItemIsSelectable | PyQt5.QtCore.Qt.ItemIsEnabled );
    self.agent_table.setItem( self.agent_table.rowCount() - 1, 5, item );

    @qtinter.asyncslot
    async def _monitor_agents( self ):
    """
    Monitors for new agents and updates the table with their information.
    """
    # Is there a callback not running at the moment?
    if not self.monitor_agent_lock.locked():
    # Lock access to the agent resource
    async with self.monitor_agent_lock:
    # Query the list of the agents!
    agent_list = await self.ghost.rpc.teamserver_agent_list_get();

    # Loop through every that hasnt been added yet!
    for agent in [ agent for agent in agent_list if agent[ 'id' ] not in self.agents ]:
    # Add the entry to the list!
    self.agents.append( agent[ 'id' ] );

    # Add the entry to the table!
    await self._monitor_agents_add( agent );