Fixed Vuln auto complete removes value

PentestMapper.py

@@ -55,7 +55,7 @@ def registerExtenderCallbacks(self, callbacks):
 
         # Creating a output after loading
         callbacks.printOutput("Author: Sourav Kalal Aka AnoF")
-        callbacks.printOutput("Version: 1.7")
+        callbacks.printOutput("Version: 1.7.1")
         callbacks.printOutput("https://github.com/Anof-cyber/Pentest-Mapper")
 
         callbacks.registerContextMenuFactory(self)
@@ -65,6 +65,7 @@ def registerExtenderCallbacks(self, callbacks):
         self.tab.add("Center", self.tabbedPane)
 
         self._log = list()
+        self._vuln = list()
         self._lock = Lock()
 
         # Creating Another Tab in the extension tab
@@ -291,8 +292,12 @@ def registerExtenderCallbacks(self, callbacks):
 
 
         # creating the table to vulnerability tab
-        self.colNames3 = ('URL', 'Parameters','Vulnerability','Severity','Request','Response')
-        self.dataModel2 = CustomDefaultTableModelHosts2(None,self.colNames3)
+        #self.colNames3 = ['URL', 'Parameters','Vulnerability','Severity','Request','Response']
+        #self.dataModel2 = CustomDefaultTableModelHosts2(None,self.colNames3)
+        self.dataModel2 = CustomTableModelVuln(self)
+
+
+
         self.table3 = JTable(self.dataModel2)
         self.table3.setAutoCreateRowSorter(True)
         self.table3.getColumnModel().getColumn(2).setPreferredWidth(0)
@@ -305,7 +310,7 @@ def registerExtenderCallbacks(self, callbacks):
         #comboColumn.setCellEditor(DefaultCellEditor(self.comboBox))
 
 
-        editor = AutocompleteTableCellEditor(self.combolist)
+        editor = AutocompleteTableCellEditor(self.combolist,self.table3)
         comboColumn = self.table3.getColumnModel().getColumn(2)
         comboColumn.setCellEditor(editor)
 
@@ -363,6 +368,7 @@ def registerExtenderCallbacks(self, callbacks):
         listener = CustomSelectionListener(self.table3, self.requestViewer1, self.responseViewer1)
         self.table3.getSelectionModel().addListSelectionListener(listener)
 
+
 
 
 
@@ -1326,15 +1332,6 @@ def isCellEditable(self, row, column):
         return 0
 
 
-# Extending the default table model for vulnerability tab column 0 (url) is not editable 
-class CustomDefaultTableModelHosts2(DefaultTableModel):
-
-    # override isCellEditable method
-    def isCellEditable(self, row, column):
-        if column == 0:
-            return 0
-        else:
-            return 1
 
 
 class CustomSelectionListener(ListSelectionListener):
@@ -1358,16 +1355,22 @@ def valueChanged(self, event):
 
 # Vulnerability Table vulnerability selection autocompletion
 class AutocompleteTableCellEditor(TableCellEditor):
-    def __init__(self, items):
+    def __init__(self, items,table3):
+        self.items = items
         self.comboBox = JComboBox(items)
         self.comboBox.setEditable(True)
         AutoCompleteDecorator.decorate(self.comboBox)
         self.listeners = []
+        self.table3 = table3
+
 
     def getCellEditorValue(self):
         return self.comboBox.getSelectedItem()
 
     def getTableCellEditorComponent(self, table, value, isSelected, row, column):
+        self.comboBox = JComboBox(self.items)
+        self.comboBox.setEditable(True)
+        AutoCompleteDecorator.decorate(self.comboBox)
         self.comboBox.setSelectedItem(value)
         return self.comboBox
 
@@ -1384,7 +1387,124 @@ def shouldSelectCell(self, event):
         return True
 
     def stopCellEditing(self):
+        selected_value = self.comboBox.getSelectedItem()
+        selected_row = self.table3.convertRowIndexToModel(self.table3.getEditingRow())
+        selected_column = self.table3.convertColumnIndexToModel(self.table3.getEditingColumn())
+        self.table3.getModel().setValueAt(selected_value, selected_row, selected_column)
+
         return True
 
     def cancelCellEditing(self):
-        return True
+        return True
+
+
+
+class RowData:
+    def __init__(self, url="", parameters="", vulnerability="", severity="", request="", response=""):
+        self.url = url
+        self.parameters = parameters
+        self.vulnerability = vulnerability
+        self.severity = severity
+        self.request = request
+        self.response = response
+
+
+
+class CustomTableModelVuln(AbstractTableModel):
+    def __init__(self, extender):
+        self.vulndata = extender._vuln
+        self.callbacks = extender.callbacks
+
+
+    def getRowCount(self):
+        try:
+            return len(self.vulndata)
+        except:
+            return 0
+
+
+    def getColumnCount(self):
+        return 6
+
+    def getColumnName(self, columnIndex):
+        if columnIndex == 0:
+            return "URL"
+        if columnIndex == 1:
+            return "Parameters"
+        if columnIndex == 2:
+            return "Vulnerability"
+        if columnIndex == 3:
+            return "Severity"
+        if columnIndex == 4:
+            return "Request"
+        if columnIndex == 5:
+            return "Response"
+        return ""
+
+
+    def getValueAt(self, rowIndex, columnIndex):
+        if rowIndex < self.getRowCount() and columnIndex < self.getColumnCount():
+            vulnEntry = self.vulndata[rowIndex]
+
+            if columnIndex == 0:
+                return vulnEntry.url
+            elif columnIndex == 1:
+                return vulnEntry.parameters
+            elif columnIndex == 2:
+                return vulnEntry.vulnerability
+            elif columnIndex == 3:
+                return vulnEntry.severity
+            elif columnIndex == 4:
+                return vulnEntry.request
+            elif columnIndex == 5:
+                return vulnEntry.response
+
+
+    def getColumnClass(self, col):
+        return str
+
+    def isCellEditable(self, row, col):
+        return True
+
+    def setValueAt(self, value, row, col):
+        rowData = self.vulndata[row]
+        if col == 0:
+            rowData.url = value
+        elif col == 1:
+            rowData.parameters = value
+        elif col == 2:
+            rowData.vulnerability = value
+        elif col == 3:
+            rowData.severity = value
+        elif col == 4:
+            rowData.request = value
+        elif col == 5:
+            rowData.response = value
+        self.fireTableCellUpdated(row, col)
+
+
+
+    def setValueAt(self, value, rowIndex, columnIndex):
+
+        if rowIndex < self.getRowCount() and columnIndex == 1:
+                vulEntry = self.vulndata[rowIndex]
+                vulEntry.parameters = value
+
+        if rowIndex < self.getRowCount() and columnIndex == 2:
+                vulEntry = self.vulndata[rowIndex]
+                vulEntry.vulnerability = value
+
+        if rowIndex < self.getRowCount() and columnIndex == 3:
+                vulEntry = self.vulndata[rowIndex]
+                vulEntry.severity = value
+        else:
+            self.callbacks.printError("Table is empty")
+
+    def addRow(self, obj):
+        row_data = RowData(*obj)
+        self.vulndata.append(row_data)
+        self.fireTableRowsInserted(len(self.vulndata) - 1, len(self.vulndata) - 1)
+
+
+
+