#!/usr/bin/python3

import os
import sys
import requests

CONSUL_API = 'http://localhost:8500'


def get_service(sess, service_id):
 r = sess.get(CONSUL_API + '/v1/agent/services', timeout=2)
 r.raise_for_status()
 services = r.json()

 for svc in services.values():
  if svc['ID'] == service_id:
   return svc

 return None


def change_service_tags(service, tags_to_add, tags_to_remove):
 with requests.Session() as sess:
  sess.headers = {'X-Consul-Token': os.getenv('CONSUL_HTTP_TOKEN')}

  svc = get_service(sess, service)
  if svc:
   new_tags = (set(svc.get('Tags', [])) | tags_to_add) - tags_to_remove
   new_svc = {
      'ID': svc['ID'],
      'Name': svc['Service'],
      'Address': svc.get('Address', ''),
      'Port': svc.get('Port', 0),
      'Meta': svc.get('Meta', {}),
      'Tags': sorted(list(new_tags)),
      'EnableTagOverride': svc.get('EnableTagOverride', False),
   }
   for k, v in new_svc.items():
      print('{} = {}'.format(k, v))
   r = sess.put(CONSUL_API + '/v1/agent/service/register', json=new_svc, timeout=2)
   r.raise_for_status()


def main(argv):
 if len(argv) < 3:
  print("Usage: consul-tag service +tag -tag...")
  return 1

 service = argv[1]
 tags_to_add = set()
 tags_to_remove = set()
 for tag in argv[2:]:
  if tag.startswith('-'):
   tags_to_remove.add(tag[1:])
  elif tag.startswith('+'):
   tags_to_add.add(tag[1:])
  else:
   tags_to_add.add(tag)

 try:
  change_service_tags(service, tags_to_add, tags_to_remove)
 except Exception as exc:
  print("Error: {}".format(exc))
  return 2

 return 0


if __name__ == '__main__':
 sys.exit(main(sys.argv))