Merge pull request #5 from elbiazo/dev

Dev

Author: elbiazo

Diff for: README.md

@@ -10,7 +10,8 @@ Generates call tree
 
 ## Releases
 
-* v0.0 -- Beta release
+* 1.0 -- Public Release
+* 0.0 -- Beta Release
 
 ## Minimum Version
 
@@ -51,14 +52,4 @@ When working with really big binaries with alot of xrefs, you would want to chan
 
 **Default Recursion Depth in Setting**
 
-![](images/2022-02-09-16-59-03.png)
-
-## Future Work:
-
-- [ ] Ability to expand selected function
-- [x] Ability to expand all functions
-- [x] Setting and viewing depth of function recursion
-- [x] Move widget to side bar
-- [ ] Generating graph on current call tree
-- [ ] Refactor Code
-- [x] Documentation
+![](images/2022-02-09-16-59-03.png)

Diff for: calltree.py

@@ -3,7 +3,16 @@
     QStandardItemModel,
     QStandardItem,
 )
+from PySide6.QtCore import QSize, Qt
 from PySide6.QtWidgets import QTreeView
+from PySide6.QtWidgets import (
+    QVBoxLayout,
+    QHBoxLayout,
+    QPushButton,
+    QLineEdit,
+    QSpinBox,
+    QTextEdit,
+)
 
 
 class BNFuncItem(QStandardItem):
@@ -13,25 +22,109 @@ def __init__(self, func):
         self.setText(func.name)
 
 
-class CallTreeWidget:
-    def __init__(self, label_name, depth):
-        self.treeview = QTreeView()
-        self.model = QStandardItemModel()
-        self.proxy_model = QSortFilterProxyModel(self.treeview)
+class CurrentFunctionLayout(QHBoxLayout):
+    def __init__(self):
+        super().__init__()
+        self.cur_func_text = QTextEdit()
+        self.cur_func_text.setReadOnly(True)
+        self.cur_func_text.setMaximumHeight(40)
+        self.cur_func_text.setAlignment(Qt.AlignLeft | Qt.AlignTop)
+        self.cur_func_text.setLineWrapMode(QTextEdit.NoWrap)
+        super().addWidget(self.cur_func_text)
+
+
+# Layout with search bar and expand/collapse buttons
+# Takes CallTreeLayout as a parameter
+class CallTreeUtilLayout(QHBoxLayout):
+    def __init__(self, calltree: object):
+        super().__init__()
+        self.calltree = calltree
+        btn_size = QSize(25, 25)
+        self.expand_all_button = QPushButton("E")
+        self.expand_all_button.setFixedSize(btn_size)
+        self.expand_all_button.clicked.connect(self.calltree.expand_all)
+
+        self.collapse_all_button = QPushButton("C")
+        self.collapse_all_button.setFixedSize(btn_size)
+        self.collapse_all_button.clicked.connect(self.calltree.collapse_all)
+
+        self.func_filter = QLineEdit()
+        self.func_filter.textChanged.connect(self.calltree.onTextChanged)
+
+        self.spinbox = QSpinBox()
+        self.spinbox.valueChanged.connect(self.spinbox_changed)
+        self.spinbox.setValue(self.calltree.func_depth)
+        super().addWidget(self.func_filter)
+        super().addWidget(self.expand_all_button)
+        super().addWidget(self.collapse_all_button)
+        super().addWidget(self.spinbox)
+
+    def spinbox_changed(self):
+        self.calltree.func_depth = self.spinbox.value()
+        if self.calltree.cur_func is not None:
+            self.calltree.update_widget(self.calltree.cur_func)
+
+
+class CallTreeLayout(QVBoxLayout):
+    def __init__(self, label_name: str, depth: int, is_caller: bool):
+        super().__init__()
+        self._cur_func = None
+        self._is_caller = is_caller
+        # Creates treeview for all the function calls
+        self._treeview = QTreeView()
+        self._model = QStandardItemModel()
+        self._proxy_model = QSortFilterProxyModel(self.treeview)
         self.proxy_model.setSourceModel(self.model)
 
         self.treeview.setModel(self.proxy_model)
+
+        # Clicking function on treeview will take you to the function
         self.treeview.doubleClicked.connect(self.goto_func)
         self._func_depth = depth
         self._binary_view = None
-        self.label_name = label_name
+        self._label_name = label_name
         self.set_label(self.label_name)
+        super().addWidget(self.treeview)
+        self.util = CallTreeUtilLayout(self)
+        super().addLayout(self.util)
 
     def onTextChanged(self, text):
         self.proxy_model.setRecursiveFilteringEnabled(True)
         self.proxy_model.setFilterRegularExpression(text)
         self.expand_all()
 
+    @property
+    def proxy_model(self):
+        return self._proxy_model
+
+    @property
+    def label_name(self):
+        return self._label_name
+
+    @property
+    def cur_func(self):
+        return self._cur_func
+
+    @cur_func.setter
+    def cur_func(self, cur_func):
+        self._cur_func = cur_func
+
+    @property
+    def is_caller(self):
+        return self._is_caller
+
+    @property
+    def treeview(self):
+        return self._treeview
+
+    @property
+    def func_depth(self):
+        return self._func_depth
+
+    @property
+    def model(self):
+        return self._model
+
     @property
     def binary_view(self):
         return self._binary_view
@@ -79,12 +172,12 @@ def set_func_calls(self, cur_func, cur_std_item, is_caller: bool, depth=0):
                             cur_func_call, new_std_item, is_caller, depth + 1
                         )
 
-    def update_widget(self, cur_func, is_caller):
+    def update_widget(self, cur_func):
         # Clear previous calls
         self.clear()
         call_root_node = self.model.invisibleRootItem()
 
-        if is_caller:
+        if self.is_caller:
             cur_func_calls = list(set(cur_func.callers))
         else:
             cur_func_calls = list(set(cur_func.callees))
@@ -97,7 +190,7 @@ def update_widget(self, cur_func, is_caller):
                 root_std_items.append(BNFuncItem(cur_func_call))
                 cur_std_item = root_std_items[-1]
                 if cur_func != cur_func_call:
-                    self.set_func_calls(cur_func_call, cur_std_item, is_caller)
+                    self.set_func_calls(cur_func_call, cur_std_item, self.is_caller)
 
         call_root_node.appendRows(root_std_items)
         self.expand_all()