Terraform version manager

metaform.py

#!/usr/bin/env python3

import os
import stat
import sys
import urllib.request
import zipfile
from html.parser import HTMLParser


METAFORM_DIR = '.metaform'
METAFORM_SUPPORTED_PLATFORMS = ['linux_amd64']
TF_VERSION_PIN_FILE = '.terraform-version'
TF_RELEASES_CACHE = 'releases'
TF_VERSIONS_CACHE = 'versions'
TF_RELEASES_URL = 'https://releases.hashicorp.com/terraform'


class TFReleaseParser(HTMLParser):
    def __init__(self):
        HTMLParser.__init__(self)
        self.data = {}

    def handle_starttag(self, tag, attrs):
        try:
            if attrs[0][0] == 'href':
                tf_rel = attrs[0][1]
                (_, tf_string, ver_string, _) = tf_rel.split('/')
                if tf_string == 'terraform':
                    for platform in METAFORM_SUPPORTED_PLATFORMS:
                        self.data[ver_string] = '{0}{1}terraform_{2}_{3}.zip'.\
                            format('/'.join(TF_RELEASES_URL.split('/')[:-1]),
                                   tf_rel,
                                   ver_string,
                                   platform)
        except (IndexError, ValueError):
            return


def parse_tf_releases():
    cached_vers = {}
    releases_cache = os.path.join(
        os.path.expanduser('~'), METAFORM_DIR, TF_RELEASES_CACHE)
    if os.path.isfile(releases_cache):
        with open(releases_cache, 'r') as f:
            for l in f:
                (ver, url) = l[:-1].split(',')
                cached_vers[ver] = url
    else:
        parser = TFReleaseParser()
        with urllib.request.urlopen(TF_RELEASES_URL) as url:
            parser.feed(url.read().decode())
        with open(releases_cache, 'w+') as f:
            for k, v in parser.data.items():
                f.write('{0},{1}\n'.format(k, v))
        cached_vers = parser.data
    return cached_vers


def store_tf_version(path, ver):
    version_cache = os.path.join(
        os.path.expanduser('~'), METAFORM_DIR, TF_VERSIONS_CACHE)
    with open(version_cache, 'a+') as f:
        f.write('{0},{1}\n'.format(path, ver))


def get_tf_versions():
    cached_dirs = {}
    version_cache = os.path.join(
        os.path.expanduser('~'), METAFORM_DIR, TF_VERSIONS_CACHE)
    if os.path.isfile(version_cache):
        with open(version_cache, 'r') as f:
            for l in f:
                (path, ver) = l[:-1].split(',')
                cached_dirs[path] = ver
    return cached_dirs


if __name__ == '__main__':
    metaform_dir = os.path.join(os.path.expanduser('~'), METAFORM_DIR)
    if not os.path.isdir(metaform_dir):
        os.makedirs(metaform_dir)

    metaform_bin_dir = os.path.join(metaform_dir, 'bin')
    if not os.path.isdir(metaform_bin_dir):
        os.makedirs(metaform_bin_dir)

    versioned_links = parse_tf_releases()

    cwd = os.getcwd()

    if not any(e.endswith('tf') for e in os.listdir(cwd)):
        print('No .tf files in cwd')
        sys.exit(1)

    tf_versions = get_tf_versions()

    pinned_version = None
    pinned_file = os.path.join(cwd, TF_VERSION_PIN_FILE)
    if os.path.isfile(pinned_file):
        with open(pinned_file, 'r') as f:
            pinned_version = f.readline().rstrip('\n')

    try:
        desired_version = tf_versions[cwd]
    except KeyError:
        desired_version = input('Which tf version does this state use?: ')
        if not desired_version:
            print('Please input desired tf version')
            sys.exit(1)
        else:
            yes = set(['yes', 'y', 'ye', ''])
            no = set(['no', 'n'])
            choice = input('Save to ./{0}?: '.format(
                TF_VERSION_PIN_FILE)).lower()
            if choice in yes:
                with open(pinned_file, 'w+') as f:
                    f.write('{0}\n'.format(desired_version))
            store_tf_version(cwd, desired_version)

    if pinned_version and pinned_version != desired_version:
        print(('Mismatch between .terraform-version and saved version:\n'
               '.terraform-version: {0}\tstored: {1}\n'
               'Please edit ~/.metaform/versions file and manually fix it'.format(
                   pinned_version, desired_version)))
        sys.exit(1)

    bin_location = os.path.join(metaform_bin_dir, desired_version)
    bin_name = os.path.join(metaform_bin_dir, desired_version, 'terraform')

    if not os.path.isfile(bin_name):
        zip_name = 'terraform_{0}.zip'.format(desired_version)
        zip_full = '{0}/{1}'.format(metaform_bin_dir, zip_name)
        print('Downloading terraform {0}...'.format(desired_version))
        urllib.request.urlretrieve(versioned_links[desired_version], zip_full)
        with zipfile.ZipFile(zip_full, 'r') as zip_ref:
            zip_ref.extractall(bin_location)
        os.remove(zip_full)
        st = os.stat(bin_name)
        os.chmod(bin_name, st.st_mode | stat.S_IEXEC)

    print('Using terraform {0}'.format(desired_version))
    os.execv(bin_name, sys.argv)