ui/models/hostmodels.py
#!/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/>.
"""
import re
from PyQt6 import QtWidgets, QtGui, QtCore
from PyQt6.QtGui import QFont
from PyQt6.QtCore import pyqtSignal, QObject
from app.ModelHelpers import resolveHeaders, itemSelectable
from app.auxiliary import * # for bubble sort
class HostsTableModel(QtCore.QAbstractTableModel):
def __init__(self, hosts = [[]], headers = [], parent = None):
QtCore.QAbstractTableModel.__init__(self, parent)
self.__headers = headers
self.__hosts = hosts
def setHosts(self, hosts):
self.__hosts = hosts
def rowCount(self, parent):
return len(self.__hosts)
def columnCount(self, parent):
if len(self.__hosts) != 0:
return len(self.__hosts[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.DecorationRole: # to show the operating system icon instead of text
if index.column() == 1: # if trying to display the operating system
os_string = self.__hosts[index.row()]['osMatch']
if os_string == '': # if there is no OS information, use the question mark icon
return QtGui.QIcon("./images/question-icon.png")
elif re.search('[lL]inux', os_string, re.I):
return QtGui.QIcon("./images/linux-icon.png")
elif re.search('[wW]indows', os_string, re.I):
return QtGui.QIcon("./images/windows-icon.png")
elif re.search('[cC]isco', os_string, re.I):
return QtGui.QIcon("./images/cisco-big.jpg")
elif re.search('HP ', os_string, re.I):
return QtGui.QIcon("./images/hp-icon.png")
elif re.search('[vV]x[wW]orks', os_string, re.I):
return QtGui.QIcon("./images/hp-icon.png")
elif re.search('[vV]m[wW]are', os_string, re.I):
return QtGui.QIcon("./images/vmware-big.jpg")
else: # if it's an unknown OS also use the question mark icon
return QtGui.QIcon("./images/question-icon.png")
if role == QtCore.Qt.ItemDataRole.DisplayRole: # how to display each cell
value = ''
row = index.row()
column = index.column()
if column == 0:
value = self.__hosts[row]['id']
elif column == 2:
value = self.__hosts[row]['osAccuracy']
elif column == 3:
if not self.__hosts[row]['hostname'] == '':
value = self.__hosts[row]['ip'] + ' ('+ self.__hosts[row]['hostname'] +')'
else:
value = self.__hosts[row]['ip']
elif column == 4:
value = self.__hosts[row]['ipv4']
elif column == 5:
value = self.__hosts[row]['ipv6']
elif column == 6:
value = self.__hosts[row]['macaddr']
elif column == 7:
value = self.__hosts[row]['status']
elif column == 8:
value = self.__hosts[row]['hostname']
elif column == 9:
value = self.__hosts[row]['vendor']
elif column == 10:
value = self.__hosts[row]['uptime']
elif column == 11:
value = self.__hosts[row]['lastboot']
elif column == 12:
value = self.__hosts[row]['distance']
elif column == 13:
value = self.__hosts[row]['checked']
elif column == 14:
value = self.__hosts[row]['state']
elif column == 15:
value = self.__hosts[row]['count']
else:
value = 'Not set in view model'
return value
if role == QtCore.Qt.ItemDataRole.FontRole:
# if a host is checked strike it out and make it italic
if index.column() == 3 and self.__hosts[index.row()]['checked'] == 'True':
checkedFont=QFont()
checkedFont.setStrikeOut(True)
checkedFont.setItalic(True)
return checkedFont
# 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 itemSelectable()
# sort function called when the user clicks on a header
def sort(self, Ncol, order):
self.layoutAboutToBeChanged.emit()
array = []
if Ncol == 0 or Ncol == 3: # if sorting by IP address (and by default)
log.debug("__hosts: {0}".format(str(self.__hosts)))
for i in range(len(self.__hosts)):
array.append(IP2Int(self.__hosts[i]['ip']))
elif Ncol == 1: # if sorting by OS
for i in range(len(self.__hosts)):
os_string = self.__hosts[i]['osMatch']
if os_string == '':
array.append('')
elif re.search('[lL]inux', os_string, re.I):
array.append('Linux')
elif re.search('[wW]indows', os_string, re.I):
array.append('Windows')
elif re.search('[cC]isco', os_string, re.I):
array.append('Cisco')
elif re.search('HP ', os_string, re.I):
array.append('Hp')
elif re.search('[vV]x[wW]orks', os_string, re.I):
array.append('Hp')
elif re.search('[vV]m[wW]are', os_string, re.I):
array.append('Vmware')
else:
array.append('')
sortArrayWithArray(array, self.__hosts) # sort the array of OS
if order == Qt.SortOrder.AscendingOrder: # reverse if needed
self.__hosts.reverse()
self.layoutChanged.emit() # update the UI (built-in signal)
### getter functions ###
def getHostIPForRow(self, row):
return self.__hosts[row]['ip']
def getHostIdForRow(self, row):
return self.__hosts[row]['id']
def getHostCheckStatusForRow(self, row):
return self.__hosts[row]['checked']
def getHostCheckStatusForIp(self, ip):
for i in range(len(self.__hosts)):
if str(self.__hosts[i]['ip']) == str(ip):
return self.__hosts[i]['checked']
def getRowForIp(self, ip):
for i in range(len(self.__hosts)):
if self.__hosts[i]['ip'] == ip:
return i