#!/usr/bin/env python2

import sys
import json
import platform
import collections
import subprocess
import os
import urlparse
import abc
import base64
import tempfile
import time

MYLINBIT = "https://my.linbit.com/"
# MYLINBIT = "http://ruby-dev13:3001"

API_URL = urlparse.urljoin(MYLINBIT, "api/v1/register_node")
LINBIT_PLUGIN_BASE = urlparse.urljoin(MYLINBIT, "yum-plugin/")
MYNAME = "linbit-manage-node.py"
SELF = urlparse.urljoin(MYLINBIT, MYNAME)
GPG_KEY = "https://packages.linbit.com/package-signing-pubkey.asc"
LINBIT_PLUGIN = urlparse.urljoin(LINBIT_PLUGIN_BASE, "linbit.py")
LINBIT_PLUGIN_CONF = urlparse.urljoin(LINBIT_PLUGIN_BASE, "linbit.conf")
NODE_REG_DATA = "/var/lib/drbd-support/registration.json"

REMOTEVERSION = 1
# VERSION has to be in the form "MAJOR.MINOR"
VERSION = "1.11"

api_retcodes = {
    "ok": 0,
    "need_node": -1,
    "need_contract": -2,
    "need_cluster": -3,
    "no_nodes_left": -4,
}


RESTResponse = collections.namedtuple("RESTResponse",
                                      "retcode options proxy nodehash")


# Utility Functions that might need update (e.g., if we add distro-types)
def getHostInfo():
    if platform.system() != "Linux":
        err("You have to run this script on a GNU/Linux based system")

    try:
        distname, version, distid = platform.linux_distribution()
    except AttributeError:
        distname = ""

    distname = distname.strip().lower()
    version = version.strip()

    # python spec is bit fuzzy about the hostname:
    # "may not be fully qualified!". To be save, use the first part only
    hostname = platform.node().strip().split('.')[0]

    if distname.startswith("red hat enterprise linux server") or \
       distname.startswith("centos") or \
       distname.startswith("oracle linux server"):
        version_split = version.split('.')
        if len(version_split) > 2:
            version = "%s.%s" % tuple(version_split[:2])
        dist = "rhel%s" % (version)
    elif distname.startswith("ubuntu"):
        rel = getDebianUbuntuRelease()
        dist = "ubuntu-%s" % (rel)  # ubuntu-lucid
    elif distname.startswith("debian"):
        rel = getDebianUbuntuRelease()
        dist = "debian-%s" % (rel)  # debian-jessie
    elif distname.startswith("suse linux enterprise server"):
        dist = "sles%s" % (version)
        # patchlevel is not included in platform-info
        patchlevel = getSuSEPatchLevel()
        if patchlevel and patchlevel != '0':
            dist = dist + "-sp" + patchlevel
    else:
        dist = False

    # it seems really hard to get MAC addresses if you want:
    # a python only solution, e.g., no extra C code
    # no extra non-built-in modules
    # support for legacy python versions
    macs = set()
    # we are Linux-only anyways...
    CLASSNET = "/sys/class/net"
    if os.path.isdir(CLASSNET):
        for dev in os.listdir(CLASSNET):
            devpath = os.path.join(CLASSNET, dev)

            if not os.path.islink(devpath):
                continue

            with open(os.path.join(devpath, "type")) as t:
                dev_type = t.readline().strip()
                if dev_type != '1':  # this filters for example ib/lo devs
                    continue

            # try to filter non permanent interfaces
            # very old kernels do not have /sys/class/net/*/addr_assign_type
            addr_assign_path = os.path.join(devpath, "addr_assign_type")
            if os.path.isfile(addr_assign_path):
                with open(addr_assign_path) as a:
                    dev_aatype = a.readline().strip()
                    if dev_aatype != '0' and dev_aatype != '1':  # NET_ADDR_PERM
                        continue
            else:  # try our best to manually filter them
                if dev.startswith("vir") or \
                   dev.startswith("vnet") or \
                   dev.startswith("bond"):
                    continue

            with open(os.path.join(devpath, "address")) as addr:
                mac = addr.readline().strip()
                macs.add(mac)

    def getFamily(dist):
        family = False
        if dist:
            if dist.startswith("rhel"):
                family = "redhat"
            if dist.startswith("sles"):
                family = "suse"
            if dist.startswith("debian"):
                family = "debian"
            if dist.startswith("ubuntu"):
                family = "debian"
        return family

    return dist, getFamily(dist), hostname, macs


