From f1aa0c8f5cd4dc5addb02c800f65d4e2f30e3a55 Mon Sep 17 00:00:00 2001
From: anthraxx <levente@leventepolyak.net>
Date: Wed, 31 Mar 2021 22:23:42 +0200
Subject: [PATCH] feature(profile): use a simple module based approach to
 define profiles

---
 profiling/context.py |  9 +++++++++
 profiling/profile.sh | 40 +++++++++++++++++++++++++++++-----------
 2 files changed, 38 insertions(+), 11 deletions(-)
 create mode 100644 profiling/context.py

diff --git a/profiling/context.py b/profiling/context.py
new file mode 100644
index 000000000..8aa0564d9
--- /dev/null
+++ b/profiling/context.py
@@ -0,0 +1,9 @@
+import pwndbg.commands.context
+
+
+def warmup():
+    pwndbg.commands.context.context()
+
+
+def run():
+    pwndbg.commands.context.context()
diff --git a/profiling/profile.sh b/profiling/profile.sh
index f377ebef9..9224765cf 100755
--- a/profiling/profile.sh
+++ b/profiling/profile.sh
@@ -1,21 +1,39 @@
 #!/bin/bash
+
+if ! (( $# )); then
+	cat <<- _EOF_
+		$0: [profile-script]
+
+		Example: $0 context.py
+_EOF_
+	exit 1
+fi
+
+module=$(basename "${1/.py/}")
+basedir=$(dirname "$0")
+
 # Quick and dirty script to profile pwndbg using cProfile.
-make test > /dev/null
-git log --abbrev-commit --pretty=oneline HEAD^..HEAD
-# To profile first run, remove -ex "context".
-gdb ./test \
-  -ex "source ../gdbinit.py" \
-  -ex "b main" -ex "r" \
-  -ex "context" \
-  -ex "python import cProfile; cProfile.run('pwndbg.commands.context.context()', 'stats')" \
-  -ex "quit"
+make -C "${basedir}" test > /dev/null
+
+gdb "${basedir}/test" \
+	-ex "source ${basedir}/../gdbinit.py" \
+	-ex "b main" \
+	-ex "r" \
+	-ex "python
+import cProfile;
+import profiling.${module} as profile;
+profile.warmup();
+cProfile.run('profile.run()', '${basedir}/stats')" \
+	-ex "quit"
 
 python3 -c "
 import pstats
-p = pstats.Stats('stats')
+p = pstats.Stats('${basedir}/stats')
 p.strip_dirs().sort_stats('tottime').print_stats(20)
 "
 
 if command -v pyprof2calltree >/dev/null 2>&1 && command -v kcachegrind >/dev/null 2>&1; then
-	pyprof2calltree -k -i stats
+	pyprof2calltree -k -i "${basedir}/stats"
 fi
+
+# vim: ts=4 sw=4 noet