GoVanguard/legion

View on GitHub
ui/models/servicemodels.py

Summary

Maintainability
F
6 days
Test Coverage
#!/usr/bin/env python

"""
LEGION (https://gotham-security.com)
Copyright (c) 2023 Gotham Security

    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 3 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, see <http://www.gnu.org/licenses/>.

"""

from PyQt6 import QtWidgets, QtGui, QtCore
from PyQt6.QtCore import pyqtSignal, QObject

from app.ModelHelpers import resolveHeaders, itemInteractive
from app.auxiliary import *                                                 # for bubble sort

class ServicesTableModel(QtCore.QAbstractTableModel):

    def __init__(self, services = [[]], headers = [], parent = None):
        QtCore.QAbstractTableModel.__init__(self, parent)
        self.__headers = headers
        self.__services = services
        
    def setServices(self, services):
        self.__services = services

    def rowCount(self, parent):
        return len(self.__services)

    def columnCount(self, parent):
        if len(self.__services) != 0:
            return len(self.__services[0])
        return 0
        
    def headerData(self, section, orientation, role):
        return resolveHeaders(role, orientation, section, self.__headers)

    # this method takes care of how the information is displayed
    def data(self, index, role):
        if role == QtCore.Qt.ItemDataRole.DecorationRole:                            # to show the open/closed/filtered icons
            if index.column() == 0 or index.column() == 2:
                tmp_state = self.__services[index.row()]['state']

                stateMap = {'open':'open', 'closed':'closed', 'filtered':'filtered'}
                defaultState = 'filtered'

                stateIconName = stateMap.get(str(tmp_state)) or defaultState
                stateIcon = "./images/{stateIconName}.gif".format(stateIconName=stateIconName)
                return QtGui.QIcon(stateIcon)

        if role == QtCore.Qt.ItemDataRole.DisplayRole:                               # how to display each cell
            value = ''
            row = index.row()
            column = index.column()

            if column == 0:
                # the spaces are needed for spacing with the icon that precedes the text
                value = '   ' + self.__services[row]['ip']
            elif column == 1:
                value = self.__services[row]['portId']
            elif column == 2:
                # the spaces are needed for spacing with the icon that precedes the text
                value = '   ' + self.__services[row]['portId']
            elif column == 3:
                value = self.__services[row]['protocol']
            elif column == 4:
                value = self.__services[row]['state']
            elif column == 5:
                value = self.__services[row]['hostId']
            elif column == 6:
                value = self.__services[row]['serviceId']
            elif column == 7:
                value = self.__services[row]['name']
            elif column == 8:
                value = self.__services[row]['product']
            elif column == 9:
                if not self.__services[row]['product'] == None and not self.__services[row]['product'] == '':
                    value = str(self.__services[row]['product'])
                
                if not self.__services[row]['version'] == None and not self.__services[row]['version'] == '':
                    value = value + ' ' + self.__services[row]['version']

                if not self.__services[row]['extrainfo'] == None and not self.__services[row]['extrainfo'] == '':
                    value = value + ' (' + self.__services[row]['extrainfo'] + ')'
            elif column == 10:
                value = self.__services[row]['extrainfo']
            elif column == 11:
                value = self.__services[row]['fingerprint']
            return value

    # method that allows views to know how to treat each item, eg: if it should be enabled, editable, selectable etc
    def flags(self, index):
        return itemInteractive()

    # sort function called when the user clicks on a header
    def sort(self, Ncol, order):
        self.layoutAboutToBeChanged.emit()
        array = []
        
        if Ncol == 0:                                                   # if sorting by ip (and by default)
            for i in range(len(self.__services)):
                array.append(IP2Int(self.__services[i]['ip']))

        elif Ncol == 1:                                                 # if sorting by port
            for i in range(len(self.__services)):
                array.append(int(self.__services[i]['portId']))

        elif Ncol == 2:                                                 # if sorting by port
            for i in range(len(self.__services)):
                array.append(int(self.__services[i]['portId']))
                
        elif Ncol == 3:                                                 # if sorting by protocol
            for i in range(len(self.__services)):
                array.append(self.__services[i]['protocol'])
                
        elif Ncol == 4:                                                 # if sorting by state
            for i in range(len(self.__services)):
                array.append(self.__services[i]['state'])
                
        elif Ncol == 7:                                                 # if sorting by name
            for i in range(len(self.__services)):
                array.append(self.__services[i]['name'])
            
        elif Ncol == 9:                                                 # if sorting by version
            for i in range(len(self.__services)):
                value = ''
                if not self.__services[i]['product'] == None and not self.__services[i]['product'] == '':
                    value = str(self.__services[i]['product'])
                
                if not self.__services[i]['version'] == None and not self.__services[i]['version'] == '':
                    value = value + ' ' + self.__services[i]['version']

                if not self.__services[i]['extrainfo'] == None and not self.__services[i]['extrainfo'] == '':
                    value = value + ' (' + self.__services[i]['extrainfo'] + ')'
                array.append(value)

        # sort the services based on the values in the array
        sortArrayWithArray(array, self.__services)
        
        if order == Qt.SortOrder.AscendingOrder:                                  # reverse if needed
            self.__services.reverse()
            
        self.layoutChanged.emit()                           # update the UI (built-in signal)

    ### getter functions ###
    
    def getPortForRow(self, row):
        return self.__services[row]['portId']
        
    def getServiceNameForRow(self, row):
        return self.__services[row]['name']
            
    def getIpForRow(self, row):
        return self.__services[row]['ip']
        
    def getProtocolForRow(self, row):
        return self.__services[row]['protocol']

    ####################################################################

