command 12.2 KB
Newer Older
Noe Nieto's avatar
Noe Nieto committed
1 2
#!/usr/bin/python3
import sys
Noe Nieto's avatar
Noe Nieto committed
3
import os, stat
Noe Nieto's avatar
Noe Nieto committed
4 5 6 7
import argparse
import tempfile
import subprocess
import json
Noe Nieto's avatar
Noe Nieto committed
8
import shutil
9
from pathlib import Path
10 11
import ipaddress

Noe Nieto's avatar
Noe Nieto committed
12

Noe Nieto's avatar
Noe Nieto committed
13 14 15 16 17 18 19
ansible_shim = """
---
- hosts: all
  become: yes
  tasks:
    - name: Shim file for ansible
      file:
Noe Nieto's avatar
Noe Nieto committed
20
        path: /etc/ansible_shim.txt
Noe Nieto's avatar
Noe Nieto committed
21 22 23
        state: touch
"""

Noe Nieto's avatar
Noe Nieto committed
24 25

XDG_BOX_CONFIG_HOME = Path(os.environ['HOME'],'.config','ldh_developer','box')
26
XDG_PB_CONFIG_HOME = Path(os.environ['HOME'],'.config','ldh_developer','playbooks')
Noe Nieto's avatar
Noe Nieto committed
27 28 29
VG_REQ_PLUGINS = {'libvirt': 'vagrant-libvirt', 'digital_ocean': 'vagrant-digitalocean'}
HERE = Path(__file__).parent
DO_SSH_KEY_PATH = Path('~/.ssh/shipwright/').expanduser()
30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58
VAGRANT_CONFIG_HOME = Path(os.environ.get('VAGRANT_HOME', str(Path.home().joinpath('.vagrant.d'))))
VAGRANT_DB_PATH = Path(VAGRANT_CONFIG_HOME, 'data', 'machine-index', 'index')


class VagrantDB(object):
    @property
    def boxes(self):
        data = json.loads(VAGRANT_DB_PATH.open().read())
        vgt_boxes = data['machines']
        return {
            vgt_boxes[box_id]['name']:{
                'uid': box_id,
                'provider': vgt_boxes[box_id]['provider'],
                'state': vgt_boxes[box_id]['state'],
                'vagrantfile_path': vgt_boxes[box_id]['vagrantfile_path'],
            } for box_id in vgt_boxes
        }

    def provider_for(self, hostname):
        try:
            return self.boxes[hostname]['provider']
        except KeyError:
            return None

    @classmethod
    def prune(cls):
        print("Removing vagrant's invalid entries ...")
        subprocess.run(['vagrant', 'global-status', '--prune'], capture_output=True)
        print("Done")
Noe Nieto's avatar
Noe Nieto committed
59 60 61 62 63 64 65 66 67


def ensure_plugins():
    cp = subprocess.run(
        ['vagrant', 'plugin', 'list'],
        stdout=subprocess.PIPE, stderr=subprocess.PIPE,
        universal_newlines=True
    )
    if cp.returncode != 0:
68
        print("Ooops, there's something wrong with vagrant!!?")
Noe Nieto's avatar
Noe Nieto committed
69 70 71 72 73 74 75 76 77
        print(cp.stdout)
        print(cp.stderr)
        sys.exit(1)

    if set(VG_REQ_PLUGINS.values()) != set([l.split()[0] for l in cp.stdout.splitlines()]):
        print('There are some missing vagrant plugins.')
        print('Install them with: apt install ' + ' '.join(VG_REQ_PLUGINS))
        print('Or run the bootstrap script')
        sys.exit(1)
Noe Nieto's avatar
Noe Nieto committed
78

79

80 81 82 83 84 85 86 87 88 89 90 91
def forbdden_for_digital_ocean(func):
    def wrapper(*args, **kwargs):
        vg_db = VagrantDB()
        hostname = args[0]
        if vg_db.provider_for(hostname) == 'digital_ocean':
            print("Error: This command is forbidden for the digital_ocean provider.")
            exit(1)
        else:
            func(*args, **kwargs)
    return wrapper


