#!/usr/bin/env python

import urllib.request, urllib.error, urllib.parse, json, base64, os, sys, time

def has_consul_token():
  return "CONSUL_HTTP_TOKEN" in os.environ

def get_consul_token():
  return os.environ.get('CONSUL_HTTP_TOKEN')

def url_put_json(url, data):
  opener = urllib.request.build_opener(urllib.request.HTTPHandler)
  request = urllib.request.Request(url, data.encode('utf-8'))
  request.add_header('Content-Type', 'application/json')

  if has_consul_token():
    request.add_header('X-Consul-Token', get_consul_token())

  request.get_method = lambda: 'PUT'
  return opener.open(request).read()

class ConsulDo:

  RUN = 0
  SKIP = 1

  KV_FOUND = 0
  KV_MISSING = 1
  KV_ERROR = 2

  def __init__(self):
    self.key = ""
    self.node = ""
    self.base_url = "http://localhost:8500"
    self.session = ""
    self.session_node = ""
    self.debug = False
    self.session_parms = {}

  def log(self, message):
    if self.debug:
      print(message)

  def get_node_name(self):
    url = "%s/v1/agent/self" % self.base_url
    request = urllib.request.Request(url)
    if has_consul_token():
      request.add_header('X-Consul-Token', get_consul_token())
    response = urllib.request.urlopen(request).read()
    data = json.loads(response)
    return data['Config']['NodeName']

  def get_store(self):
    url = "%s/v1/kv/service/%s/leader" % (self.base_url, self.key)
    try:
      if has_consul_token():
        headers = {'X-Consul-Token': get_consul_token()}
        req = urllib.request.Request(url, None, headers)
        response = urllib.request.urlopen(req).read()
      else:
        response = urllib.request.urlopen(url).read()
    except urllib.error.HTTPError as e:
      if e.code == 404:
        return self.KV_MISSING
      else:
        return self.KV_ERROR
    except Exception as e:
      self.log("Could not get datastore: %s" % str(e))
      return self.KV_ERROR
    else:
      data = json.loads(response)
      if "Session" in data[0]:
        self.session = data[0]["Session"]
      if "Value" in data[0]:
        self.session_node = json.loads(base64.b64decode(data[0]["Value"]))["node"]
      return self.KV_FOUND

  def create_session(self):
    url = "%s/v1/session/create" % self.base_url
    try:
      self.log("Creating session")
      self.session_parms['LockDelay'] = '30s'
      self.session_parms['Behavior'] = 'release' # default
### if a session TTL is warranted, this is where it belongs. 
### I am not convinced this is a 'good' thing.  I believe it
### will possibly break the intention of the distributed lock.
#      self.session_parms['TTL'] = '30s' 
      response = url_put_json(url, json.dumps(self.session_parms))
    except Exception as e:
      self.log("Could not create session: %s" % str(e))
      sys.exit(self.SKIP)
    else:
      session_id = json.loads(response)["ID"]
      self.log("Setting session id to %s" % session_id)
      self.session = session_id

  def acquire_session(self):
    url = "%s/v1/kv/service/%s/leader?acquire=%s" % (self.base_url, self.key, self.session)
    try:
      self.log("Acquiring session")
      response = url_put_json(url, json.dumps({'node': self.node, 'updated': time.time()}))
    except Exception as e:
      self.log("Could not acquire session: %s" % str(e))
      sys.exit(self.SKIP)
    else:
      if response == "true":
        self.log("Got the session!")
        sys.exit(self.RUN)
      else:
        self.log("Couldn't get the session")
        sys.exit(self.SKIP)


  def become_leader(self):
    self.log("Trying to become leader")
    self.create_session()
    self.acquire_session()


  def run(self):
    self.log("Running...")
    self.node = self.get_node_name()
    self.log("Found nodename: %s" % self.node)

    r = self.get_store()
    if r == self.KV_FOUND:
      self.log("Found KV store")
      if self.session != "":
        self.log("Found session: %s" % self.session)
        if self.session_node == self.node:
          self.log("Found session node, we're the leader")
          sys.exit(self.RUN)
        else:
          self.log("Found session node, but %s is the leader" % self.session_node)
          sys.exit(self.SKIP)
      else:
        self.log("No session found")
        self.become_leader()
    elif r == self.KV_MISSING:
      self.log("No KV entry found")
      self.become_leader()
    else:
      self.log("Could not read KV store")
      sys.exit(self.SKIP)


if __name__ == '__main__':
  c = ConsulDo()
  if len(sys.argv) < 2:
    print("Usage: consul-do KEY NODE [DEBUG]")
    sys.exit(1)

  c.key = sys.argv[1]

  if len(sys.argv) >= 3:
    c.debug = True
  c.run()
