From 0ef18a7cbaf2e7176031f8b386ea77825c5efd29 Mon Sep 17 00:00:00 2001
From: Jim Lloyd <jim@coalitionsproject.org>
Date: Mon, 6 Aug 2012 14:09:35 -0700
Subject: Adds --dummy-certs option to specify certdir

If --dummy-certs=CERTSDIR is provided, use CERTSDIR as the location
for generating/finding the dummy certs. And in this case, preserve
the CERTSDIR directory on exit.
---
 libmproxy/cmdline.py |  6 ++++++
 libmproxy/proxy.py   | 24 ++++++++++++++++++------
 2 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/libmproxy/cmdline.py b/libmproxy/cmdline.py
index 6d26b74d..302cfd29 100644
--- a/libmproxy/cmdline.py
+++ b/libmproxy/cmdline.py
@@ -279,4 +279,10 @@ def common_options(parser):
     )
     parser.add_option_group(group)
 
+    group.add_option(
+        "--dummy-certs", action="store",
+        type = "str", dest = "certdir", default=None,
+        help = "Generated dummy certs directory."
+    )
+
     proxy.certificate_option_group(parser)
diff --git a/libmproxy/proxy.py b/libmproxy/proxy.py
index 7fb84530..ae0e4415 100644
--- a/libmproxy/proxy.py
+++ b/libmproxy/proxy.py
@@ -36,12 +36,12 @@ class Log(controller.Msg):
 
 
 class ProxyConfig:
-    def __init__(self, certfile = None, cacert = None, clientcerts = None, cert_wait_time=0, no_upstream_cert=False, body_size_limit = None, reverse_proxy=None, transparent_proxy=None):
+    def __init__(self, certfile = None, cacert = None, clientcerts = None, cert_wait_time=0, no_upstream_cert=False, body_size_limit = None, reverse_proxy=None, transparent_proxy=None, certdir = None):
         assert not (reverse_proxy and transparent_proxy)
         self.certfile = certfile
         self.cacert = cacert
         self.clientcerts = clientcerts
-        self.certdir = None
+        self.certdir = certdir
         self.cert_wait_time = cert_wait_time
         self.no_upstream_cert = no_upstream_cert
         self.body_size_limit = body_size_limit
@@ -399,8 +399,13 @@ class ProxyServer(tcp.TCPServer):
         except socket.error, v:
             raise ProxyServerError('Error starting proxy server: ' + v.strerror)
         self.masterq = None
-        self.certdir = tempfile.mkdtemp(prefix="mitmproxy")
-        config.certdir = self.certdir
+        if config.certdir:
+            self.certdir = config.certdir
+            self.remove_certdir = False
+        else:
+            self.certdir = tempfile.mkdtemp(prefix="mitmproxy")
+            config.certdir = self.certdir
+            self.remove_certdir = True
         self.apps = AppRegistry()
 
     def start_slave(self, klass, masterq):
@@ -417,7 +422,8 @@ class ProxyServer(tcp.TCPServer):
 
     def handle_shutdown(self):
         try:
-            shutil.rmtree(self.certdir)
+            if self.remove_certdir:
+                shutil.rmtree(self.certdir)
         except OSError:
             pass
 
@@ -513,6 +519,11 @@ def process_proxy_options(parser, options):
         if not os.path.exists(options.clientcerts) or not os.path.isdir(options.clientcerts):
             parser.error("Client certificate directory does not exist or is not a directory: %s"%options.clientcerts)
 
+    if options.certdir:
+        options.certdir = os.path.expanduser(options.certdir)
+        if not os.path.exists(options.certdir) or not os.path.isdir(options.certdir):
+            parser.error("Dummy cert directory does not exist or is not a directory: %s"%options.certdir)
+
     return ProxyConfig(
         certfile = options.cert,
         cacert = cacert,
@@ -521,5 +532,6 @@ def process_proxy_options(parser, options):
         body_size_limit = body_size_limit,
         no_upstream_cert = options.no_upstream_cert,
         reverse_proxy = rp,
-        transparent_proxy = trans
+        transparent_proxy = trans,
+        certdir = options.certdir
     )
-- 
cgit v1.2.3