def setupConfig(urlhandler, dist, family, config, free_running=False):
    # Write repository configuration
    if family == "debian":
        repo_file = "/etc/apt/sources.list.d/linbit.list"
    elif family == "redhat":
        repo_file = "/etc/yum.repos.d/linbit.repo"
    elif family == "suse":
        repo_file = "/etc/zypp/repos.d/linbit.repo"

    if not free_running:
        printcolour("Repository configuration:\n", GREEN)
        print "It is perfectly fine if you do not want to enable any repositories now."
        print "The configuration for disabled repositories gets written,"
        print "but the repositories are disabled for now."
        print "You can edit the configuration (e.g., %s) file later to enable them.\n" % (repo_file)

    repo_content = []
    repo_names = []
    if family == "debian":
        for line in config:
            config_split = line.split()
            for name in config_split[3:]:
                repo_names.append(name)

        enabled = ask_enable(repo_names, free_running)
        for line in config:
            config_split = line.split()
            for name in config_split[3:]:
                line = " ".join(config_split[:3])
                line += " " + name
                if not enabled.get(name):
                    line = "# " + line
                repo_content.append(line + '\n\n')

    elif family == "redhat" or family == "suse":
        lines = []
        for line in config:
            lines.append(line.strip())
            repo_names.append(line.split('/')[-2])

        enabled = ask_enable(repo_names, free_running)
        for line in lines:
            name = line.split('/')[-2]
            repo_content.append("[%s]\n" % (name))
            repo_content.append("name=LINBIT Packages for %s - $basearch\n" % (name))
            repo_content.append("%s\n" % (line))
            if family == "suse":
                repo_content.append("type=rpm-md\n")
            if enabled.get(name):
                repo_content.append("enabled=1\n")
            else:
                repo_content.append("enabled=0\n")
            repo_content.append("gpgkey=%s\n" % (GPG_KEY))
            repo_content.append("gpgcheck=1\n")
            repo_content.append("\n")

    printcolour("Writing repository config:\n", GREEN)
    if len(repo_content) == 0:
        if len(config) == 0:
            repo_content.append("# Could not find any repositories for your distribution\n")
            repo_content.append("# Please contact support@linbit.com\n")
        else:
            repo_content.append("# Repositories found, but none enabled\n")
    success = writeFile(repo_file, repo_content, free_running=free_running)
    if success:
        OK('Repository configuration written')

    # Download yum plugin on yum based systems
    # if family == "redhat" && False:
    if family == "redhat":
        printcolour("Downloading LINBIT yum plugin\n", GREEN)
        FINAL_PLUGIN = LINBIT_PLUGIN
        if dist.startswith('rhel6'):
            FINAL_PLUGIN += '.6'
        elif dist.startswith('rhel7'):
            FINAL_PLUGIN += '.7'
        f = urlhandler.fileHandle(FINAL_PLUGIN)
        plugin = [pluginline for pluginline in f]
        writeFile("/usr/share/yum-plugins/linbit.py", plugin,
                  showcontent=False, askforwrite=False, free_running=free_running)

        printcolour("Downloading LINBIT yum plugin config\n", GREEN)
        f = urlhandler.fileHandle(LINBIT_PLUGIN_CONF)
        cfg = [cfgline for cfgline in f]
        writeFile("/etc/yum/pluginconf.d/linbit.conf", cfg,
                  showcontent=False, askforwrite=False, free_running=free_running)

    return True