Noe Nieto's avatar
Noe Nieto committed
92 93 94
def only_if_box_exist(func):
    def wrapper(*args, **kwargs):
        hostname = args[0]
Noe Nieto's avatar
Noe Nieto committed
95
        vagrant_wd = XDG_BOX_CONFIG_HOME.joinpath(hostname)
Noe Nieto's avatar
Noe Nieto committed
96 97 98
        if vagrant_wd.exists():
            func(*args, **kwargs)
        else:
99
            print("Error: no box managed by Shipwright has the hostname '{}'".format(hostname))
Noe Nieto's avatar
Noe Nieto committed
100 101 102
            exit(1)
    return wrapper

Noe Nieto's avatar
Noe Nieto committed
103

104 105 106 107 108 109 110
def prune_after_call(func):
    def wrapper(*args, **kwargs):
        func(*args, **kwargs)
        VagrantDB.prune()
    return wrapper


111 112 113 114 115 116 117 118
def is_ipaddress(ip_addr):
    try:
        ipaddress.ip_address(ip_addr)
    except ValueError:
        return False
    return True


119 120
def print_status():
    """
121
    Prints the status of all the boxes managed by Shipwright
122
    """
123
    box_paths = [str(d) for d in XDG_BOX_CONFIG_HOME.iterdir() if d.is_dir()]
124
    count = 0
125 126 127 128 129 130
    if VAGRANT_CONFIG_HOME.exists():
        print('-----------------------------------------------------------------')
        print('| hostname       | State     | Provider       | IP Address      |')
        print('-----------------------------------------------------------------')
        db = VagrantDB()
        for hostname, v in db.boxes.items():
131
            vagrant_wd = v['vagrantfile_path']
132
            if vagrant_wd in box_paths:
133 134 135 136 137 138 139 140 141 142 143 144 145
                # get the ip address
                cp = subprocess.run(
                    ['vagrant','ssh-config'], cwd=vagrant_wd,
                    stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                    universal_newlines=True
                )
                if cp.returncode != 0:
                    ip_addr = "<N/A>"
                else:
                    for fld in cp.stdout.split():
                        if is_ipaddress(str(fld)):
                            ip_addr = str(fld)
                            break
146
                print(f"| {hostname:15}| {v['state']:10}| {v['provider']:15}| {ip_addr:15} |")
147
                count+=1
148 149
    print('-----------------------------------------------------------------')
    print(f'Total {count} box(es)')
150 151


152
@prune_after_call
153
def box_create(hostname, ram, cpus, provider, token):
154 155 156
    """
    This creates a box using the template
    """
Noe Nieto's avatar
Noe Nieto committed
157
    vagrant_wd = XDG_BOX_CONFIG_HOME.joinpath(hostname)
158
    if vagrant_wd.exists():
159
        print("Shipyard error: There's already a box with that hostname")
160 161
        exit(0)

Noe Nieto's avatar
Noe Nieto committed
162 163
    vagrant_wd.mkdir(parents=True)
    os.chmod(vagrant_wd, stat.S_IRWXU)
164
    vagrant_wd.joinpath('shim.yml').open(mode='w').write(ansible_shim)
165

Noe Nieto's avatar
Noe Nieto committed
166 167 168 169 170
    if provider == 'digital_ocean':
        if not DO_SSH_KEY_PATH.joinpath('id_rsa').exists():
            print("Creating ssh key required by Vagrant's Digital Ocean plugin ...")
            DO_SSH_KEY_PATH.mkdir(parents=True, exist_ok=True)
            os.chmod(DO_SSH_KEY_PATH, stat.S_IRWXU)
171
            subprocess.run(['ssh-keygen', '-f', 'id_rsa', '-t', 'rsa', '-C', 'Created automatically by Shipwright', '-N', ''], cwd=DO_SSH_KEY_PATH)
