from flask import Flask, request, jsonify
import json
import base64

app = Flask(__name__)

@app.route('/mutate', methods=['POST'])
def mutate():
   request_info = request.get_json()
   pod = request_info['request']['object']
   uid = request_info['request']['uid']

   patch = []
   if 'metadata' in pod and 'labels' in pod['metadata']:
       labels = pod['metadata']['labels']
       if 'env' not in labels:
           patch.append({
               "op": "add",
               "path": "/metadata/labels/env",
               "value": "training"
           })
   else:
       patch.append({
           "op": "add",
           "path": "/metadata/labels",
           "value": {"env": "training"}
       })

   admission_response = {
       "apiVersion": "admission.k8s.io/v1",
       "kind": "AdmissionReview",
       "response": {
           "uid": uid,
           "allowed": True,
           "patch": base64.b64encode(json.dumps(patch).encode()).decode(),
           "patchType": "JSONPatch"
       }
   }
   print (admission_response)
   print (jsonify(admission_response))
   return jsonify(admission_response)

@app.route('/validate', methods=['POST'])
def validate():
   request_info = request.get_json()
   pod = request_info['request']['object']
   uid = request_info['request']['uid']

   allowed = True
   reason = ""
   admission_response = {
       "apiVersion": "admission.k8s.io/v1",
       "kind": "AdmissionReview",
       "response": {
           "uid": uid,
           "allowed": allowed
       }
   }

   admission_response["response"]["status"] = {
           "message": reason
       }
   return jsonify(admission_response)

if __name__ == '__main__':
   context = ('tls.crt', 'tls.key')
   print(context)
   app.run(host='0.0.0.0', port=443, ssl_context=context)