def main():
    free_running = False
    proxy_only = False
    non_interactive = False

    args = {}
    args["version"] = REMOTEVERSION
    # urlhandler = requestsHandler()
    urlhandler = urllib2Handler()

    if os.path.isfile(NODE_REG_DATA):
        with open(NODE_REG_DATA) as infile:
            jsondata = json.load(infile)
            args["nodehash"] = jsondata["nodehash"]

    opts = sys.argv[1:]
    for opt in opts:
        if opt == "-p":
            proxy_only = True
            sys.argv.remove("-p")
            if not args["nodehash"]:
                err('Your node is not registered, first run this script without "-p".'
                    "\nMake sure %s exists!" % (NODE_REG_DATA))
            args["proxy_only"] = True

    e_user = os.getenv('LB_USERNAME', None)
    e_pwd = os.getenv('LB_PASSWORD', None)
    e_cluster = os.getenv('LB_CLUSTER_ID', None)
    e_contract = os.getenv('LB_CONTRACT_ID', None)

    e_all = e_user and e_pwd and e_cluster and e_contract
    e_one = e_user or e_pwd or e_cluster or e_contract

    if e_one and not e_all:
        err('You have to set all (or none) of the required environment variables (LB_USERNAME, LB_PASSWORD, LB_CLUSTER_ID, and LB_CONTRACT_ID)')
    if e_all and proxy_only:
        err('You are not allowed to mix "-p" and non-interactive mode')

    non_interactive = e_all
    free_running = proxy_only or non_interactive

    if proxy_only:
        token = "N:" + args["nodehash"]
        headers = createHeaders(token)
    else:
        force_user_input = False
        print MYNAME, "(Version: %s)" % (VERSION)
        checkVersion(urlhandler)

        while True:
            if non_interactive:
                token = '%s:%s' % (e_user, e_pwd)
            else:
                token = getToken(force_user_input)

            username = token.split(':')[0]
            token = "C:" + token

            headers = createHeaders(token, username)

            # create a first request to test UN/PWD
            ret = urlhandler.postRESTRequest(headers, args)
            if ret.retcode == "failed":
                msg = "Username and/or Credential are wrong"
                if non_interactive:
                    err(msg)
                else:
                    warn(msg)
                force_user_input = True
            else:
                OK("Login successful")
                break

    dist, family, hostname, macs = getHostInfo()

    if len(macs) == 0:
        err("Could not detect MAC addresses of your node")

    if not dist and not free_running:
        print "Distribution information could not be retrieved"
        contactInfo(executeCommand("uname -a"))
        print "You can still register your node, but the script will not"
        print "write a repository configuration for this node"
        cont_or_exit()
        dist = "Unknown"

    if not isRoot():
        if free_running:
            err("You have to execute this script as super user")
        print "You are not running this script as super user"
        print ""
        print "There are two choices:"
        print "-) Abort now and restart the script (please use su/sudo)"
        print "-) Continue:"
        print "  - Registration itself does not require super user permissions"
        print "  - BUT the repository configuration will only be printed"
        print "    and written to /tmp"
        print ""
        cont_or_exit()

    # XXX
    # fake redhat
    # dist = "rhel7.2"
    # family = "redhat"
    #
    # fake suse
    # dist = "sles11-sp3"
    # family = "suse"
    #
    # fake debian
    # dist = "debian-wheezy"
    # family = "debian"
    # XXX

    args["hostname"] = hostname
    args["distribution"] = dist
    args["mac_addresses"] = ','.join(macs)
    if non_interactive:
        args["contract_id"] = e_contract
        args["cluster_id"] = e_cluster

    ret = urlhandler.postRESTRequest(headers, args)
    postsfailed = 0
    while ret.retcode != api_retcodes["ok"]:
        if ret.retcode == "failed":
            postsfailed += 1
            if postsfailed >= 3:
                err("Could not post request, giving up")

        if ret.retcode == api_retcodes["need_contract"]:
            if non_interactive:
                err('Not a valid CONTRACT_ID')
            if len(ret.options) == 0:
                err("Sorry, but you do not have any valid contract for this credential")

            print "The following contracts are available:"
            selection = getOptions(ret.options, what="contract")
            args["contract_id"] = selection
        elif ret.retcode == api_retcodes["need_cluster"]:
            if non_interactive:
                err('Not a valid CLUSTER_ID')
            selection = getOptions(ret.options, allow_new=True, what="cluster")
            args["cluster_id"] = selection
        elif ret.retcode == api_retcodes["need_node"]:
            warn("The script could not determinde all required information")
            contactInfo(args)
            sys.exit(1)
        elif ret.retcode == api_retcodes["no_nodes_left"]:
            err("Sorry, but you do not have any nodes left for this contract")

        # print "D*E*B*U*G, press enter"
        # raw_input()
        ret = urlhandler.postRESTRequest(headers, args)

    if ret.retcode == api_retcodes["ok"]:
        if not free_running:
            printcolour("Writing registration data:\n", GREEN)
        args_save = {}
        args_save["nodehash"] = ret.nodehash
        for tosave in ["distribution", "hostname", "mac_addresses"]:
            args_save[tosave] = args[tosave]
        writeFile(NODE_REG_DATA, args_save, showcontent=False,
                  free_running=free_running, asjson=True)
        if dist != "Unknown" and family:
            if ret.proxy:
                if not free_running:
                    printcolour("Writing proxy license:\n", GREEN)
                license = [x + '\n' for x in base64.b64decode(ret.proxy).split('\n')]
                writeFile("/etc/drbd-proxy.license", license,
                          showcontent=False, free_running=free_running)

            if not free_running or non_interactive:
                setupConfig(urlhandler, dist, family, config=ret.options, free_running=non_interactive)
            if not free_running:  # RCK THINK
                epilogue(family, urlhandler)
    else:
        err(sys.argv[0] + " exited abnormally")

    if not free_running:
        OK("Congratulations! Your node was successfully configured.")
    sys.exit(0)


