From 49eae23f920be56178a6cedbceb977e0de432f3b Mon Sep 17 00:00:00 2001
From: David Wilson <dw@botanicus.net>
Date: Fri, 4 May 2018 22:51:54 +0100
Subject: [PATCH] issue #218: ansibe: use Secret and Blob types.

---
 ansible_mitogen/connection.py | 16 ++++++++++------
 ansible_mitogen/services.py   |  9 ++++++---
 2 files changed, 16 insertions(+), 9 deletions(-)

diff --git a/ansible_mitogen/connection.py b/ansible_mitogen/connection.py
index 93ad0df6..2b031a9f 100644
--- a/ansible_mitogen/connection.py
+++ b/ansible_mitogen/connection.py
@@ -58,6 +58,10 @@ def _connect_local(spec):
         }
     }
 
+def wrap_or_none(klass, value):
+    if value is not None:
+        return klass(value)
+
 
 def _connect_ssh(spec):
     if C.HOST_KEY_CHECKING:
@@ -71,7 +75,7 @@ def _connect_ssh(spec):
             'check_host_keys': check_host_keys,
             'hostname': spec['remote_addr'],
             'username': spec['remote_user'],
-            'password': spec['password'],
+            'password': wrap_or_none(mitogen.core.Secret, spec['password']),
             'port': spec['port'],
             'python_path': spec['python_path'],
             'identity_file': spec['private_key_file'],
@@ -142,7 +146,7 @@ def _connect_su(spec):
         'enable_lru': True,
         'kwargs': {
             'username': spec['become_user'],
-            'password': spec['become_pass'],
+            'password': wrap_or_none(mitogen.core.Secret, spec['become_pass']),
             'python_path': spec['python_path'],
             'su_path': spec['become_exe'],
             'connect_timeout': spec['timeout'],
@@ -156,7 +160,7 @@ def _connect_sudo(spec):
         'enable_lru': True,
         'kwargs': {
             'username': spec['become_user'],
-            'password': spec['become_pass'],
+            'password': wrap_or_none(mitogen.core.Secret, spec['become_pass']),
             'python_path': spec['python_path'],
             'sudo_path': spec['become_exe'],
             'connect_timeout': spec['timeout'],
@@ -171,7 +175,7 @@ def _connect_mitogen_su(spec):
         'method': 'su',
         'kwargs': {
             'username': spec['remote_user'],
-            'password': spec['password'],
+            'password': wrap_or_none(mitogen.core.Secret, spec['password']),
             'python_path': spec['python_path'],
             'su_path': spec['become_exe'],
             'connect_timeout': spec['timeout'],
@@ -185,7 +189,7 @@ def _connect_mitogen_sudo(spec):
         'method': 'sudo',
         'kwargs': {
             'username': spec['remote_user'],
-            'password': spec['password'],
+            'password': wrap_or_none(mitogen.core.Secret, spec['password']),
             'python_path': spec['python_path'],
             'sudo_path': spec['become_exe'],
             'connect_timeout': spec['timeout'],
@@ -581,7 +585,7 @@ class Connection(ansible.plugins.connection.ConnectionBase):
         """
         self.call(ansible_mitogen.target.write_path,
                   mitogen.utils.cast(out_path),
-                  mitogen.utils.cast(data),
+                  mitogen.core.Blob(data),
                   mode=mode,
                   utimes=utimes)
 
diff --git a/ansible_mitogen/services.py b/ansible_mitogen/services.py
index 96661c4c..98432fe7 100644
--- a/ansible_mitogen/services.py
+++ b/ansible_mitogen/services.py
@@ -489,8 +489,11 @@ class FileService(mitogen.service.Service):
     # odd-sized messages waste one tiny write() per message on the trailer.
     # Therefore subtract 10 bytes pickle overhead + 24 bytes header.
     IO_SIZE = mitogen.core.CHUNK_SIZE - (mitogen.core.Stream.HEADER_LEN + (
-        len(mitogen.core.Message.pickled(' ' * mitogen.core.CHUNK_SIZE).data) -
-        mitogen.core.CHUNK_SIZE
+        len(
+            mitogen.core.Message.pickled(
+                mitogen.core.Blob(' ' * mitogen.core.CHUNK_SIZE)
+            ).data
+        ) - mitogen.core.CHUNK_SIZE
     ))
 
     def _schedule_pending_unlocked(self, state):
@@ -507,7 +510,7 @@ class FileService(mitogen.service.Service):
             s = fp.read(self.IO_SIZE)
             if s:
                 state.unacked += len(s)
-                sender.send(s)
+                sender.send(mitogen.core.Blob(s))
             else:
                 # File is done. Cause the target's receive loop to exit by
                 # closing the sender, close the file, and remove the job entry.