tf.py: World's thinnest Terraform wrapper

tf.py

#!/usr/bin/env python3
#
# tf.py
# World's thinnest Terraform wrapper
# https://github.com/marns/tf.py
#
# Usage:
# 1. Create an alias to tf.py for your shell of choice (optional), e.g.:
# > echo "alias tf='./tf.py'" >> ~/.zshrc
# > source ~/.zshrc
#
# 2. Run Terraform commands as usual, subbing in `tf` and set your workspace:
# > ENV=dev tf plan
#
# 3. Add your own commands as needed.
#
# Apologies to TensorFlow
#
# Author:  Martin Smith
# License: MIT

import sys
import os
import subprocess
import platform
from pathlib import Path

class tf():

    # Config
    plans_dir = ".plans"
    terraform_version = "0.14.8"

    # Utility
    @staticmethod
    def print_cmd(args):
        display_cmd = args
        if isinstance(display_cmd, list):
            display_cmd = ' '.join(args)
        print('> ' + display_cmd)

    @staticmethod
    def run(args, shell=False):
        tf.print_cmd(args)
        return subprocess.run(args, shell=shell).returncode

    @staticmethod
    def run_capture_output(args):
        return subprocess.run(args, shell=True, capture_output=True, text=True).stdout.strip("\n")

    @staticmethod
    def check_platform():
        os = platform.system()
        if os == "Darwin":
            md5 = tf.run_capture_output("md5 -q `which terraform`")
            required_md5 = "f0c6df2c7ce9778a5cbaf3540cf45bfc"
        elif os == "Linux":
            # I don't actually know the Linux hash, if you get here you should set this value
            md5 = tf.run_capture_output("md5sum - < `which terraform` | tr -d ' -'")
            required_md5 = "f91834c3ac0f78c77914498765d3e0b7"
            return False
        else:
            print("System %s not supported. You can add it!" % os)
            return False
        
        if md5 != required_md5:
            print("Inconsistent Terraform version. Make sure you're running %s on a known OS" % tf.terraform_version)
            return False

        return True

    def check_env(self):
        env = os.environ.get('ENV')
        if not env: print("ENV must be set.")
        return env

    def select_workspace(self, env):
        # Check current workspace. Perhaps not as thorough, but `terraform workspace select` can be very slow.
        current_workspace = tf.run_capture_output("terraform workspace show")
        if current_workspace == env:
            print("Terraform workspace '%s' already selected" % env)
            return True
        
        return tf.run("terraform workspace select %s" % env, shell=True)

    # Terraform command handlers
    def plan(self, env, args):
        Path(self.plans_dir).mkdir(parents=True, exist_ok=True)

        args = [ "terraform", "plan", "-var-file=env/%s.tfvars" % env, "-out", "%s/%s.plan" % (self.plans_dir, env), *args ]
        return tf.run(args)

    def apply(self, env, args):
        plan_filename = "%s/%s.plan" % (self.plans_dir, env)
        args = [ "terraform", "apply", plan_filename, *args ]
        
        return_code = tf.run(args)
        if return_code: return return_code
        
        os.remove(plan_filename)
        print('Removed ' + plan_filename)

    # tf.py
    def do_cmd(self, env):
        l = len(sys.argv)
        if l < 2: return False

        # Get function name to intercept from the first parameter.
        # If it's a member of this class, hand off to it.
        cmd_name = sys.argv[1]
        cmd_func = getattr(self, cmd_name, None)
        if cmd_func:
            args = sys.argv[2:]
            print(">> tf", cmd_name, ' '.join(args))
            if cmd_func(env, args):
                # Returned something other than None, return code 0, etc
                print("Command completed with errors.")
            return True

        # Not handling this command
        return False

    def main(self):
        
        # Verify / get environment var
        env = self.check_env()
        if not env: return env

        # Validate against the expected Terraform version
        if not tf.check_platform(): return False

        # Select Terraform workspace
        if not self.select_workspace(env): return False

        # Check for command handlers
        result = self.do_cmd(env)
        if result: return result
        
        # No special commands, hand off to Terraform with args
        args = ["terraform", *sys.argv[1:]]
        tf.print_cmd(args)
        os.execvp("terraform", args)

if __name__ == '__main__':
    tf().main()