Support a custom certificates directory in test/server.py

2 files changed, 15 insertions, 9 deletions

diff --git a/test/proxy_core.py b/test/proxy_core.py
index 5a23278..da90455 100644
--- a/test/proxy_core.py
+++ b/test/proxy_core.py
@@ -8,17 +8,19 @@ The core for a "virtual network" proxy
 Be sure to run this inside your intended certificates directory.
 """
 
-import os, socket, ssl, sys, threading, time
+import os, socket, ssl, subprocess, sys, threading, time
 from http.server import HTTPServer, BaseHTTPRequestHandler
 from socketserver import ThreadingMixIn
-from subprocess import Popen, PIPE
 
-gen_cert_req, lock = 'openssl req -new -key cert.key -subj /CN=%s', threading.Lock()
-sign_cert_req = 'openssl x509 -req -days 3650 -CA ca.crt -CAkey ca.key -set_serial %d -out %s'
+gen_cert_req, lock = 'openssl req -new -key %scert.key -subj /CN=%s', threading.Lock()
+sign_cert_req = 'openssl x509 -req -days 3650 -CA %sca.crt -CAkey %sca.key -set_serial %d -out %s'
 
+def popen(command, *args, **kwargs):
+    return subprocess.Popen((command % args).split(' '), **kwargs)
 
 class ProxyRequestHandler(BaseHTTPRequestHandler):
 	"""Handles a network request made to the proxy"""
+	certdir = ''
 	
 	def log_error(self, format, *args):
 		# suppress "Request timed out: timeout('timed out',)"
@@ -29,17 +31,17 @@ class ProxyRequestHandler(BaseHTTPRequestHandler):
 
 	def do_CONNECT(self):
 		hostname = self.path.split(':')[0]
-		certpath = '%s.crt' % (hostname if hostname != 'ca' else 'CA')
+		certpath = '%s%s.crt' % (certdir, hostname if hostname != 'ca' else 'CA')
 
 		with lock:
 			if not os.path.isfile(certpath):
-				p1 = Popen((gen_cert_req % hostname).split(' '), stdout=PIPE).stdout
-				Popen((sign_cert_req % (time.time() * 1000, certpath)).split(' '), stdin=p1, stderr=PIPE).communicate()
+				p1 = popen(gen_cert_req, certdir, hostname, stdout=subprocess.PIPE).stdout
+				popen(sign_cert_req, certdir, certdir, time.time() * 1000, certpath, stdin=p1, stderr=subprocess.PIPE).communicate()
 
 		self.send_response(200)
 		self.end_headers()
 
-		self.connection = ssl.wrap_socket(self.connection, keyfile='cert.key', certfile=certpath, server_side=True)
+		self.connection = ssl.wrap_socket(self.connection, keyfile=certdir+'cert.key', certfile=certpath, server_side=True)
 		self.rfile = self.connection.makefile('rb', self.rbufsize)
 		self.wfile = self.connection.makefile('wb', self.wbufsize)
 
diff --git a/test/server.py b/test/server.py
index d396495..83a72fa 100644
--- a/test/server.py
+++ b/test/server.py
@@ -39,6 +39,8 @@ mime_types = {
 }
 
 class RequestHijacker(ProxyRequestHandler):
+	certdir = global_certdir
+	
 	def handle_request(self, req_body):
 		path_components = self.path.split('?', maxsplit=1)
 		path = path_components[0]
@@ -99,8 +101,10 @@ class RequestHijacker(ProxyRequestHandler):
 		
 
 
-def do_an_internet(port=1337):
+def do_an_internet(certdir, port):
 	"""Start up the proxy/server"""
+	global global_certdir
+	global_certdir = certdir
 	
 	httpd = ThreadingHTTPServer(('', port), RequestHijacker)
 	httpd.serve_forever()