#!/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 );