class ServiceNamesTableModel(QtCore.QAbstractTableModel):

    def __init__(self, serviceNames = [[]], headers = [], parent = None):
        QtCore.QAbstractTableModel.__init__(self, parent)
        self.__headers = headers
        self.__serviceNames = serviceNames
        
    def setServices(self, serviceNames):
        self.__serviceNames = serviceNames

    def rowCount(self, parent):
        return len(self.__serviceNames)

    def columnCount(self, parent):
        if len(self.__serviceNames) != 0:
            return len(self.__serviceNames[0])
        return 0
        
    def headerData(self, section, orientation, role):
        return resolveHeaders(role, orientation, section, self.__headers)

    def data(self, index, role):   # This method takes care of how the information is displayed

        if role == QtCore.Qt.ItemDataRole.DisplayRole:                               # how to display each cell
            row = index.row()
            column = index.column()
            if column == 0:
                return self.__serviceNames[row]['name']

    # method that allows views to know how to treat each item, eg: if it should be enabled, editable, selectable etc
    def flags(self, index):
        return QtCore.Qt.ItemFlag.ItemIsEnabled | QtCore.Qt.ItemFlag.ItemIsSelectable | QtCore.Qt.ItemFlag.ItemIsEditable

    # sort function called when the user clicks on a header
    def sort(self, Ncol, order):
        
        self.layoutAboutToBeChanged.emit()
        array = []
        
        if Ncol == 0:  # if sorting by service name (and by default)
            for i in range(len(self.__serviceNames)):
                array.append(self.__serviceNames[i]['name'])

        # sort the services based on the values in the array
        sortArrayWithArray(array, self.__serviceNames)

        if order == Qt.SortOrder.AscendingOrder:                                  # reverse if needed
            self.__serviceNames.reverse()
            
        self.layoutChanged.emit()                            # update the UI (built-in signal)

    ### getter functions ###

    def getServiceNameForRow(self, row):
        return self.__serviceNames[row]['name']

    def getRowForServiceName(self, serviceNames):
        for i in range(len(self.__serviceNames)):
            if self.__serviceNames[i]['name'] == serviceNames:
                return i