Merge pull request #1307 from mvanotti/getSourceAddress

Add getSourceAddress to PathConstraint
JonathanSalwan · Feb 29, 2024 · 06610ca · 06610ca
2 parents a03142d + 2b835cb
commit 06610ca
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 0 deletions.
diff --git a/src/libtriton/bindings/python/objects/pyPathConstraint.cpp b/src/libtriton/bindings/python/objects/pyPathConstraint.cpp
@@ -83,6 +83,9 @@ instruction address.
 - <b>string getComment(void)</b><br>
 Returns the comment (if exists) of the path constraint.
 
+- <b>integer getSourceAddress(void)</b><br>
+Returns the source address of the branch.
+
 - <b>integer getTakenAddress(void)</b><br>
 Returns the address of the taken branch.
 
@@ -147,6 +150,16 @@ namespace triton {
       }
 
 
+      static PyObject* PathConstraint_getSourceAddress(PyObject* self, PyObject* noarg) {
+        try {
+          return PyLong_FromUint64(PyPathConstraint_AsPathConstraint(self)->getSourceAddress());
+        }
+        catch (const triton::exceptions::Exception& e) {
+          return PyErr_Format(PyExc_TypeError, "%s", e.what());
+        }
+      }
+
+
       static PyObject* PathConstraint_getTakenAddress(PyObject* self, PyObject* noarg) {
         try {
           return PyLong_FromUint64(PyPathConstraint_AsPathConstraint(self)->getTakenAddress());
@@ -207,6 +220,7 @@ namespace triton {
       PyMethodDef PathConstraint_callbacks[] = {
         {"getBranchConstraints",        PathConstraint_getBranchConstraints,      METH_NOARGS,    ""},
         {"getComment",                  PathConstraint_getComment,                METH_NOARGS,    ""},
+        {"getSourceAddress",            PathConstraint_getSourceAddress,          METH_NOARGS,    ""},
         {"getTakenAddress",             PathConstraint_getTakenAddress,           METH_NOARGS,    ""},
         {"getTakenPredicate",           PathConstraint_getTakenPredicate,         METH_NOARGS,    ""},
         {"getThreadId",                 PathConstraint_getThreadId,               METH_NOARGS,    ""},

diff --git a/src/testers/unittests/test_path_constraint.py b/src/testers/unittests/test_path_constraint.py
@@ -50,6 +50,10 @@ def test_getTakenPredicate(self):
         pc = self.ctx.getPathConstraints()[0]
         self.assertEqual(pc.getTakenPredicate().evaluate(), 1)
 
+    def test_getSourceAddress(self):
+        pc = self.ctx.getPathConstraints()[0]
+        self.assertEqual(pc.getSourceAddress(), 17)
+
     def test_getTakenAddress(self):
         pc = self.ctx.getPathConstraints()[0]
         self.assertEqual(pc.getTakenAddress(), 108)