#!/usr/bin/env python3

import json
import sys
import os
import _jsonnet
import requests

# Rack awareness logic
node_labels = []
rack_labels = []
env_var = 'RACK_NODE_LABELS'
token_path = '/var/run/secrets/kubernetes.io/serviceaccount/token'
if env_var in os.environ:
    rack_labels = json.loads(os.environ[env_var])
    service_host = os.environ['KUBERNETES_SERVICE_HOST']
    tcp_port = os.environ['KUBERNETES_SERVICE_PORT_HTTPS']
    hostname = os.environ['HOSTNAME']
    namespace = os.environ['POD_NAMESPACE']

    with open(token_path) as f:
        kube_token = f.read()
    headers = {'Authorization': 'Bearer %s' % kube_token}

    # get label value for rack awareness node label
    api_url = 'https://%s:%s/api/v1/' % (service_host, tcp_port)

    r = requests.get(api_url + 'namespaces/%s/pods/%s' % (namespace, hostname), headers=headers, verify=False)
    if r.status_code != 200:
        sys.exit(r.json())
    node_name = r.json()['spec']['nodeName']

    r = requests.get(api_url + 'nodes/%s' % node_name, headers=headers, verify=False)
    if r.status_code != 200:
        sys.exit(r.json())

    node_labels = r.json()['metadata']['labels']
    for label in rack_labels:
        if label not in node_labels:
            sys.exit('node label: %s not found in pod node metadata' % label)

pod_name = sys.argv[1]
data_file = sys.argv[2]
jsonnet_file = sys.argv[3]
output_file = sys.argv[4]

with open(data_file) as df:
    data = json.load(df)

with open(jsonnet_file, "r") as jf:
    jsonnet_str = jf.read()

json_str = _jsonnet.evaluate_snippet(
    "snippet",
    jsonnet_str,
    ext_vars={
        'podName': pod_name,
        'data': json.dumps(data),
        'nodeLabels': json.dumps(node_labels),
        'rackLabels': json.dumps(rack_labels)
    },
)

json_obj = json.loads(json_str)
with open(output_file, 'w') as of:
    json.dump(json_obj, of)