Noe Nieto's avatar
Noe Nieto committed
172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188

        vagrant_template = open(HERE.joinpath('digital_ocean.tpl'), 'r').read()
        with vagrant_wd.joinpath('Vagrantfile').open(mode='w') as vfile:
            vfile.write(vagrant_template.format(
                hostname=hostname,
                do_token=token,
                username=os.environ['USER']
            ))
    else:
        vagrant_template = open(HERE.joinpath('libvirt.tpl'), 'r').read()
        with vagrant_wd.joinpath('Vagrantfile').open(mode='w') as vfile:
            vfile.write(vagrant_template.format(
                hostname=hostname,
                ram=ram,
                cpus=cpus
            ))

189 190 191 192
    # Execute vagrant with the generated vagrantfile
    subprocess.run(['vagrant','up'], cwd=vagrant_wd)


193 194
@prune_after_call
@forbdden_for_digital_ocean
Noe Nieto's avatar
Noe Nieto committed
195
@only_if_box_exist
196
def box_destroy(hostname):
197 198 199
    """
    Destroy a box, or all
    """
Noe Nieto's avatar
Noe Nieto committed
200
    vagrant_wd = XDG_BOX_CONFIG_HOME.joinpath(hostname)
201 202 203
    r = subprocess.run(['vagrant','destroy'], cwd=vagrant_wd)
    if r.returncode == 0:
        shutil.rmtree(vagrant_wd)
Noe Nieto's avatar
Noe Nieto committed
204 205


206 207
@prune_after_call
@forbdden_for_digital_ocean
Noe Nieto's avatar
Noe Nieto committed
208
@only_if_box_exist
Noe Nieto's avatar
Noe Nieto committed
209
def run_playbook(hostname, playbook, retry_file=None):
Noe Nieto's avatar
Noe Nieto committed
210
    """
211
    run an Ansible playbook against this box
Noe Nieto's avatar
Noe Nieto committed
212
    """
Noe Nieto's avatar
Noe Nieto committed
213
    vagrant_wd = XDG_BOX_CONFIG_HOME.joinpath(hostname)
Noe Nieto's avatar
Noe Nieto committed
214
    inventory_file = vagrant_wd.joinpath('.vagrant', 'provisioners', 'ansible', 'inventory', 'vagrant_ansible_inventory')
Noe Nieto's avatar
Noe Nieto committed
215
    _args = [
Noe Nieto's avatar
Noe Nieto committed
216
        'ansible-playbook', '-u', 'vagrant', '-i', f'{inventory_file}',
217
        '-e', f'shipwright_box_config_home={vagrant_wd}',
218
        f'{XDG_PB_CONFIG_HOME.joinpath(playbook).resolve()}'
Noe Nieto's avatar
Noe Nieto committed
219 220 221
    ]

    if retry_file:
222
        _args.extend(['--limit', f'@{XDG_PB_CONFIG_HOME.joinpath(retry_file).resolve()}'])
Noe Nieto's avatar
Noe Nieto committed
223 224

    subprocess.run(_args, cwd=vagrant_wd)
Noe Nieto's avatar
Noe Nieto committed
225

226

227
@forbdden_for_digital_ocean
Noe Nieto's avatar
Noe Nieto committed
228
@only_if_box_exist
229
def box_start(hostname):
Noe Nieto's avatar
Noe Nieto committed
230
    """
231
    Starts the box with the provided hostname
Noe Nieto's avatar
Noe Nieto committed
232
    """
Noe Nieto's avatar
Noe Nieto committed
233
    vagrant_wd = XDG_BOX_CONFIG_HOME.joinpath(hostname)
Noe Nieto's avatar
Noe Nieto committed
234 235 236
    subprocess.run(['vagrant','up'], cwd=vagrant_wd)