# Utility functions that are unlikely to require change
def checkVersion(urlhandler):
    import re
    printcolour("Checking if version is up to date\n", GREEN)
    outdated = False

    # we do not want to fail if anything is wrong here...
    try:
        f = urlhandler.fileHandle(SELF)
        selfpy = [selfline for selfline in f]
        p = re.compile('^VERSION.*(\d+)\.(\d+).*')

        upstream_major = sys.maxint
        upstream_minor = 0

        for line in selfpy:
            m = p.match(line.strip())
            if m:
                upstream_major = int(m.group(1))
                upstream_minor = int(m.group(2))
                break

        v = VERSION.split('.')
        my_major = int(v[0])
        my_minor = int(v[1])

        if my_major < upstream_major:
            outdated = True
        elif my_major == upstream_major and my_minor < upstream_minor:
            outdated = True

        if outdated:
            warn("Your version is outdated")
            tmpf = tempfile.mkstemp(suffix='_' + MYNAME)[1]
            writeFile(tmpf, selfpy, showcontent=False, askforwrite=False,
                      hinttocopy=False)
            OK("New version downloaded to %s" % (tmpf))
        else:
            OK("Your version is up to date")
    except:
        warn("Version check failed, but continuing anyways")

    if outdated:
        sys.exit(0)


def executeCommand(command):
    pyvers = sys.version_info
    if pyvers[0] == 2 and pyvers[1] == 6:
        output = subprocess.Popen(command, shell=True,
                                  stdout=subprocess.PIPE).communicate()[0]
    else:
        output = subprocess.check_output(command, shell=True)
    return output


# content is a list of of lines
def writeFile(name, content, showcontent=True, askforwrite=True,
              free_running=False, asjson=False, hinttocopy=True):
    origname = name
    if not isRoot():
        name = os.path.join("/tmp", os.path.basename(name))

    if showcontent and not free_running:
        print "Content:"
        for line in content:
            print line,
        print ""

    if askforwrite and not free_running:
        if askYesNo("Write to file (%s)?" % (name)):
            if os.path.isfile(name):
                print "File: " + name + " exists"
                if not askYesNo("Overwrite file?"):
                    return False
        else:
            return False

    dirname = os.path.dirname(name)
    if not os.path.exists(dirname):
        os.makedirs(dirname)

    with open(name, "w") as outfile:
        if asjson:
            json.dump(content, outfile)
        else:
            for line in content:
                outfile.write(line)

    if not isRoot() and hinttocopy:
        printcolour("Important: ", MAGENTA)
        print
        print "Please review", name, "and copy file to", origname

    return True


def getToken(force_user_input):
    if len(sys.argv) == 1 or force_user_input:
        import getpass
        while True:
            username = raw_input("Username: ")
            if username:
                break

        while True:
            pwd = getpass.getpass("Credential (will not be echoed): ")
            if pwd:
                break
        return "%s:%s" % (username.strip(), pwd.strip())
    elif len(sys.argv) == 2:
        return sys.argv[-1]


