Merge pull request #198 from mkinney/add-mech-connector

Add mech connector

README.md

@@ -17,7 +17,7 @@ pyinfra automates/provisions/manages/deploys infrastructure super fast at massiv
 + 💻 **Agentless execution** by speaking native SSH/Docker/subprocess depending on the target.
 + ❗️ **Two stage process** that enables `--dry` runs before making any changes.
 + 📦 **Extendable** with _any_ Python package as configured & written in standard Python.
-+ 🔌 **Integrated** with Docker, Vagrant & Ansible out of the box.
++ 🔌 **Integrated** with Docker, Vagrant/Mech & Ansible out of the box.
 
 When you run pyinfra you'll see something like ([non animated version](https://raw.githubusercontent.com/Fizzadar/pyinfra/master/docs/static/example_deploy.png)):
 

docs/apidoc/pyinfra.api.connectors.mech.rst

@@ -0,0 +1,7 @@
+pyinfra.api.connectors.mech module
+=====================================
+
+.. automodule:: pyinfra.api.connectors.mech
+   :members:
+   :undoc-members:
+   :show-inheritance:

docs/cli.md

@@ -97,3 +97,12 @@ Now that nginx is installed on the box, we can use pyinfra to control the ``ngin
 ```sh
 pyinfra inventory.py init.service nginx running=true enabled=true
 ```
+
+#### Additional debug info
+
+For additional debug info, use one of these options:
+
++ `--debug` Print debug info.
++ `--debug-data` Print host/group data before connecting and exit.
++ `--debug-facts` Print facts after generating operations and exit.
++ `--debug-operations` Print operations after generating and exit.

pyinfra/api/connectors/__init__.py

@@ -1,6 +1,6 @@
 import six
 
-from . import ansible, docker, local, ssh, vagrant
+from . import ansible, docker, local, mech, ssh, vagrant
 
 
 # Connectors that handle execution of pyinfra operations
@@ -13,6 +13,7 @@
 # Connectors that handle generation of inventories
 INVENTORY_CONNECTORS = {  # pragma: no cover
     'docker': docker,
+    'mech': mech,
     'vagrant': vagrant,
     'ansible': ansible,
 }

pyinfra/api/connectors/mech.py

@@ -0,0 +1,174 @@
+import json
+
+from os import path
+from threading import Thread
+
+from six.moves.queue import Queue
+
+from pyinfra import local, logger
+from pyinfra.api.exceptions import InventoryError
+from pyinfra.api.util import memoize
+from pyinfra.progress import progress_spinner
+
+
+def _get_mech_ssh_config(queue, progress, target):
+    logger.debug('Loading SSH config for {0}'.format(target))
+
+    # Note: We have to work-around the fact that "mech ssh-config somehost"
+    # does not return the correct "Host" value. When "mech" fixes this
+    # issue we can simply this code.
+    lines = local.shell(
+        'mech ssh-config {0}'.format(target),
+        splitlines=True,
+    )
+
+    newlines = []
+    for line in lines:
+        if line.startswith('Host '):
+            newlines.append('Host ' + target)
+        else:
+            newlines.append(line)
+
+    queue.put(newlines)
+
+    progress(target)
+
+
+@memoize
+def get_mech_config(limit=None):
+    logger.info('Getting Mech config...')
+
+    if limit and not isinstance(limit, (list, tuple)):
+        limit = [limit]
+
+    # Note: There is no "--machine-readable" option to 'mech status'
+    with progress_spinner({'mech ls'}) as progress:
+        output = local.shell(
+            'mech ls',
+            splitlines=True,
+        )
+        progress('mech ls')
+
+    targets = []
+
+    for line in output:
+
+        address = ''
+
+        data = line.split()
+        target = data[0]
+
+        if len(data) == 5:
+            address = data[1]
+
+        # Skip anything not in the limit
+        if limit is not None and target not in limit:
+            continue
+
+        # For each vm that has an address, fetch it's SSH config in a thread
+        if address != '' and address[0].isdigit():
+            targets.append(target)
+
+    threads = []
+    config_queue = Queue()
+
+    with progress_spinner(targets) as progress:
+        for target in targets:
+            thread = Thread(
+                target=_get_mech_ssh_config,
+                args=(config_queue, progress, target),
+            )
+            threads.append(thread)
+            thread.start()
+
+    for thread in threads:
+        thread.join()
+
+    queue_items = list(config_queue.queue)
+
+    lines = []
+    for output in queue_items:
+        lines.extend(output)
+
+    return lines
+
+
+@memoize
+def get_mech_options():
+    if path.exists('@mech.json'):
+        with open('@mech.json', 'r') as f:
+            return json.loads(f.read())
+    return {}
+
+
+def _make_name_data(host):
+    mech_options = get_mech_options()
+    mech_host = host['Host']
+
+    data = {
+        'ssh_hostname': host['HostName'],
+    }
+
+    for config_key, data_key in (
+        ('Port', 'ssh_port'),
+        ('User', 'ssh_user'),
+        ('IdentityFile', 'ssh_key'),
+    ):
+        if config_key in host:
+            data[data_key] = host[config_key]
+
+    # Update any configured JSON data
+    if mech_host in mech_options.get('data', {}):
+        data.update(mech_options['data'][mech_host])
+
+    # Work out groups
+    groups = mech_options.get('groups', {}).get(mech_host, [])
+
+    if '@mech' not in groups:
+        groups.append('@mech')
+
+    return '@mech/{0}'.format(host['Host']), data, groups
+
+
+def make_names_data(limit=None):
+    mech_ssh_info = get_mech_config(limit)
+
+    logger.debug('Got Mech SSH info: \n{0}'.format(mech_ssh_info))
+
+    hosts = []
+    current_host = None
+
+    for line in mech_ssh_info:
+        if not line:
+            if current_host:
+                hosts.append(_make_name_data(current_host))
+
+            current_host = None
+            continue
+
+        key, value = line.strip().split(' ', 1)
+
+        if key == 'Host':
+            if current_host:
+                hosts.append(_make_name_data(current_host))
+
+            # Set the new host
+            current_host = {
+                key: value,
+            }
+
+        elif current_host:
+            current_host[key] = value
+
+        else:
+            logger.debug('Extra Mech SSH key/value ({0}={1})'.format(
+                key, value,
+            ))
+
+    if current_host:
+        hosts.append(_make_name_data(current_host))
+
+    if not hosts:
+        raise InventoryError('No running Mech instances found!')
+
+    return hosts

