sync-do-hosts

#!/usr/bin/env python3

import os
import re
import click
import digitalocean

BEGIN = '\n# BEGIN DigitalOcean\n'
END = '# END\n'


@click.command()
@click.option('token_file', '--tokenFile', '-T', default='~/.dotokens')
@click.option('--token', '-t', multiple=True)
@click.option('--file', '-f', default='/etc/hosts')
def sync(token_file, token, file):
    path = os.path.expanduser(token_file)

    if token:
        tokens = token
    elif os.path.isfile(path):
        with open(path, 'r') as f:
            tokens = f.read().splitlines()
    else:
        raise click.UsageError('no token file or tokens supplied')

    click.echo('Retrieving hosts from DigitalOcean', nl=False)
    hosts = {}
    for t in tokens:
        if (t.startswith('#')):
            continue

        t = t.split(' ')[0]
        manager = digitalocean.Manager(token=t)
        droplets = manager.get_all_droplets()
        hosts.update({d.ip_address: d.name for d in droplets})
        click.echo('.', nl=False)
    click.echo('\nWriting {} hosts to {}'.format(len(hosts), file))

    with open(file, 'r') as f:
        contents = f.read()

        if all(s in contents for s in [BEGIN, END]):
            contents = re.sub(r'{0}.+{1}'.format(BEGIN, END), '', contents, flags=re.DOTALL)

    with open(file, 'w+') as f:
        f.write(contents)
        f.write(BEGIN)
        f.writelines(['{ip}\t{host}\n'.format(ip=ip, host=host) for (ip, host) in hosts.items()])
        f.write(END)


if __name__ == '__main__':
    sync()