From ed413d1ea904a293ed62bce3d2b32c534c539eed Mon Sep 17 00:00:00 2001
From: theguy147 <37738506+theguy147@users.noreply.github.com>
Date: Sun, 10 Oct 2021 17:26:01 +0200
Subject: [PATCH] feat: measure coverage (#738)

---
 .gitignore       |  2 ++
 Makefile         | 11 ++++++++++-
 requirements.txt |  1 +
 tests/helpers.py | 24 +++++++++++++++++-------
 4 files changed, 30 insertions(+), 8 deletions(-)

diff --git a/.gitignore b/.gitignore
index 19e3cbea..62825775 100644
--- a/.gitignore
+++ b/.gitignore
@@ -10,3 +10,5 @@ tests/pylint3.html
 _build
 debug.log
 .pytest_cache
+.coverage
+htmlcov
diff --git a/Makefile b/Makefile
index 5b9e8779..06f335e1 100644
--- a/Makefile
+++ b/Makefile
@@ -8,6 +8,7 @@ PYLINT_COMMON_PARAMETERS := --jobs=$(PYLINT_JOBS) --suggestion-mode=$(PYLINT_SUG
 PYLINT_GEF_PARAMETERS := --disable=$(PYLINT_DISABLE) --enable=$(PYLINT_ENABLE) $(PYLINT_COMMON_PARAMETERS)
 PYLINT_TEST_PARAMETERS := --disable=$(PYLINT_DISABLE) --enable=$(PYLINT_TEST_ENABLE) $(PYLINT_COMMON_PARAMETERS)
 TARGET := $(shell lscpu | head -1 | sed -e 's/Architecture:\s*//g')
+COVERAGE_DIR ?= /tmp/cov
 GEF_PATH ?= $(shell readlink -f gef.py)
 TMPDIR ?= /tmp
 PYTEST_PARAMETERS := --verbose -n $(NB_CORES)
@@ -31,12 +32,20 @@ testbins: $(TMPDIR) $(wildcard tests/binaries/*.c)
 
 clean:
 	TMPDIR=$(TMPDIR) $(MAKE) -j $(NB_CORES) -C tests/binaries clean
-	@rm -rf $(TMPDIR)
+	@rm -rf $(TMPDIR)/gef-* $(TMPDIR)/gef.py || true
 
 lint:
 	python3 -m pylint $(PYLINT_GEF_PARAMETERS) $(GEF_PATH)
 	python3 -m pylint $(PYLINT_TEST_PARAMETERS) $(wildcard tests/*.py)
 
+coverage:
+	@! ( [ -d $(COVERAGE_DIR) ] && echo "COVERAGE_DIR=$(COVERAGE_DIR) exists already")
+	@mkdir -p $(COVERAGE_DIR)
+	@COVERAGE_DIR=$(COVERAGE_DIR) $(MAKE) test
+	@coverage combine $(COVERAGE_DIR)/*
+	@coverage html --include "*/gef.py"
+	@rm -rf $(COVERAGE_DIR)
+
 $(TMPDIR):
 	mkdir -p $@
 
diff --git a/requirements.txt b/requirements.txt
index 1b63b555..22d4a563 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -5,3 +5,4 @@ ropper
 unicorn
 pytest
 pytest-xdist
+coverage
diff --git a/tests/helpers.py b/tests/helpers.py
index 55771e53..c7642640 100644
--- a/tests/helpers.py
+++ b/tests/helpers.py
@@ -14,6 +14,7 @@ STRIP_ANSI_DEFAULT = True
 DEFAULT_CONTEXT = "-code -stack"
 ARCH = (os.getenv("GEF_CI_ARCH") or platform.machine()).lower()
 CI_VALID_ARCHITECTURES = ("x86_64", "i686", "aarch64", "armv7l")
+COVERAGE_DIR = os.getenv("COVERAGE_DIR", "")
 
 CommandType = NewType("CommandType", Union[str, Iterable[str]])
 
@@ -37,16 +38,25 @@ def gdb_run_cmd(cmd: CommandType, before: CommandType = (), after: CommandType =
                 target: Path = DEFAULT_TARGET, strip_ansi: bool = STRIP_ANSI_DEFAULT) -> str:
     """Execute a command inside GDB. `before` and `after` are lists of commands to be executed
     before (resp. after) the command to test."""
-    command = [
-        "gdb", "-q", "-nx",
-        "-ex", f"source {GEF_PATH}",
-        "-ex", "gef config gef.debug True"
-    ]
-
+    command = ["gdb", "-q", "-nx"]
+    if COVERAGE_DIR:
+        coverage_file = Path(COVERAGE_DIR) / os.getenv("PYTEST_XDIST_WORKER")
+        command += _add_command([
+            "pi from coverage import Coverage",
+            f"pi cov = Coverage(data_file=\"{coverage_file}\","
+            "auto_data=True, branch=True)",
+            "pi cov.start()",
+        ])
+    command += _add_command([
+        f"source {GEF_PATH}",
+        "gef config gef.debug True",
+    ])
     command += _add_command(before)
     command += _add_command(cmd)
     command += _add_command(after)
-    command += ["-ex", "quit", "--", target]
+    if COVERAGE_DIR:
+        command += _add_command(["pi cov.stop()", "pi cov.save()"])
+    command += ["-ex", "quit", "--", str(target)]
 
     lines = subprocess.check_output(command, stderr=subprocess.STDOUT).strip().splitlines()
     output = b"\n".join(lines)