237 238
@prune_after_call
@forbdden_for_digital_ocean
Noe Nieto's avatar
Noe Nieto committed
239
@only_if_box_exist
240
def box_halt(hostname):
Noe Nieto's avatar
Noe Nieto committed
241
    """
242
    Stops the box with the provided hostname
Noe Nieto's avatar
Noe Nieto committed
243
    """
Noe Nieto's avatar
Noe Nieto committed
244
    vagrant_wd = XDG_BOX_CONFIG_HOME.joinpath(hostname)
Noe Nieto's avatar
Noe Nieto committed
245 246 247
    subprocess.run(['vagrant','halt'], cwd=vagrant_wd)


Noe Nieto's avatar
Noe Nieto committed
248
@only_if_box_exist
249
def box_ssh_connect(hostname):
250
    """
251
    Open a SSH session to the hostname
252
    """
Noe Nieto's avatar
Noe Nieto committed
253
    vagrant_wd = XDG_BOX_CONFIG_HOME.joinpath(hostname)
Noe Nieto's avatar
Noe Nieto committed
254
    subprocess.run(['vagrant','ssh'], cwd=vagrant_wd)
255 256


Noe Nieto's avatar
Noe Nieto committed
257
@only_if_box_exist
258
def box_get_working_directory(hostname):
Noe Nieto's avatar
Noe Nieto committed
259
    """
260
    Print the directory of a box by it's hostname
Noe Nieto's avatar
Noe Nieto committed
261
    """
Noe Nieto's avatar
Noe Nieto committed
262
    vagrant_wd = XDG_BOX_CONFIG_HOME.joinpath(hostname)
Noe Nieto's avatar
Noe Nieto committed
263 264
    print(vagrant_wd)

265

Noe Nieto's avatar
Noe Nieto committed
266
@only_if_box_exist
267
def box_gio_open(hostname):
Noe Nieto's avatar
Noe Nieto committed
268
    """
269
    Open the box directory in nautilis or any other file manager
Noe Nieto's avatar
Noe Nieto committed
270
    """
Noe Nieto's avatar
Noe Nieto committed
271
    vagrant_wd = XDG_BOX_CONFIG_HOME.joinpath(hostname)
Noe Nieto's avatar
Noe Nieto committed
272 273
    subprocess.run(['gio','open', vagrant_wd])

274
COMMANDS = {
275 276
    'status': "Prints the status of all the boxes' managed by box",
    'create': 'Creates a new box using either libvirt(default) or Digital Ocean',
Noe Nieto's avatar
Noe Nieto committed
277
    'up': 'Start box',
278
    'start': 'Start box',
Noe Nieto's avatar
Noe Nieto committed
279 280 281 282
    'halt': 'Stop/shutdown a box',
    'restart': 'Restarts the box',
    'destroy': 'Destroy the box and do cleanup',
    'ssh': 'Connect to the box using SSH',
283
    'playbook': 'Run a playbook against a box',
284
    'which': 'Print the path to the Vagrantfile',
Noe Nieto's avatar
Noe Nieto committed
285
    'open': 'Show the working directory of the box using the file manager.'
286
}
Noe Nieto's avatar
Noe Nieto committed
287

Noe Nieto's avatar
Noe Nieto committed
288

Noe Nieto's avatar
Noe Nieto committed
289
if __name__ == '__main__':
290
    cmd_parser = argparse.ArgumentParser(
291 292
        prog='shipwright box',
        description="This utility helps you mange virtual boxes with Vagrant without dealing directly with configuration details",
Noe Nieto's avatar
Noe Nieto committed
293
    )
294 295
    cmd_parser.add_argument('command', choices=COMMANDS,)
    cmd_args = cmd_parser.parse_args(sys.argv[1:2])
Noe Nieto's avatar
Noe Nieto committed
296

Noe Nieto's avatar
Noe Nieto committed
297 298
    # Ensure XDG_BOX_CONFIG_HOME exists
    XDG_BOX_CONFIG_HOME.mkdir(parents=True, exist_ok=True)
