kube-context

#!/usr/bin/python
# bugs to http://github.com/hubt
from __future__ import print_function

import os,sys,yaml,subprocess,shlex
import argparse

def shell(cmd):
  return subprocess.check_output(shlex.split(cmd))

def set_context(name):
  cmd = "kubectl config use-context " + name
  print("Executing command:\n  " + cmd)
  shell(cmd)
  
def create_context():
  print("Context Name: ",end="")
  name = sys.stdin.readline().strip()
  cluster = config['clusters'][prompt([ x['name'] for x in config['clusters']],header="Select Cluster")]['name']
  print(cluster)
  print("Namespace: ", end="")
  namespace = sys.stdin.readline().strip()
  user = config['users'][prompt([ x['name'] for x in config['users']],header="Select user")]['name']
  print(user)
  config['contexts'].append({
     "name":name,
     "context": {
       "cluster": cluster,
       "namespace": namespace,
       "user": user
     }
  })
  print(config['contexts'])
  sys.exit(0)

# given an array of options, prompts the user to select on by name or index. returns selected index or -1
def prompt(options,header="Select"):
  if header:
    print("{:>3}  {}".format("",header))
  i = 1
  for o in options:
    print("{:>3}) {}".format(i,o))
    i += 1
  print("Select(1-{}): ".format(i-1))
  selection = sys.stdin.readline().strip()
  i = 0
  for s in options:
    if selection == s:
      return i
    i += 1
  n = int(selection)
  if n >= 1 and n <= len(options):
   return n-1
  raise BaseException("Unknown selection")

parser = argparse.ArgumentParser(description="""
  alternate kubernetes configuration manager
  does not do configuration merging of multiple config files, but is more friendly
  """)
parser.add_argument("--create",default=False,action="store_true",help="create a new kubernetes context entry")
parser.add_argument("--kubeconfig",default=False,nargs=1,help="kubernetes config file to use")
parser.add_argument("context",default=False,nargs='?',help="kubernetes context to use")

args = parser.parse_args()
#print(args)
if args.kubeconfig:
  config_file = args.kubeconfig
else:
  config_file = os.environ['HOME'] + "/.kube/config"
print("Loading config file: " + config_file)
config = yaml.load(open(config_file))

if args.create:
  create_context()
# passed in parameter, use it as the context to set
if len(sys.argv) == 2:
  set_context(sys.argv[1])
  sys.exit(0)

# otherwise bring up a menu
  
contexts = config['contexts']
# sort by cluster/namespace
no_namespace="<none>"
contexts.sort(key=lambda c: c['context']['cluster'] + " "+ (c['context']['namespace'] if 'namespace' in c['context'] else no_namespace ))
i = 1
context_lines = []
print("{:>8} {:>2}  {:<20} {:<20} {:<15} {:<15}".format("CURRENT","","CONTEXT","CLUSTER","NAMESPACE","USER")) 
for x in contexts:
  current = ""
  if 'namespace' in x['context']:
    ns = x['context']['namespace']
  else: 
    ns = no_namespace
  if x['name'] == config['current-context']:
    current = "--->"
  line = "{:>8} {:>2}) {:<20} {:<20} {:<15} {:<15}".format(current,i,x['name'],x['context']['cluster'], ns,x['context']['user'])
  print(line)
  context_lines.append(line) 
  i += 1
print("Current context: " + config['current-context'])

print("Set Context (1-{0} or name): ".format(i-1),end="")
new_context = sys.stdin.readline().strip()
if new_context.strip() == '':
  print("No change")
  sys.exit(0)
context_list = filter(lambda c: c['name'] == new_context, contexts)
if len(context_list) > 0:
  context_name = context_list[0]['name']
  print("Setting context to {0}".format(context_name))
else:
  c = int(new_context) - 1
  if c < 0 or c > len(contexts) - 1:
    print("Context {} out of range".format(new_context))
    sys.exit(1)
  context_name = contexts[c]
  context_name = contexts[c]['name'] 
  print("Setting context to {0}".format(context_lines[c]))
set_context(context_name)