From 5985e6ffc94d0a4306e1096389f153a68d2f7b9a Mon Sep 17 00:00:00 2001
From: "Conor.McManus" <Conor.McManus@gfk.com>
Date: Wed, 17 Apr 2019 10:53:05 +0200
Subject: [PATCH] Add workspace manager to lock git branch with workspace

---
 atmos.py | 23 ++++++++++++++++++++++-
 1 file changed, 22 insertions(+), 1 deletion(-)

diff --git a/atmos.py b/atmos.py
index f95cf20..873bc6e 100755
--- a/atmos.py
+++ b/atmos.py
@@ -12,6 +12,7 @@ def main(argv):
         determine_actions(args, params)
 
 def determine_actions(args, params):
+    workspace_manager()
     workspace = get_env()
     env_actions = ["plan", "apply", "destroy"] # Commands that require env context
     cmd = 'terraform {args}'.format(args=args.command)
@@ -24,11 +25,31 @@ def determine_actions(args, params):
 
     if (args.t):
         generate_creds()
-
+    
     print('Terraform {args} using env vars in {env}'.format(args=args.command, env=workspace))
     with subprocess.Popen(shlex.split(cmd)) as proc:
         exit # Start process but kill py program
 
+def workspace_manager():
+    if is_git_directory != False:
+        branch = subprocess.getoutput("git rev-parse --abbrev-ref HEAD")
+        if branch == "master":
+            branch = "default"
+        else:
+            if branch not in get_valid_envs():
+                branch = "qa"
+
+        if get_env() != branch:
+            print("WARNING: Terraform workspace & git branch have diverged. Changing workspace to git branch...")
+            subprocess.call(["terraform", "workspace", "new", branch], stderr=subprocess.STDOUT, stdout=open(os.devnull, 'w'))
+            subprocess.call(["terraform", "workspace", "select", branch], stderr=subprocess.STDOUT, stdout=open(os.devnull, 'w'))
+
+def is_git_directory(path = '.'):
+    if subprocess.call(["git", "branch"], stderr=subprocess.STDOUT, stdout=open(os.devnull, 'w')) != 0:
+        return(False)
+    else:
+        return(True)
+
 def generate_creds():
     current_workspace = get_env()
     workspaces = ['default']