def err(string):
    printcolour("ERR: ", RED)
    print string
    sys.exit(1)


def warn(string):
    printcolour("W: ", MAGENTA)
    print string


def contactInfo(args):
    print "Please report this issue to:"
    print "\tdrbd-support@linbit.com"
    print ""
    print "Make sure to include the following infomation:"
    print "%s - Version: %s" % (os.path.basename(sys.argv[0]), VERSION)
    print args


def askYesNo(question):
    printcolour("--> ", CYAN)
    ret = raw_input(question + " [y/N] ")
    ret = ret.strip().lower()
    if ret == 'y' or ret == "yes":
        return True
    else:
        return False


def cont_or_exit():
    if not askYesNo("Continue?"):
        sys.exit(0)


def createHeaders(token, username=None):
    headers = {
        'Content-Type': 'application/json',
        'Authorization': 'Token token=' + token.strip()
    }
    agent = "ManageNode/%s" % (VERSION)
    if username:
        agent += " (U:%s)" % (username)
    headers['User-agent'] = agent

    return headers


def isRoot():
    return os.getuid() == 0


def getOptions(options, allow_new=False, what="contract"):
    # dicts have no guaranteed order, so we use an array to keep track of the
    # keys
    lst = []
    new = 0
    e = -1  # set it in case len(options) == 0

    print "Will this node form a cluster with...\n"
    for e, k in enumerate(sorted(options)):
        lst.append(k)
        if what == "contract":
            print "%d) Contract: %s (ID: %s)" % (e + 1, options[k], k)
        elif what == "cluster":
            print "%d) Nodes: %s (Cluster-ID: %s)" % (e + 1, options[k], k)
        else:
            err("Unknown selection option")

    if allow_new:
        printcolour("%d) *Be first node of a new cluster*\n" % (e + 2), CYAN)
        new = 1
    print ""

    while True:
        printcolour("--> ", CYAN)
        nr = raw_input("Please enter a number in range and press return: ")
        try:
            nr = int(nr.strip()) - 1  # we are back to CS/array notion
            if nr >= 0 and nr < len(options) + new:
                if allow_new and nr == e + 1:
                    return -1
                else:
                    return lst[nr]
        except ValueError:
            pass


def getDebianUbuntuRelease():
    output = executeCommand("lsb_release -c")
    return output.split(':')[1].strip().lower()


def getSuSEPatchLevel():
    with open("/etc/SuSE-release") as SuSErel:
        for line in SuSErel:
            line = line.strip()
            if line.startswith("PATCHLEVEL"):
                return line.split('=')[1].strip()
    return False


def epilogue(family, urlhandler):
    printcolour("Final Notes:", GREEN)
    print ""

    tmpf = tempfile.mkstemp()[1]

    f = urlhandler.fileHandle(GPG_KEY)
    key = [keyline for keyline in f]
    writeFile(tmpf, key,
              showcontent=False, askforwrite=False, hinttocopy=False)

    if isRoot() and askYesNo("Add linbit signing key to keyring now?"):
        addkey = False
        if family == "redhat" or family == "suse":
            addkey = "rpm --import %s" % (tmpf)
        elif family == "debian":
            addkey = "apt-key add %s" % (tmpf)

        if addkey:
            output = executeCommand(addkey)
            print output
    else:
        print "Download package signing key from %s and import it manually!" % (GPG_KEY)

    print "Now update your package information and install"
    print "LINBIT's kernel module and/or user space utilities"


# implement REST calls, actually, the 'requests' module is the way to go.
# unfortunately  we have to support (very) old version of python and
# do not want to force customers to install (cooler) extra modules...
class URLHanlder(object):
    __metaclass__ = abc.ABCMeta

    @abc.abstractmethod
    def postRESTRequest(self, headers, payload):
        return

    @abc.abstractmethod
    def fileHandle(self, url):
        return


