main.py

#Import the necessary libraries
import sys, os
import pandas as pd
from PyQt6.QtWidgets import *
from PyQt6 import uic
from PyQt6.QtCore import QTimer
from PyQt6.QtGui import QIcon, QColor

class Interface(QMainWindow):
    def __init__(self):
        super().__init__()

        # Interface connection 
        uic.loadUi(self.resolve_route('interface.ui'), self)

        #Icon
        icon_path = self.resolve_route('icon.ico')
        app_icon = QIcon(icon_path)
        self.setWindowIcon(app_icon)

        # Setting title
        self.setWindowTitle("Vulnerability Status Checker")

        # Table UI
        self.table = self.findChild(QTableWidget, 'tableWidget')
        self.table.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)

        #Lcds UI
        self.lcd_solved = self.findChild(QLCDNumber, 'lcdNumber_2')
        self.lcd_total = self.findChild(QLCDNumber, 'lcdNumber_3')
        self.lcd_vulnerabilities = self.findChild(QLCDNumber, 'lcdNumber')

        #Search Engine UI
        self.search_engine = self.findChild(QLineEdit, 'lineEdit')
        self.search_engine.textChanged.connect(self.query_table)

        #Buttons UI
        self.btn_file1 = self.findChild(QPushButton, 'pushButton') 
        self.btn_file2 = self.findChild(QPushButton, 'pushButton_2') 
        self.btn_run = self.findChild(QPushButton, 'pushButton_3') 
        self.btn_show_daily = self.findChild(QPushButton, 'pushButton_4')
        self.btn_show_database = self.findChild(QPushButton, 'pushButton_5')
        self.btn_show_new_assets_cves = self.findChild(QPushButton, 'pushButton_6')
        self.btn_export_new_cves = self.findChild(QPushButton, 'pushButton_7')
        self.bnt_save_changes = self.findChild(QPushButton, 'pushButton_8')

        #LineEdit
        self.line_ed_search = self.findChild(QLineEdit, 'lineEdit') 
        self.line_ed_status = self.findChild(QLineEdit, 'lineEdit_2') 

        #CheckBox
        self.cbox_filter = self.findChild(QCheckBox, 'checkBox')
        self.cbox_filter.setEnabled(False)
        
        #Connections UI 
        self.btn_file1.clicked.connect(lambda: self.import_file('Open File 1 - Daily')) 
        self.btn_file2.clicked.connect(lambda: self.import_file('Open File 2 - Database'))
        self.btn_run.clicked.connect(lambda: self.run_check())
        self.btn_show_daily.clicked.connect(lambda: self.show_table(self.btn_show_daily.text()))
        self.btn_show_database.clicked.connect(lambda: self.show_table(self.btn_show_database.text()))
        self.btn_show_new_assets_cves.clicked.connect(lambda: self.show_table(self.btn_show_new_assets_cves.text()))
        self.btn_export_new_cves.clicked.connect(lambda: self.export_new_cves())
        self.bnt_save_changes.clicked.connect(lambda: self.save_database_changes())
        self.cbox_filter.stateChanged.connect(self.toggle_columns)
        

        # Variables
        self.file_type = ''
        self.status = ['File exported successfully', 'File uploaded successfully']
        self.columns_to_filter = None
        self.green_color = QColor(204, 255, 204)
        self.red_color = QColor(255, 204, 204)
        self.aux_filter = ''

    # Resolve route
    def resolve_route(self, ruta_relativa):
            if hasattr(sys, '_MEIPASS'):
                return os.path.join(sys._MEIPASS, ruta_relativa)
            return os.path.join(os.path.abspath('.'), ruta_relativa)
    
    # Ajuste de widgets
    def resizeEvent(self, event):
        table_width = self.width() - 20
        table_height = self.height() -188
        self.table.setGeometry(10, 178, table_width, table_height)

    def show_table(self, option):
        self.cbox_filter.setChecked(False)
        self.cbox_filter.setEnabled(False)
        self.aux_filter = ''
        self.line_ed_search.setText("")

        #Show daily file
        if option == 'Show Daily':
            # Reading file 
            try:
                self.aux_filter = 'filter_daily'
                self.line_ed_status.setText("")
                df_daily = self.df_daily
                self.rows = df_daily.shape[0]
                self.columns = df_daily.shape[1]
                
                # UI table settings 
                self.table.setRowCount(self.rows)
                self.table.setColumnCount(self.columns)
                self.name_columns = list(df_daily.columns)
                self.table.setHorizontalHeaderLabels(self.name_columns)
                self.fill_table(df_daily, self.rows, self.columns)
                self.line_ed_status.setStyleSheet(f'background-color: {self.green_color.name()}')
                self.line_ed_status.setText(f"Displaying daily archive information")
                self.cbox_filter.setEnabled(True)
            except Exception as e:
                self.table.setRowCount(0)
                self.table.setColumnCount(0)
                self.line_ed_status.setStyleSheet(f'background-color: {self.red_color.name()}')
                self.line_ed_status.setText(f'Error: {type(e).__name__}. Description: {str(e)}')

        #Show database file
        elif option == 'Show Database':
            self.line_ed_status.setText("")
            # Reading file 
            try:
                self.aux_filter = 'filter_database'
                df_database = self.df_database
                self.rows = df_database.shape[0]
                self.columns = df_database.shape[1]
                
                # UI table settings 
                self.table.setRowCount(self.rows)
                self.table.setColumnCount(self.columns)
                self.name_columns = list(df_database.columns)
                self.table.setHorizontalHeaderLabels(self.name_columns)
                self.fill_table(df_database, self.rows, self.columns)
                self.line_ed_status.setStyleSheet(f'background-color: {self.green_color.name()}')
                self.line_ed_status.setText(f"Displaying database archive information")
                self.cbox_filter.setEnabled(True)

            except Exception as e:
                self.table.setRowCount(0)
                self.table.setColumnCount(0)
                self.line_ed_status.setStyleSheet(f'background-color: {self.red_color.name()}')
                self.line_ed_status.setText(f'Error: {type(e).__name__}. Description: {str(e)}')

        #Show new assets and cves file
        elif option == 'Show Assets':
            self.line_ed_status.setText("")
            # Reading file  
            try:
                self.aux_filter = 'filter_new_cves'
                #df_new_cves = pd.read_csv(self.file_name, sep=';', encoding='latin-1')
                df_new_cves = pd.read_excel(self.file_name)
                self.rows = df_new_cves.shape[0]
                self.columns = df_new_cves.shape[1]
                
                # UI table settings 
                self.table.setRowCount(self.rows)
                self.table.setColumnCount(self.columns)
                self.name_columns = list(df_new_cves.columns)
                self.table.setHorizontalHeaderLabels(self.name_columns)
                self.fill_table(df_new_cves, self.rows, self.columns)
                self.line_ed_status.setStyleSheet(f'background-color: {self.green_color.name()}')
                self.line_ed_status.setText(f"Displaying assets and cves archive information")
                self.cbox_filter.setEnabled(True)

            except Exception as e:
                self.table.setRowCount(0)
                self.table.setColumnCount(0) 
                self.line_ed_status.setStyleSheet(f'background-color: {self.red_color.name()}')               
                self.line_ed_status.setText(f'Error: {type(e).__name__}. Description: {str(e)}')

        # Show check run
        elif option == 'Run Check':
            # Reading file 
            try:
                self.aux_filter = 'filter_run'
                df_database = self.df_database
                self.rows = df_database.shape[0]
                self.columns = df_database.shape[1]
                
                # UI table settings 
                self.table.setRowCount(self.rows)
                self.table.setColumnCount(self.columns)
                self.name_columns = list(df_database.columns)
                self.table.setHorizontalHeaderLabels(self.name_columns)
                self.fill_table(df_database, self.rows, self.columns)
                self.cbox_filter.setEnabled(True)


            except Exception as e:
                self.table.setRowCount(0)
                self.table.setColumnCount(0)  
                self.line_ed_status.setStyleSheet(f'background-color: {self.red_color.name()}')              
                self.line_ed_status.setText(f'Error: {type(e).__name__}. Desc: {str(e)}')
    
    # Llenado de la tabla según sea el caso 
    def fill_table(self, df, rows, columns):
         for row in range(rows):
                    for column in range(columns):
                        data = str(df.iloc[row,column])
                        self.table.setItem(row, column, QTableWidgetItem(data))


    # Cargado del archivo Daily y el archivo Database
    def import_file(self, file_type):
        self.cbox_filter.setChecked(False)
        self.cbox_filter.setEnabled(False)
        self.line_ed_search.setText("")
        file_path, _ = QFileDialog.getOpenFileName(self, f"{file_type}", "", "Archivos de Excel (*.xlsx *.xls)")
        if file_path:
            try:
                if file_type == "Open File 1 - Daily":
                    #self.df_daily = pd.read_csv(file_path, sep=';', encoding='latin-1')
                    self.df_daily = pd.read_excel(file_path,  dtype={'Device name': str, 'Vulnerability ID': str})
                    self.line_ed_status.setStyleSheet(f'background-color: {self.green_color.name()}')
                    self.line_ed_status.setText(f"Daily file successfully uploaded")
                    self.lcd_vulnerabilities.display('0')
                    self.lcd_solved.display('0')
                    self.lcd_total.display('0')
                    self.table.setRowCount(0)
                    self.table.setColumnCount(0)
                elif file_type == "Open File 2 - Database":
                    self.df_database = pd.read_excel(file_path, dtype={'Device name': str, 'Vulnerability ID': str})
                    self.line_ed_status.setStyleSheet(f'background-color: {self.green_color.name()}')
                    self.line_ed_status.setText(f"Database file successfully uploaded")
                    self.lcd_vulnerabilities.display('0')
                    self.lcd_solved.display('0')
                    self.lcd_total.display('0')
                    self.table.setRowCount(0)
                    self.table.setColumnCount(0)
            except Exception as e:
                self.table.setRowCount(0)
                self.table.setColumnCount(0)
                self.line_ed_status.setStyleSheet(f'background-color: {self.red_color.name()}')
                self.line_ed_status.setText(f'Error: {type(e).__name__}. Description: {str(e)}')
    
    # Ejecución del chequeo del estado de las vulnerabildades
    def run_check(self):
        try:
            self.temp_df = self.df_daily[['Device name', 'Vulnerability ID']].apply(lambda x: x.str.lower().str.strip()).merge(
                self.df_database[['Device name', 'Vulnerability ID']].apply(lambda x: x.str.lower().str.strip()),
                on=['Device name', 'Vulnerability ID'],
                how='outer',
                indicator='indicador'
            )
            
            self.temp_df_unique = self.df_daily[['Vulnerability ID']].apply(lambda x: x.str.lower().str.strip()).merge(
                self.df_database[['Vulnerability ID']].apply(lambda x: x.str.lower().str.strip()),
                on=['Vulnerability ID'],
                how='outer',
                indicator='indicador'
            )

            self.temp_df_assets = self.df_daily[['Device name']].apply(lambda x: x.str.lower().str.strip()).merge(
                self.df_database[['Device name']].apply(lambda x: x.str.lower().str.strip()),
                on=['Device name'],
                how='outer',
                indicator='indicador'
            )
            
            self.unique_cves = self.temp_df_unique[self.temp_df_unique['indicador']=='left_only']
            self.unique_cves['Vulnerability ID'] = self.unique_cves['Vulnerability ID'].astype(str)

            print(f"UNIQUES CVES: {self.unique_cves['Vulnerability ID'].info}")
            self.unique_assets = self.temp_df_assets[self.temp_df_assets['indicador']=='left_only']
            self.unique_assets['Device name'] = self.unique_assets['Device name'].astype(str)
            print(f"UNIQues asstes: {self.unique_assets['Device name'].info}")


            # Clasificación de las vulnerabilidades
            self.only_daily = self.temp_df[self.temp_df['indicador']=='left_only'][['Device name', 'Vulnerability ID']]
            self.only_database = self.temp_df[self.temp_df['indicador']=='right_only'][['Device name', 'Vulnerability ID']]
            self.only_same = self.temp_df[self.temp_df['indicador']=='both'][['Device name', 'Vulnerability ID']]

            # Crear una nueva columna que combine activo y cve_id
            self.temp_database = self.df_database.copy()
            self.temp_database = self.temp_database[['Device name', 'Vulnerability ID']].apply(lambda x: x.str.lower().str.strip())
            self.temp_database['key'] = self.temp_database['Device name'] + '_' + self.temp_database['Vulnerability ID']
            self.only_daily['key'] = self.only_daily['Device name'] + '_' + self.only_daily['Vulnerability ID']
            self.only_database['key'] = self.only_database['Device name'] + '_' + self.only_database['Vulnerability ID']
            self.only_same['key'] = self.only_same['Device name'] + '_' + self.only_same['Vulnerability ID']

            # Actualización del estado Solucionado 
            format = "%d/%m/%Y" #pd.Timestamp("today").strftime()
            today = pd.to_datetime('today').date().strftime(format)
            self.df_database.loc[(self.temp_database['key'].isin(self.only_database['key'])), 'actualizacion_interlan'] = 'Solucionado'
            self.df_database.loc[(self.temp_database['key'].isin(self.only_database['key'])), 'Fecha actualizacion'] = today

            # Actualización del estado Vulnerable
            self.df_database.loc[(self.temp_database['key'].isin(self.only_same['key'])), 'actualizacion_interlan'] = 'Vulnerable'
            self.df_database.loc[(self.temp_database['key'].isin(self.only_same['key'])), 'Fecha actualizacion'] = today
            self.line_ed_status.setStyleSheet(f'background-color: {self.green_color.name()}')
            self.line_ed_status.setText("Verification has been successfully completed")

            self.set_labels(self.df_database)
            
            return self.show_table(self.btn_run.text())
    
        except Exception as e:
            self.table.setRowCount(0)
            self.table.setColumnCount(0)
            self.line_ed_status.setStyleSheet(f'background-color: {self.red_color.name()}')
            self.line_ed_status.setText(f'Error: {type(e).__name__}. Description: {str(e)}')
    
    # Exportación del archivo que contiene los activos y sus respectivos cves que no estan relacionados en la base de datos principal
    def export_new_cves(self):
        try:
            self.cbox_filter.setChecked(False)
            self.cbox_filter.setEnabled(False)
            self.line_ed_search.setText("")
            self.file_name, _ = QFileDialog.getSaveFileName(self,"Save new assets","","Archivos de Excel (*.xlsx *.xls)")
            self.file_name_unique, _ = QFileDialog.getSaveFileName(self,"Save new cves","","Archivos de Excel (*.xlsx *.xls)")
            if self.file_name:
                
                 # Archivo unicos de assets
                self.temp_df_daily = self.df_daily.copy()
                self.temp_df_daily[['Device name', 'Vulnerability ID']] = self.temp_df_daily[['Device name', 'Vulnerability ID']].astype(str)
                self.temp_df_daily[['Device name', 'Vulnerability ID']] = self.temp_df_daily[['Device name', 'Vulnerability ID']].apply(lambda x: x.str.lower().str.strip())
                self.unique_assets['Device name'] = self.unique_assets['Device name'].apply(lambda x: str(x).lower().strip())
                self.temp_df_daily = self.temp_df_daily[self.temp_df_daily['Device name'].isin(self.unique_assets['Device name'])]
                #Generación de archivo de activos
                self.temp_df_daily.to_excel(self.file_name, index = False)

                # Archivo unico de cves
                self.temp_df_daily_unique = self.df_daily.copy()
                self.temp_df_daily_unique[['Device name', 'Vulnerability ID']] = self.temp_df_daily_unique[['Device name', 'Vulnerability ID']].astype(str)
                self.temp_df_daily_unique[['Device name', 'Vulnerability ID']] = self.temp_df_daily_unique[['Device name', 'Vulnerability ID']].apply(lambda x: x.str.lower().str.strip())
                self.unique_cves['Device name'] = self.unique_cves['Vulnerability ID'].apply(lambda x: str(x).lower().strip())
                self.temp_df_daily_unique = self.temp_df_daily_unique[self.temp_df_daily_unique['Vulnerability ID'].isin(self.unique_cves['Vulnerability ID'])]
                #Generación de archivo de CVEs nuevos
                self.temp_df_daily_unique.to_excel(self.file_name_unique, index= False )

                # Vaciado de tabla
                self.table.setRowCount(0)
                self.table.setColumnCount(0)
                self.line_ed_status.setStyleSheet(f'background-color: {self.green_color.name()}')
                self.line_ed_status.setText(f'File exported successfully: {self.file_name}')

        except Exception as e:
            self.table.setRowCount(0)
            self.table.setColumnCount(0)
            self.line_ed_status.setStyleSheet(f'background-color: {self.red_color.name()}')
            self.line_ed_status.setText(f'Error: {type(e).__name__}. Descripción: {str(e)}')

    # Guardado de la base de datos modificada
    def save_database_changes(self):
        try:
            self.cbox_filter.setChecked(False)
            self.cbox_filter.setEnabled(False)
            self.line_ed_search.setText("")
            self.file_changes, _ = QFileDialog.getSaveFileName(self,"Save changes as","","Archivos de Excel (*.xlsx *.xls)")
            if (self.file_changes and len(self.temp_df)>0):
                #Se genera el archivo con las modificaciones de la verificación en la base de datos original
                self.df_database.to_excel(self.file_changes, index= False)
                self.table.setRowCount(0)
                self.table.setColumnCount(0)
                self.line_ed_status.setStyleSheet(f'background-color: {self.green_color.name()}')
                self.line_ed_status.setText(f'File saved successfully: {self.file_changes}')
        
        except Exception as e:
            self.table.setRowCount(0)
            self.table.setColumnCount(0)
            self.line_ed_status.setStyleSheet(f'background-color: {self.red_color.name()}')
            self.line_ed_status.setText(f'Error: {type(e).__name__}. Descripción: {str(e)}')       

    # Motor de consultas en la tabla
    def query_table(self, query):
        # Obtener el modelo de datos de la tabla
        model = self.table.model()
        # Iterar sobre las filas del modelo
        for row in range(model.rowCount()):
            # Concatenar el contenido de la fila en una cadena
            row_data = []
            for column in range(model.columnCount()):
                index = model.index(row, column)
                data = str(model.data(index))
                row_data.append(data)
            row_content = ''.join(row_data)
            # Ocultar o mostrar la fila según la consulta de búsqueda
            match = query.lower() in row_content.lower()
            self.table.setRowHidden(row, not match)
        # Actualizar la vista de la tabla
        self.table.update()

    def toggle_columns(self, checked):
        try: 
            if (self.aux_filter == 'filter_daily' or self.aux_filter == 'filter_new_cves'):
                self.columns_to_filter = [0, 1, 3, 4, 6, 7, 8, 12]
                for i in range(self.table.columnCount()):
                    self.table.setColumnHidden(i, i not in self.columns_to_filter and checked)
                self.table.resizeColumnsToContents()
            elif (self.aux_filter == 'filter_database' or self.aux_filter == 'filter_run'):
                self.columns_to_filter = [0, 10, 28, 29]
                for i in range(self.table.columnCount()):
                    self.table.setColumnHidden(i, i not in self.columns_to_filter and checked)
                self.table.resizeColumnsToContents()
        except Exception as e:
            self.line_ed_status.setStyleSheet(f'background-color: {self.red_color.name()}')
            self.line_ed_status.setText(f'Error: {type(e).__name__}. Description: {str(e)}')

    #  Visualización de las variables en los lcds
    def set_labels(self, database):
        num_vulnerabilities = self.get_number_vulnerabilities(database)
        num_solved = self.get_number_solved(database)
        self.lcd_vulnerabilities.display(num_vulnerabilities)
        self.lcd_solved.display(num_solved)
        self.lcd_total.display(num_solved + num_vulnerabilities)

    # Cálculo de las vulnerabilidades a mostrar en lcd
    def get_number_vulnerabilities(self, database):
        num_vulnerabilities = database[database['actualizacion_interlan']=='Vulnerable'].shape[0] 
        return num_vulnerabilities

    # Cálculo de los casos solucionados a mostrar en lcd
    def get_number_solved(self, database):
        num_solved = database[database['actualizacion_interlan']=='Solucionado'].shape[0]
        return num_solved

if __name__ == '__main__':
    app = QApplication(sys.argv)
    frame = Interface()
    frame.show()
    sys.exit(app.exec())