From 2a6dac19b56864b7420d9ceeedd0ba020272fb01 Mon Sep 17 00:00:00 2001
From: Matt Johnston <matt@ucc.asn.au>
Date: Sun, 30 Jan 2022 13:37:20 +0800
Subject: [PATCH] Use venv for test_aslr

Otherwise we can't find the psutil dependency
---
 test/Makefile.in      |  6 ++++--
 test/test_aslr.py     |  6 +++++-
 test/test_dropbear.py | 14 ++++++++++++++
 3 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/test/Makefile.in b/test/Makefile.in
index 441ac41..b2c8d43 100644
--- a/test/Makefile.in
+++ b/test/Makefile.in
@@ -1,12 +1,14 @@
 srcdir=@srcdir@
 
+SHELL=bash
+
 all: test
 
 test: venv/bin/pytest fakekey
-	./venv/bin/pytest --hostkey=fakekey --dbclient=../dbclient --dropbear=../dropbear $(srcdir)
+	(source ./venv/bin/activate; pytest --hostkey=fakekey --dbclient=../dbclient --dropbear=../dropbear $(srcdir) )
 
 one: venv/bin/pytest fakekey
-	./venv/bin/pytest --hostkey=fakekey --dbclient=../dbclient --dropbear=../dropbear $(srcdir) -k exit
+	(source ./venv/bin/activate; pytest --hostkey=fakekey --dbclient=../dbclient --dropbear=../dropbear $(srcdir) -k exit)
 
 fakekey:
 	../dropbearkey -t ecdsa -f $@
diff --git a/test/test_aslr.py b/test/test_aslr.py
index 6f997b8..1234e35 100644
--- a/test/test_aslr.py
+++ b/test/test_aslr.py
@@ -8,7 +8,11 @@ def test_reexec(request, dropbear):
 	Tests that two consecutive connections have different address layouts.
 	This indicates that re-exec makes ASLR work
 	"""
-	cmd = (Path(request.node.fspath).parent / "parent_dropbear_map.py").resolve()
+	map_script = (Path(request.node.fspath).parent / "parent_dropbear_map.py").resolve()
+	# run within the same venv, for python deps
+	activate = own_venv_command()
+	cmd = f"{activate}; {map_script}"
+	print(cmd)
 	r = dbclient(request, cmd, capture_output=True, text=True)
 	map1 = r.stdout.rstrip()
 	print(r.stderr, file=sys.stderr)
diff --git a/test/test_dropbear.py b/test/test_dropbear.py
index 05ea953..98d748e 100644
--- a/test/test_dropbear.py
+++ b/test/test_dropbear.py
@@ -57,6 +57,20 @@ def dbclient(request, *args, **kwargs):
 		# wait for response
 		return subprocess.run(full_args, **kwargs)
 
+def own_venv_command():
+	""" Returns a command to run as a prefix to get the same venv
+	as the current running Python. Returns '' on not a virtualenv
+	"""
+	try:
+		print(os.environ)
+		venv = os.environ['VIRTUAL_ENV']
+		print(venv)
+	except KeyError:
+		return ""
+
+	# note: bash/zsh unix specific
+	return f"source {venv}/bin/activate"
+
 class HandleTcp(socketserver.ThreadingMixIn, socketserver.TCPServer):
 	""" Listens for a single incoming request, sends a response if given,
 	and returns the inbound data.