From a4b6deb10026d9615a31f42c4bc7771a29dc7bf3 Mon Sep 17 00:00:00 2001
From: Sergey Polzunov <sergey@polzunov.com>
Date: Wed, 30 Sep 2015 17:00:29 +0200
Subject: [PATCH] adding key password support for SSL keys

---
 cabby/abstract.py    | 10 ++++++++--
 cabby/cli/commons.py |  4 ++++
 cabby/cli/poll.py    |  6 ++++--
 cabby/dispatcher.py  |  9 +++++++--
 4 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/cabby/abstract.py b/cabby/abstract.py
index a5e756b..2e31456 100644
--- a/cabby/abstract.py
+++ b/cabby/abstract.py
@@ -43,8 +43,8 @@ def __init__(self, host=None, discovery_path=None, port=None,
         self.log = logging.getLogger("%s.%s" % (self.__module__,
                                                 self.__class__.__name__))
 
-    def set_auth(self, cert_file=None, key_file=None, username=None,
-                 password=None, jwt_auth_url=None):
+    def set_auth(self, cert_file=None, key_file=None, key_password=None,
+                 username=None, password=None, jwt_auth_url=None):
         '''Set authentication credentials.
 
         ``jwt_auth_url`` is required for JWT based authentication. If
@@ -58,6 +58,11 @@ def set_auth(self, cert_file=None, key_file=None, username=None,
         :param str key_file: a path to SSL key file
         :param str username: username, used in basic auth or JWT auth
         :param str password: password, used in basic auth or JWT auth
+        :param str key_password: same argument as in
+            ``ssl.SSLContext.load_cert_chain`` - may be a function to call
+            to get the password for decrypting the private key or
+            string/bytes/bytearray. It will only be called if the private
+            key is encrypted and a password is necessary.
         :param str jwt_auth_url: URL used to obtain JWT token
         '''
 
@@ -66,6 +71,7 @@ def set_auth(self, cert_file=None, key_file=None, username=None,
             'key_file': key_file,
             'username': username,
             'password': password,
+            'key_password': key_password,
             'jwt_url': jwt_auth_url
         }
 
diff --git a/cabby/cli/commons.py b/cabby/cli/commons.py
index f8308ca..81d2f87 100644
--- a/cabby/cli/commons.py
+++ b/cabby/cli/commons.py
@@ -45,6 +45,9 @@ def get_basic_arg_parser():
 
     parser.add_argument("--cert", dest="cert", help="certificate file")
     parser.add_argument("--key", dest="key", help="private key file")
+    parser.add_argument(
+        "--key-password", dest="key_password",
+        help="private key password")
 
     parser.add_argument(
         "--username", dest="username",
@@ -132,6 +135,7 @@ def run_client(parser, run_func):
         key_file=args.key,
         username=args.username,
         password=args.password,
+        key_password=args.key_password,
         jwt_auth_url=args.jwt_auth_url,
     )
 
diff --git a/cabby/cli/poll.py b/cabby/cli/poll.py
index 99dae5c..e8c742a 100755
--- a/cabby/cli/poll.py
+++ b/cabby/cli/poll.py
@@ -72,13 +72,15 @@ def save_to_dir(dest_dir, collection, content_block, as_raw):
     filename = generate_filename(collection, content_block)
     path = os.path.abspath(os.path.join(dest_dir, filename))
 
-    with open(path, 'w') as f:
+    with open(path, 'wb') as f:
         if as_raw:
             content = content_block.raw.to_xml(pretty_print=True)
         else:
             content = content_block.content
 
-        f.write(content)
+        f.write(
+            content if isinstance(content, bytes)
+            else content.encode('utf-8'))
 
     log.info("Content block saved to %s", path)
 
diff --git a/cabby/dispatcher.py b/cabby/dispatcher.py
index 4e01828..f6c2e87 100644
--- a/cabby/dispatcher.py
+++ b/cabby/dispatcher.py
@@ -30,7 +30,9 @@
 
 
 def _set_auth_details(tclient, cert_file=None, key_file=None,
-                      username=None, password=None, is_jwt=False):
+                      key_password=None, username=None,
+                      password=None, is_jwt=False):
+
     tls_auth = (cert_file and key_file)
     basic_auth = (not is_jwt and username and password)
 
@@ -41,13 +43,15 @@ def _set_auth_details(tclient, cert_file=None, key_file=None,
             'key_file': key_file,
             'cert_file': cert_file,
             'username': username,
-            'password': password
+            'password': password,
+            'key_password': key_password
         }
     elif tls_auth:
         tclient.set_auth_type(HttpClient.AUTH_CERT)
         credentials = {
             'key_file': key_file,
             'cert_file': cert_file,
+            'key_password': key_password
         }
     elif basic_auth:
         tclient.set_auth_type(HttpClient.AUTH_BASIC)
@@ -117,6 +121,7 @@ def send_taxii_request(url, request, headers, auth_details=None,
         key_file=auth_details.get('key_file'),
         username=auth_details.get('username'),
         password=auth_details.get('password'),
+        key_password=auth_details.get('key_password'),
         is_jwt=bool(auth_details.get('jwt_url'))
     )