Fireball2 mode working!

lib/ansible/runner/connection_plugins/fireball2.py

@@ -19,7 +19,9 @@ import json
 import os
 import base64
 import socket
+import struct
 from ansible.callbacks import vvv
+from ansible.runner.connection_plugins.ssh import Connection as SSHConnection
 from ansible import utils
 from ansible import errors
 from ansible import constants
@@ -27,32 +29,68 @@ from ansible import constants
 class Connection(object):
     ''' raw socket accelerated connection '''
 
-    def __init__(self, runner, host, port, *args, **kwargs):
+    def __init__(self, runner, host, port, user, password, private_key_file, *args, **kwargs):
+
+        self.ssh = SSHConnection(
+            runner=runner,
+            host=host, 
+            port=port, 
+            user=user, 
+            password=password, 
+            private_key_file=private_key_file
+        )
 
         self.runner = runner
+        self.host = host
+        self.context = None
+        self.conn = None
+        self.key = utils.key_for_hostname(host)
+        self.fbport = constants.FIREBALL2_PORT
+        self.is_connected = False
 
         # attempt to work around shared-memory funness
         if getattr(self.runner, 'aes_keys', None):
             utils.AES_KEYS = self.runner.aes_keys
 
-        self.host = host
-        self.context = None
-        self.conn = None
-        self.cipher = AES256Cipher()
+    def _execute_fb_module(self):
+        args = "password=%s" % base64.b64encode(self.key.__str__())
+        self.ssh.connect()
+        return self.runner._execute_module(self.ssh, "/root/.ansible/tmp", 'fireball2', args, inject={"password":self.key})
 
-        if  port is None:
-            self.port = constants.FIREBALL2_PORT
-        else:
-            self.port = port
-
-    def connect(self):
+    def connect(self, allow_ssh=True):
         ''' activates the connection object '''
 
-        self.conn = socket.socket()
-        self.conn.connect((self.host,self.port))
+        if self.is_connected:
+            return self
 
+        try:
+            self.conn = socket.socket()
+            self.conn.connect((self.host,self.fbport))
+        except:
+            if allow_ssh:
+                print "Falling back to ssh to startup accelerated mode"
+                res = self._execute_fb_module()
+                return self.connect(allow_ssh=False)
+            else:
+                raise errors.AnsibleError("Failed to connect to %s:%s" % (self.host,self.fbport))
+        self.is_connected = True
         return self
 
+    def send_data(self, data):
+        packed_len = struct.pack('Q',len(data))
+        return self.conn.sendall(packed_len + data)
+
+    def recv_data(self):
+        header_len = 8 # size of a packed unsigned long long
+        data = b""
+        while len(data) < header_len:
+            data += self.conn.recv(1024)
+        data_len = struct.unpack('Q',data[:header_len])[0]
+        data = data[header_len:]
+        while len(data) < data_len:
+            data += self.conn.recv(1024)
+        return data
+
     def exec_command(self, cmd, tmp_path, sudo_user, sudoable=False, executable='/bin/sh'):
         ''' run a command on the remote host '''
 
@@ -65,12 +103,12 @@ class Connection(object):
             executable=executable,
         )
         data = utils.jsonify(data)
-        data = self.cipher.encrypt(data)
-        if self.conn.sendall(data):
+        data = utils.encrypt(self.key, data)
+        if self.send_data(data):
             raise errors.AnisbleError("Failed to send command to %s:%s" % (self.host,self.port))
         
-        response = self.conn.recv(2048)
-        response = self.cipher.decrypt(response)
+        response = self.recv_data()
+        response = utils.decrypt(self.key, response)
         response = utils.parse_json(response)
 
         return (response.get('rc',None), '', response.get('stdout',''), response.get('stderr',''))
@@ -83,18 +121,18 @@ class Connection(object):
         if not os.path.exists(in_path):
             raise errors.AnsibleFileNotFound("file or module does not exist: %s" % in_path)
 
-        data = base64.file(in_path).read()
+        data = file(in_path).read()
         data = base64.b64encode(data)
         data = dict(mode='put', data=data, out_path=out_path)
 
         # TODO: support chunked file transfer
         data = utils.jsonify(data)
-        data = self.cipher.encrypt(data)
-        if self.conn.sendall(data):
+        data = utils.encrypt(self.key, data)
+        if self.send_data(data):
             raise errors.AnsibleError("failed to send the file to %s:%s" % (self.host,self.port))
 
-        response = self.conn.recv(2048)
-        response = self.cipher.decrypt(response)
+        response = self.recv_data()
+        response = utils.decrypt(self.key, data)
         response = utils.parse_json(response)
 
         # no meaningful response needed for this
@@ -105,12 +143,12 @@ class Connection(object):
 
         data = dict(mode='fetch', in_path=in_path)
         data = utils.jsonify(data)
-        data = self.cipher.encrypt(data)
-        if self.conn.sendall(data):
+        data = utils.encrypt(self.key, data)
+        if self.send_data(data):
             raise errors.AnsibleError("failed to initiate the file fetch with %s:%s" % (self.host,self.port))
 
-        response = self.socket.recv(2048)
-        response = self.cipher.decrypt(response)
+        response = self.recv_data()
+        response = utils.decrypt(self.key, data)
         response = utils.parse_json(response)
         response = response['data']
         response = base64.b64decode(response)