class requestsHandler(URLHanlder):
    def __init__(self):
        try:
            import requests
            self.callm = requests
        except ImportError:
            err("'requests' module is not installed")

    def postRESTRequest(self, headers, payload):
        r = self.callm.post(API_URL, data=json.dumps(payload), headers=headers)

        if r.status_code != self.callm.codes.ok:
            warn("Could not post request.")
            return RESTResponse("failed", False, False, False)

        ret = r.json()
        return RESTResponse(retcode=int(ret["ret"]), options=ret["options"],
                            proxy=ret["proxy_license_file"],
                            nodehash=ret["nodehash"])

    def fileHandle(self, url):
        return self.callm.get(url, stream=True)


class urllib2Handler(URLHanlder):
    def __init__(self):
        try:
            import urllib2
            self.callm = urllib2
        except ImportError:
            err("'urllib2' module is not installed")

    def postRESTRequest(self, headers, payload):
        r = self.callm.Request(API_URL,
                               data=json.dumps(payload), headers=headers)
        try:
            f = self.callm.urlopen(r)
        except self.callm.URLError as e:
            if str(e) == "HTTP Error 403: Forbidden":
                return RESTResponse("failed", False, False, False)
            else:
                err("urllib2 returned: " + str(e))

        ret = f.read()
        ret = json.loads(ret)

        return RESTResponse(retcode=int(ret["ret"]), options=ret["options"],
                            proxy=ret["proxy_license_file"],
                            nodehash=ret["nodehash"])

    def fileHandle(self, url):
        return self.callm.urlopen(url)


# following from Python cookbook, #475186
BLACK, RED, GREEN, YELLOW, BLUE, MAGENTA, CYAN, WHITE = range(8)


def has_colours(stream):
    if not hasattr(stream, "isatty"):
        return False
    if not stream.isatty():
        return False  # auto color only on TTYs
    try:
        import curses
        curses.setupterm()
        return curses.tigetnum("colors") > 2
    except:
        # guess false in case of error
        return False


def printcolour(text, colour=WHITE):
    if has_colours:
        seq = "\x1b[1;%dm" % (30+colour) + text + "\x1b[0m"
        sys.stdout.write(seq)
    else:
        sys.stdout.write(text)


def OK(text):
    print '[',
    printcolour("OK", GREEN)
    print ']',
    print text


def ask_enable(names, free_running=False):
    """Asks the user which repos they wish to enable.

    Args:
        names: A list of repo names to be enabled/disabled.
        free_running: A blooean indicating if there will be no user input.
    Returns:
        An array of dicts, keys are repo names, values are True if repo is
            enabled.
    """

    enabled_by_default = False

    if free_running:
        enabled_by_default = True

    repos = []
    # Sort reverse to try to show newest versions first.
    for name in sorted(names, reverse=True):
        repos.append([name, enabled_by_default])

    # Skip asking questions in non-interacting mode.
    while not free_running:
        idx_offset = 1  # For converting between zero and one indexed arrays.
        os.system("clear")
        print "\n  Here are the repositories you can enable:\n"
        for index, repo in enumerate(repos):

            name, value = repo

            status = "Disabled"
            display_color = RED

            if value:
                status = "Enabled"
                display_color = GREEN

            printcolour(
                "    {0}) {1}({2})\n".format(index + idx_offset, name, status),
                display_color
            )

        print("\n  Enter the number of the repository you "
              "wish to enable/disable. Hit 0 when you are done.\n")
        choice = raw_input("  Enable/Disable: ")

        # Ignore random button mashing.
        if not choice:
            continue

        try:
            choice = int(choice.strip())
        except ValueError:
            print "\n  You must enter a number!\n"
            time.sleep(1)
            continue

        if choice == 0:
            break

        choice_idx = choice - idx_offset
        try:
            # Toggle Enabled/Disabled
            repos[choice_idx][1] = not repos[choice_idx][1]
        except IndexError:
            # User will see if state of the repos change,
            # No need to complain.
            pass

    repo_map = {}
    for repo in repos:
        name, enabled = repo
        repo_map[name] = enabled

    return repo_map


if __name__ == "__main__":
    has_colours = has_colours(sys.stdout)
    try:
        main()
    except KeyboardInterrupt:
        print ""
        warn("Received Keyboard Interrupt signal, exiting...")