tests/test_connector_mech.py

@@ -0,0 +1,127 @@
+import json
+
+from unittest import TestCase
+
+from mock import mock_open, patch
+
+from pyinfra.api.connectors.mech import get_mech_options, make_names_data
+from pyinfra.api.exceptions import InventoryError
+
+FAKE_MECH_OPTIONS = {
+    'groups': {
+        'ubuntu16': ['mygroup'],
+    },
+    'data': {
+        'centos7': {
+            'somedata': 'somevalue',
+        },
+    },
+}
+FAKE_MECH_OPTIONS_DATA = json.dumps(FAKE_MECH_OPTIONS)
+
+
+def fake_mech_shell(command, splitlines=None):
+    if command == 'mech ls':
+        return [
+             'NAME	        ADDRESS	                                BOX	     VERSION	PATH',  # noqa: E501
+             'ubuntu16	       192.168.2.226             bento/ubuntu-16.04	 201912.04.0	/Users/bob/somedir/ubuntu16',  # noqa: E501
+             'centos7	       192.168.2.227                 bento/centos-7	 201912.05.0	/Users/bob/somedir/centos7',  # noqa: E501
+             'centos6	       poweroff                      bento/centos-6	 201912.04.0	/Users/bob/somedir/centos6',  # noqa: E501
+             'fedora31	               	                    bento/fedora-31	 201912.04.0	/Users/bob/somedir/fedora31',  # noqa: E501
+        ]
+    elif command == 'mech ssh-config ubuntu16':
+        return [
+            'Host ubuntu16',
+            '  User vagrant',
+            '  Port 22',
+            '  UserKnownHostsFile /dev/null',
+            '  StrictHostKeyChecking no',
+            '  PasswordAuthentication no',
+            '  IdentityFile path/to/key',
+            '  IdentitiesOnly yes',
+            '  LogLevel FATAL',
+            '  HostName 192.168.2.226',
+            '',
+        ]
+    elif command == 'mech ssh-config centos7':
+        return [
+            'Host centos7',
+            '  User vagrant',
+            '  Port 22',
+            '  UserKnownHostsFile /dev/null',
+            '  StrictHostKeyChecking no',
+            '  PasswordAuthentication no',
+            '  IdentityFile path/to/key',
+            '  IdentitiesOnly yes',
+            '  LogLevel FATAL',
+            '  HostName 192.168.2.227',
+            '',
+        ]
+    elif command == 'mech ssh-config centos6':
+        return [
+            'ERROR: Error: The virtual machine is not powered on: /Users/bob/debian8/.mech/debian-8.11-amd64.vmx',  # noqa: E501
+            'This Mech machine is reporting that it is not yet ready for SSH. Make',
+            'sure your machine is created and running and try again. Additionally,',
+            'check the output of `mech status` to verify that the machine is in the',
+            'state that you expect.',
+        ]
+
+    return []
+
+
+@patch('pyinfra.api.connectors.mech.local.shell', fake_mech_shell)
+class TestMechConnector(TestCase):
+    def tearDown(self):
+        get_mech_options.cache = {}
+
+    @patch(
+        'pyinfra.api.connectors.mech.open',
+        mock_open(read_data=FAKE_MECH_OPTIONS_DATA),
+        create=True,
+    )
+    @patch('pyinfra.api.connectors.mech.path.exists', lambda path: True)
+    def test_make_names_data_with_options(self):
+        data = make_names_data()
+
+        assert data == [
+            (
+                '@mech/ubuntu16',
+                {
+                    'ssh_port': '22',
+                    'ssh_user': 'vagrant',
+                    'ssh_hostname': '192.168.2.226',
+                    'ssh_key': 'path/to/key',
+                },
+                ['mygroup', '@mech'],
+            ), (
+                '@mech/centos7',
+                {
+                    'ssh_port': '22',
+                    'ssh_user': 'vagrant',
+                    'ssh_hostname': '192.168.2.227',
+                    'ssh_key': 'path/to/key',
+                    'somedata': 'somevalue',
+                },
+                ['@mech'],
+            ),
+        ]
+
+    def test_make_names_data_with_limit(self):
+        data = make_names_data(limit=('ubuntu16',))
+
+        assert data == [
+            (
+                '@mech/ubuntu16',
+                {
+                    'ssh_port': '22',
+                    'ssh_user': 'vagrant',
+                    'ssh_hostname': '192.168.2.226',
+                    'ssh_key': 'path/to/key',
+                },
+                ['@mech'],
+            ),
+        ]
+
+    def test_make_names_data_no_matches(self):
+        with self.assertRaises(InventoryError):
+            make_names_data(limit='nope')