class MemoryTest(unittest.TestCase):
    def setUp(self) -> None:
        self.memory = Memory()
        # just attach to own process, we're overriding the read values anyway
        self.memory.process = psutil.Process()

    @patch.object(psutil.Process, "memory_maps", autospec=True)
    @patch.object(Memory, "read", autospec=True)
    def test_scan(self, read_mock: Mock, mem_map_mock: Mock):
        read_mock.side_effect = [123, 123]
        mem_map_mock.return_value = [InternalMmap("0-8", 8, "r--w")]
        self.assertEqual(self.memory.scan("123"), [Value(0, 0, 123), Value(0, 4, 123)])

    @patch.object(psutil.Process, "memory_maps", autospec=True)
    @patch.object(Memory, "read", autospec=True)
    def test_scan_empty_search(self, read_mock: Mock, mem_map_mock: Mock):
        read_mock.side_effect = [123, 123]
        mem_map_mock.return_value = [InternalMmap("0-8", 8, "r--w")]

    def test_scan_cull_empty(self):
        self.assertEqual(self.memory._scan_cull(Value(0, 0, 123)), [])

    def test_scan_cull(self):
        maps = [InternalMmap("0-32", 32, "r--w")]
        # 15 because 32/4 = 8, since we need twice the scan area, we have a total of 16, then we
        # append a false value too to make sure the filtering works
        packed = [Value(0, 0, 123)] * 15
        packed.append(Value(0, 0, 12))

        with patch.object(Memory, "read", side_effect=packed), \
                patch("builtins.open", mock_open()), \
                patch("psutil.Process.__init__", return_value=None), \
                patch("psutil.Process.memory_maps", return_value=maps):
            self.memory.attach(123)
            self.assertEqual(len(self.memory.scan(123)), 8)
            self.assertEqual(len(self.memory.scan(123)), 7)

    def test_attach(self):
        with patch("builtins.open", mock_open()), patch("psutil.Process.__init__",
                                                        return_value=None):
            self.memory.attach(123)
            self.assertEqual(self.memory.pid, 123)

    def test_detach(self):
        with patch("builtins.open", mock_open()), patch("psutil.Process.__init__",
                                                        return_value=None):
            self.memory.attach(123)
            self.memory.detach()
            self.assertEqual(self.memory.pid, 0)
Exemple #2
0
class MainWindow(QMainWindow, Ui_MainWindow):
    def __init__(self):
        super().__init__()
        self.memory = Memory()

        self.setupUi(self)
        self.scan_widget.setEnabled(False)

        self.found_table.setModel(FoundAddressModel(self))
        self.found_table.horizontalHeader().setSectionResizeMode(
            0, QHeaderView.Stretch)

        self.new_scan.clicked.connect(self.new_scan_clicked)
        self.next_scan.clicked.connect(self.next_scan_clicked)
        self.actionAttach.triggered.connect(self.attach_triggered)
        self.found_table.doubleClicked.connect(self.found_table_double_clicked)
        self.saved_results.doubleClicked.connect(
            self.saved_model_double_clicked)

    @pyqtSlot()
    def new_scan_clicked(self):
        if self.next_scan.isEnabled():
            self.next_scan.setEnabled(False)
            self.found_table.model().set_values([])
            self.memory.reset_scan()
        else:
            # the isnumeric is a simple workaround, since I'll support strings eventually
            # until then it can stay
            if len(self.search_for.text()) > 0 and self.search_for.text(
            ).lstrip("-").isnumeric():
                self.next_scan_clicked()
                self.next_scan.setEnabled(True)

    @pyqtSlot()
    def next_scan_clicked(self):
        # TODO: fix looking for negative values, probably have to fiddle with the type system
        try:
            scan_type = {
                0: operator.eq,
                1: operator.gt,
                2: operator.lt,
                3: operator.ne
            }
            match_index = self.scan_matching.currentIndex()

            values = self.memory.scan(self.search_for.text(),
                                      Type(self.scan_byte_size.currentIndex()),
                                      scan_type[match_index])
            self.amount_found.setText("Found: {}".format(len(values)))
            self.found_table.model().set_values(values)
        except NoSuchProcess:
            QMessageBox(QMessageBox.NoIcon, "Error", "No readable memory",
                        QMessageBox.Ok, self).exec_()

    @pyqtSlot()
    def attach_triggered(self):
        process_view = ProcessView()
        pid = process_view.exec_()
        if pid != 0:
            self.memory.attach(pid)
            self.scan_widget.setEnabled(True)

    @pyqtSlot(QModelIndex)
    def found_table_double_clicked(self, index: QModelIndex):
        if not index.isValid():
            return
        clicked_value = self.found_table.model().get_value(index.row())
        self.saved_results.model().append_row(clicked_value)

    @pyqtSlot(QModelIndex)
    def saved_model_double_clicked(self, index: QModelIndex):
        if not index.isValid():
            return
        column = index.column()

        if column == SavedAddressHeaderEnum.VALUE:
            form = WriteForm(self)
            x = form.exec_()
            text = form.value.text()
            if x and len(text) != 0:
                item: TreeItem = self.saved_results.model().get_item(index)
                value = item.get_internal_pointer()
                try:
                    value.write(
                        value.type.parse_value(text, form.ishex.isChecked()))
                except struct.error:
                    QMessageBox(
                        QMessageBox.NoIcon, "Error",
                        "You can't write a larger number to "
                        "memory than you currently have set "
                        "as a type.\nSee limits.h for more "
                        "information", QMessageBox.Ok, self).exec_()
                except ValueError:
                    # this is a kinda temporary fix, I'm planning on adding a validator or some
                    # shit later, not entirely sure how to add them yet so this will do
                    QMessageBox(
                        QMessageBox.NoIcon, "Error",
                        "Error parsing value, you probably "
                        "entered text when trying to write "
                        "an integer...", QMessageBox.Ok, self).exec_()