diff --git a/src/urh/models/GeneratorTableModel.py b/src/urh/models/GeneratorTableModel.py index 54b6d2e192..133a1e29cf 100644 --- a/src/urh/models/GeneratorTableModel.py +++ b/src/urh/models/GeneratorTableModel.py @@ -1,6 +1,7 @@ from PyQt5.QtCore import Qt, QModelIndex from PyQt5.QtGui import QColor +from urh import constants from urh.models.ProtocolTreeItem import ProtocolTreeItem from urh.models.TableModel import TableModel from urh.signalprocessing.ProtocolAnalyzerContainer import ProtocolAnalyzerContainer @@ -107,6 +108,12 @@ def duplicate_row(self, row: int): self.protocol.duplicate_line(row) self.update() + def add_empty_row_behind(self, row_index: int, num_bits: int): + self.protocol.insert_empty_message(row=row_index+1, + pause=constants.SETTINGS.value("default_fuzzing_pause", 10**6, int), + num_bits=num_bits) + self.update() + def get_selected_label_index(self, row: int, column: int): if self.row_count == 0: return -1 diff --git a/src/urh/signalprocessing/ProtocolAnalyzerContainer.py b/src/urh/signalprocessing/ProtocolAnalyzerContainer.py index 007067c4ba..e27791ad0e 100644 --- a/src/urh/signalprocessing/ProtocolAnalyzerContainer.py +++ b/src/urh/signalprocessing/ProtocolAnalyzerContainer.py @@ -6,6 +6,7 @@ import numpy import time +from urh import constants from urh.models.ProtocolTreeItem import ProtocolTreeItem from urh.signalprocessing.Message import Message from urh.signalprocessing.Modulator import Modulator @@ -66,10 +67,19 @@ def duplicate_line(self, row: int): except Exception as e: logger.error("Duplicating line ", str(e)) + def insert_empty_message(self, row: int, pause: int, num_bits: int): + try: + msg = Message([0]*num_bits, pause=pause, message_type=self.default_message_type) + self.messages.insert(row, msg) + except Exception as e: + logger.error(str(e)) + def fuzz(self, mode: FuzzMode, default_pause=None): result = [] appd_result = result.append + added_message_indices = [] + for i, msg in enumerate(self.messages): labels = msg.active_fuzzing_labels appd_result(msg) @@ -106,12 +116,14 @@ def fuzz(self, mode: FuzzMode, default_pause=None): rssi=msg.rssi, message_type=message_type, modulator_indx=msg.modulator_indx, decoder=msg.decoder, fuzz_created=True) + added_message_indices.append(i+j+1) appd_result(fuz_msg) if j % 10000 == 0: self.qt_signals.current_fuzzing_message_changed.emit(j) self.qt_signals.fuzzing_finished.emit() self.messages = result # type: list[Message] + return added_message_indices def fuzz_successive(self, default_pause=None): """ @@ -119,7 +131,7 @@ def fuzz_successive(self, default_pause=None): Sequentiell heißt, ein Label wird durchgefuzzt und alle anderen Labels bleiben auf Standardwert. Das entspricht dem Vorgang nacheinander immer nur ein Label aktiv zu setzen. """ - self.fuzz(FuzzMode.successive, default_pause=default_pause) + return self.fuzz(FuzzMode.successive, default_pause=default_pause) def fuzz_concurrent(self, default_pause=None): """ @@ -127,14 +139,14 @@ def fuzz_concurrent(self, default_pause=None): gleichzeitig iteriert. Wenn ein Label keine FuzzValues mehr übrig hat, wird der erste Fuzzing Value (per Definition der Standardwert) genommen. """ - self.fuzz(FuzzMode.concurrent, default_pause=default_pause) + return self.fuzz(FuzzMode.concurrent, default_pause=default_pause) def fuzz_exhaustive(self, default_pause=None): """ Führt ein vollständiges Fuzzing durch. D.h. wenn es mehrere Label pro Message gibt, werden alle möglichen Kombinationen erzeugt (Kreuzprodukt!) """ - self.fuzz(FuzzMode.exhaustive, default_pause=default_pause) + return self.fuzz(FuzzMode.exhaustive, default_pause=default_pause) def create_fuzzing_label(self, start, end, msg_index) -> ProtocolLabel: fuz_lbl = self.messages[msg_index].message_type.add_protocol_label(start=start, end=end) diff --git a/src/urh/ui/actions/Fuzz.py b/src/urh/ui/actions/Fuzz.py index 25d839e5d7..afb6a999df 100644 --- a/src/urh/ui/actions/Fuzz.py +++ b/src/urh/ui/actions/Fuzz.py @@ -1,5 +1,6 @@ import copy +import time from PyQt5.QtWidgets import QUndoCommand from urh import constants @@ -13,19 +14,26 @@ def __init__(self, proto_analyzer_container: ProtocolAnalyzerContainer, fuz_mode self.fuz_mode = fuz_mode self.setText("{0} Fuzzing".format(self.fuz_mode)) - self.orig_messages = copy.deepcopy(self.proto_analyzer_container.messages) + self.added_message_indices = [] def redo(self): if constants.SETTINGS.value('use_default_fuzzing_pause', True, bool): default_pause = constants.SETTINGS.value("default_fuzzing_pause", 10**6, int) else: default_pause = None + if self.fuz_mode == "Successive": - self.proto_analyzer_container.fuzz_successive(default_pause=default_pause) + added_indices = self.proto_analyzer_container.fuzz_successive(default_pause=default_pause) elif self.fuz_mode == "Concurrent": - self.proto_analyzer_container.fuzz_concurrent(default_pause=default_pause) + added_indices = self.proto_analyzer_container.fuzz_concurrent(default_pause=default_pause) elif self.fuz_mode == "Exhaustive": - self.proto_analyzer_container.fuzz_exhaustive(default_pause=default_pause) + added_indices = self.proto_analyzer_container.fuzz_exhaustive(default_pause=default_pause) + else: + added_indices = [] + + self.added_message_indices.extend(added_indices) def undo(self): - self.proto_analyzer_container.messages = self.orig_messages + for index in reversed(self.added_message_indices): + del self.proto_analyzer_container.messages[index] + self.added_message_indices.clear() diff --git a/src/urh/ui/views/GeneratorTableView.py b/src/urh/ui/views/GeneratorTableView.py index b57bffd974..908ea61fad 100644 --- a/src/urh/ui/views/GeneratorTableView.py +++ b/src/urh/ui/views/GeneratorTableView.py @@ -1,7 +1,7 @@ from PyQt5.QtCore import Qt, QRect, pyqtSignal, pyqtSlot from PyQt5.QtGui import QDragMoveEvent, QDragEnterEvent, QPainter, QBrush, QColor, QPen, QDropEvent, QDragLeaveEvent, \ - QContextMenuEvent -from PyQt5.QtWidgets import QActionGroup + QContextMenuEvent, QIcon +from PyQt5.QtWidgets import QActionGroup, QInputDialog from PyQt5.QtWidgets import QHeaderView, QAbstractItemView, QStyleOption, QMenu @@ -156,6 +156,11 @@ def create_context_menu(self) -> QMenu: fuzzing_action.triggered.connect(self.on_fuzzing_action_triggered) menu.addSeparator() + add_message_action = menu.addAction("Add empty message...") + add_message_action.setIcon(QIcon.fromTheme("list-add")) + add_message_action.triggered.connect(self.on_add_message_action_triggered) + + if self.model().row_count > 0: column_menu = menu.addMenu("Add column") insert_column_left_action = column_menu.addAction("on the left") @@ -232,3 +237,13 @@ def on_encoding_action_triggered(self): for row in self.selected_rows: self.model().protocol.messages[row].decoder = self.encoding_actions[self.sender()] self.encodings_updated.emit() + + @pyqtSlot() + def on_add_message_action_triggered(self): + row = self.rowAt(self.context_menu_pos.y()) + num_bits, ok = QInputDialog.getInt(self, self.tr("How many bits shall the new message have?"), + self.tr("Number of bits:"), 42, 1) + if ok: + self.model().add_empty_row_behind(row, num_bits) + if self.model().rowCount() == 1: + self.resize_columns() diff --git a/tests/test_generator.py b/tests/test_generator.py index a514de64af..1fa7bc8da8 100644 --- a/tests/test_generator.py +++ b/tests/test_generator.py @@ -106,11 +106,11 @@ def test_close_signal(self): self.assertEqual(1, 1) def test_create_table_context_menu(self): - # Context menu should be empty if table is empty + # Context menu should only contain one item (add new message) self.assertEqual(self.form.generator_tab_controller.table_model.rowCount(), 0) self.form.generator_tab_controller.ui.tableMessages.context_menu_pos = QPoint(0, 0) menu = self.form.generator_tab_controller.ui.tableMessages.create_context_menu() - self.assertEqual(len(menu.actions()), 0) + self.assertEqual(len(menu.actions()), 1) # Add data to test entries in context menu self.add_signal_to_form("ask.complex") @@ -122,13 +122,35 @@ def test_create_table_context_menu(self): self.assertGreater(self.form.generator_tab_controller.table_model.rowCount(), 0) menu = self.form.generator_tab_controller.ui.tableMessages.create_context_menu() n_items = len(menu.actions()) - self.assertGreater(n_items, 0) + self.assertGreater(n_items, 1) # If there is a selection, additional items should be present in context menu gframe.ui.tableMessages.selectRow(0) menu = self.form.generator_tab_controller.ui.tableMessages.create_context_menu() self.assertGreater(len(menu.actions()), n_items) + def test_add_empty_row_behind(self): + self.assertEqual(self.form.generator_tab_controller.table_model.rowCount(), 0) + gframe = self.form.generator_tab_controller + gframe.ui.cbViewType.setCurrentIndex(0) + gframe.table_model.add_empty_row_behind(-1, 30) + self.assertEqual(self.form.generator_tab_controller.table_model.rowCount(), 1) + + # Add data to test + self.add_signal_to_form("ask.complex") + + index = gframe.tree_model.createIndex(0, 0, gframe.tree_model.rootItem.children[0].children[0]) + mimedata = gframe.tree_model.mimeData([index]) + gframe.table_model.dropMimeData(mimedata, 1, -1, -1, gframe.table_model.createIndex(0, 0)) + self.assertEqual(self.form.generator_tab_controller.table_model.rowCount(), 2) + self.assertNotEqual(len(self.form.generator_tab_controller.table_model.display_data[1]), 30) + gframe.table_model.add_empty_row_behind(0, 30) + self.assertEqual(self.form.generator_tab_controller.table_model.rowCount(), 3) + self.assertEqual(len(self.form.generator_tab_controller.table_model.display_data[1]), 30) + self.assertNotEqual(len(self.form.generator_tab_controller.table_model.display_data[2]), 30) + + + def test_create_fuzzing_list_view_context_menu(self): # Context menu should be empty if table is empty self.assertEqual(self.form.generator_tab_controller.table_model.rowCount(), 0)