From 6059e0c1dbffb6c234e1b03e1c77b23b11065110 Mon Sep 17 00:00:00 2001
From: David Wilson <dw@botanicus.net>
Date: Thu, 15 Feb 2018 00:27:23 +0545
Subject: [PATCH] ansible: support become:True and become_user

---
 ansible_mitogen/connection/mitogen.py | 17 +++++++++++++++--
 ansible_mitogen/strategy/mitogen.py   |  7 ++++---
 2 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/ansible_mitogen/connection/mitogen.py b/ansible_mitogen/connection/mitogen.py
index 24642ebe..61f92d71 100644
--- a/ansible_mitogen/connection/mitogen.py
+++ b/ansible_mitogen/connection/mitogen.py
@@ -54,7 +54,7 @@ class Connection(ansible.plugins.connection.ConnectionBase):
     router = None
     context = None
 
-    become_methods = []
+    become_methods = ['sudo']
     transport = 'mitogen'
 
     @property
@@ -66,10 +66,23 @@ class Connection(ansible.plugins.connection.ConnectionBase):
             return
 
         self.router, self.parent = mitogen.unix.connect('/tmp/mitosock')
-        self.context = mitogen.service.call(self.parent, 500, {
+        host = mitogen.service.call(self.parent, 500, {
+            'method': 'ssh',
             'hostname': self._play_context.remote_addr,
+            'username': self._play_context.remote_user,
+            'port': self._play_context.port,
+            'ssh_path': self._play_context.ssh_executable,
         })
 
+        if not self._play_context.become:
+            self.context = host
+        else:
+            self.context = mitogen.service.call(self.parent, 500, {
+                'method': 'sudo',
+                'username': self._play_context.become_user,
+                'via': host,
+            })
+
     def py_call(self, func, *args, **kwargs):
         self._connect()
         return self.context.call(func, *args, **kwargs)
diff --git a/ansible_mitogen/strategy/mitogen.py b/ansible_mitogen/strategy/mitogen.py
index 8adf146e..7be52f59 100644
--- a/ansible_mitogen/strategy/mitogen.py
+++ b/ansible_mitogen/strategy/mitogen.py
@@ -49,13 +49,14 @@ class ContextProxyService(mitogen.service.Service):
         self._context_by_id = {}
 
     def validate_args(self, args):
-        return (isinstance(args, dict) and
-                isinstance(args.get('hostname'), basestring))
+        return isinstance(args, dict)
 
     def dispatch(self, dct, msg):
+        print dct.get('via')
         key = repr(sorted(dct.items()))
         if key not in self._context_by_id:
-            self._context_by_id[key] = self.router.ssh(**dct)
+            method = getattr(self.router, dct.pop('method'))
+            self._context_by_id[key] = method(**dct)
         return self._context_by_id[key]