Noe Nieto's avatar
Noe Nieto committed
299

Noe Nieto's avatar
Noe Nieto committed
300
    ensure_plugins()
301
    if cmd_args.command == 'status':
302
        VagrantDB.prune()
303
        print_status()
304 305
        sys.exit()

306
    sub_parser = argparse.ArgumentParser()
307
    if cmd_args.command == 'create':
308
        sub_parser.usage = 'box create <hostname> [--ram RAM] [--cpus CPUS] [--provider libvirt|digital_ocean] [--token]'
309
        sub_parser.add_argument('hostname', help='This is the hostname for the new the box')
Noe Nieto's avatar
Noe Nieto committed
310
        sub_parser.add_argument('--ram', help='Hoy much RAM for this box (default is 512MB for libvirt and 1GB for digital ocean.)', default=512)
311
        sub_parser.add_argument('--cpus', help='Hoy many CPUs for this box (default is 1)', default=1)
312
        sub_parser.add_argument('--provider', help='Tells Shipwright which vagrant provider to use. Default is libvirt', default='libvirt')
Noe Nieto's avatar
Noe Nieto committed
313
        sub_parser.add_argument('--token', help="This is your personal access token for Digital Ocean's API")
314
        sub_args = sub_parser.parse_args(sys.argv[2:])
Noe Nieto's avatar
Noe Nieto committed
315 316
        if sub_args.provider == 'digital_ocean' and sub_args.token is None:
            print('Error: When using the digital_ocean provider you must provide the access token\n')
317
            print('Example: shipwright box create foobar --provider digital_ocean --token 0d74f0a...')
Noe Nieto's avatar
Noe Nieto committed
318
            sys.exit(1)
319
        box_create(sub_args.hostname, sub_args.ram, sub_args.cpus, sub_args.provider, sub_args.token)
320 321 322
        sys.exit()

    if cmd_args.command == 'playbook':
323
        sub_parser.add_argument('hostname', help='This is the hostname of the box')
324
        sub_parser.usage = 'box playbook hostname path/to/playbook.yml'
Noe Nieto's avatar
Noe Nieto committed
325
        sub_parser.add_argument('playbook', help='The path to the playbook file')
Noe Nieto's avatar
Noe Nieto committed
326
        sub_parser.add_argument('--retry', help='The path to the playbook retry file.', required=False)
327
        sub_args = sub_parser.parse_args(sys.argv[2:])
Noe Nieto's avatar
Noe Nieto committed
328
        run_playbook(sub_args.hostname, sub_args.playbook, sub_args.retry)
329 330
        sys.exit()

331 332 333 334 335 336 337
    if cmd_args.command in ['up', 'start']:
        sub_parser.add_argument('hostnames', help='The hostname(s) of the boxes to start', nargs='+')
        sub_args = sub_parser.parse_args(sys.argv[2:])
        for hostname in sub_args.hostnames:
            box_start(hostname)

    sub_parser.add_argument('hostname', help='This is the hostname of the box')
338
    sub_args = sub_parser.parse_args(sys.argv[2:])
339
    if cmd_args.command in ['halt', 'stop']:
340
        box_halt(sub_args.hostname)
Noe Nieto's avatar
Noe Nieto committed
341
    elif cmd_args.command == 'restart':
342 343
        box_halt(sub_args.hostname)
        box_start(sub_args.hostname)
344
    elif cmd_args.command in ['destroy', 'rm']:
345
        box_destroy(sub_args.hostname)
346
    elif cmd_args.command == 'ssh':
347
        box_ssh_connect(sub_args.hostname)
348
    elif cmd_args.command == 'which':
349
        box_get_working_directory(sub_args.hostname)
350
    elif cmd_args.command == 'open':
351
        box_gio_open(sub_args.hostname)
352
    else:
353
        cmd_parser.print_help()
354
        exit(100)