#!/usr/bin/python

import pycurl
import StringIO
import json
import sys
import time
import os

class register_agent(object):
    def __init__(self, rootpw, ui_addr):
        self.timeout = 50     # E.g. 50 periodes for
        self.curl_timeout = 6 #... for 6 sec, e.g. 5 minutes
        self.cookies = "/tmp/vstorage_ui_agent.cookies"
        self.username = "root"
        self.rootpw = rootpw
        self.ui_addr = ui_addr
        self.ui_defport = 8888
        self.token = ""
        self.http = "http"
        self.c = None

    def check_timeout(self, err = ""):
        if err:
            print err
        if self.timeout == 0:
            print "Timeout for VStorageUI registration exceeded"
            return False
        self.timeout -= 1
        return True

    def retry_token(self, err):
        if not self.check_timeout(err = err):
            return False
        time.sleep(self.curl_timeout)
        return self.get_token()

    def close_all(self, objs):
        for o in objs:
            o.close()

    def get_cookies(self):
        print "Started VStorageUI cookies %s acquire procedure" % self.http
        try:
            os.unlink(self.cookies)
        except:
            pass
        buf = StringIO.StringIO()
        c = pycurl.Curl()
        c.setopt(pycurl.CONNECTTIMEOUT, self.curl_timeout)
        c.setopt(pycurl.TIMEOUT, self.curl_timeout)
        c.setopt(pycurl.POST, 1)
        c.setopt(pycurl.HTTPHEADER, ["Content-type:application/json"])
        c.setopt(pycurl.POSTFIELDS, '{"username": "%s", "password": "%s"}' % (self.username, self.rootpw))
        c.setopt(pycurl.COOKIEJAR, self.cookies)
        c.setopt(pycurl.WRITEFUNCTION, buf.write)
#        c.setopt(c.VERBOSE, True)
        url = "%s://%s:%i/api/v2/login/" % (self.http, self.ui_addr, self.ui_defport)
        c.setopt(c.URL, url)
        if self.http == "https":
            c.setopt(c.SSL_VERIFYPEER, 0)
            c.setopt(c.SSL_VERIFYHOST, 2)
            c.setopt(c.FOLLOWLOCATION, 1)

        try:
            c.perform()
        except (Exception, pycurl.error) as e:
            self.close_all((c, buf))
            del(c)
            if not self.check_timeout(err = "Failed to get VStorage UI cookies: %s" % e):
                return False
            if type(e) == pycurl.error and e.args[0] != pycurl.E_OPERATION_TIMEOUTED:
                time.sleep(self.curl_timeout)
            return self.get_cookies()

        retcode = c.getinfo(pycurl.RESPONSE_CODE)
        self.close_all((c, buf))

        if retcode == 307:
            self.http = "https"
            return self.get_cookies()

        if retcode == 401:
            print "Failed to authenticate to %s:%i as %s" % (self.ui_addr, self.ui_defport, self.username)
            return False

        if retcode == 200:
            return True

        if not self.check_timeout(err = "VStorageUI cookies HTTP code: %s" % retcode):
            return False

        print "Will retry VStorageUI cookies acquire procedure..."
        time.sleep(self.curl_timeout)
        return self.get_cookies()

    def get_token(self):
        print "Started VStorageUI token acquire procedure"
        if not self.get_cookies():
            print "Failed to get UI cookies"
            return False
        t = pycurl.Curl()
        buf = StringIO.StringIO()
        t.setopt(t.URL, "%s://%s:%i/api/v2/nodes/registration/token/" % (self.http, self.ui_addr, self.ui_defport))
        t.setopt(pycurl.CONNECTTIMEOUT, self.curl_timeout)
        t.setopt(pycurl.TIMEOUT, self.curl_timeout)
        t.setopt(pycurl.COOKIEFILE, self.cookies)
        t.setopt(pycurl.WRITEFUNCTION, buf.write)
#        t.setopt(c.VERBOSE, True)
        if self.http == "https":
            t.setopt(t.SSL_VERIFYPEER, 0)
            t.setopt(t.SSL_VERIFYHOST, 2)
            t.setopt(t.FOLLOWLOCATION, 1)

        try:
            t.perform()
        except (Exception, pycurl.error) as e:
            self.close_all((t, buf))
            if not self.check_timeout(err = "Failed to get VStorage UI token: %s" % e):
                return False
            if type(e) == pycurl.error and e.args[0] != pycurl.E_OPERATION_TIMEOUTED:
                time.sleep(self.curl_timeout)
            return self.get_token()

        retcode = t.getinfo(pycurl.RESPONSE_CODE)
        answer = buf.getvalue()
        self.close_all((t, buf))

        if retcode != 200:
            return self.retry_token("VStorageUI token acquire HTTP code: %s" % retcode)

        # Parse token
        try:
            self.token = json.loads(answer)["token"]
            print "Token TTL is %s" % json.loads(answer)["ttl"]
            print "Token is %s" % self.token
        except Exception, e:
            return self.retry_token("Failed to parse VStorageUI token: %s\nServer Output:\n%s" % (e, answer))

        # Case for empty token
        if not self.token:
            return self.retry_token("Got empty VStorageUI token")

        return True

try:
    ip = sys.argv[1]
    passwd = sys.argv[2]
except:
    print "Usage: %s IP_ADDRESS PASSWORD" % sys.argv[0]
    sys.exit(1)

a = register_agent(passwd, ip)
if not a.get_token():
    sys.exit(1)

rc = os.system("echo %s | /usr/libexec/vstorage-ui-agent/bin/configure-agent.sh -x %s" % (a.token, ip))

if rc:
    print "Registration script failed"
    sys.exit(1)

print "All done!"
