#!/usr/bin/env python

#
# RHN to RHSM migration client for Red Hat Storage
# Copyright (c) 2014 Red Hat, Inc.  Distributed under GPL.
#
# Authors:
#       Anthony Towns <atowns@redhat.com>
#       Stanislav Graf <sgraf@redhat.com>
#
#  see the output of "--help" for the valid options.
#

import optparse
import os
import rpm
import subprocess
import sys


def error(msg):
    print msg
    sys.exit(1)


def run(cmd):
    proc = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE,
                            stderr=subprocess.STDOUT)
    output, dummy = proc.communicate()
    return (proc.returncode, output)


rhs_product_info = {
    "2.0": {
        "Server": {
            "rhel": ("rhel-x86_64-server-6.2.z",
                     "rhel-6.2-server-for-rhs-2.0-rpms"),
            "sfs": ("rhel-x86_64-server-sfs-6.2.z",
                    "rhel-6.2-scalefs-server-for-rhs-2.0-rpms"),
            "*rhss": ("rhel-x86_64-server-6-rhs-2.0",
                      "rhs-2.0-for-rhel-6-server-rpms"),
        },
        "Console": {
            "rhel": ("rhel-x86_64-server-6",
                     "rhel-6-server-rpms"),
            "eap6": ("jbappplatform-6-x86_64-server-6-rpm",
                     "jb-eap-6-for-rhel-6-server-rpms"),
            "*rhsc": ("rhel-x86_64-server-6-rhs-rhsc-2.0",
                      "rhsc-2.0-for-rhel-6-server-rpms"),
        },
    },
    "2.1": {
        "Server": {
            "rhel": ("rhel-x86_64-server-6.4.z",
                     "rhel-6.4-server-for-rhs-2.1-rpms"),
            "sfs": ("rhel-x86_64-server-sfs-6.4.z",
                    "rhel-6.4-scalefs-server-for-rhs-2.1-rpms"),
            "*rhss": ("rhel-x86_64-server-6-rhs-2.1",
                      "rhs-2.1-for-rhel-6-server-rpms"),
        },
        "Console": {
            "rhel": ("rhel-x86_64-server-6",
                     "rhel-6-server-rpms"),
            "eap6": ("jbappplatform-6-x86_64-server-6-rpm",
                     "jb-eap-6-for-rhel-6-server-rpms"),
            "*rhsc": ("rhel-x86_64-server-6-rhs-rhsc-2.1",
                      "rhsc-2.1-for-rhel-6-server-rpms"),
            "*nagios": ("rhel-x86_64-server-6-rhs-rhsc-2.1",
                        "rhsc-2.1-for-rhel-6-server-rpms")
        },
    },
    "3": {
        "Server": {
            "rhel": (None,
                     "rhel-6-server-rpms"),
            "sfs": (None,
                    "rhel-scalefs-for-rhel-6-server-rpms"),
            "*rhss": (None,
                      "rhs-3-for-rhel-6-server-rpms"),
            "_bigdata": (None,
                         "rhs-big-data-3-for-rhel-6-server-rpms"),
            "_nagios": (None,
                        "rhs-nagios-3-for-rhel-6-server-rpms"),
        },
        "Console": {
            "rhel": (None,
                     "rhel-6-server-rpms"),
            "eap6": (None,
                     "jb-eap-6-for-rhel-6-server-rpms"),
            "*rhsc": (None,
                      "rhsc-3-for-rhel-6-server-rpms"),
            "*nagios": (None,
                        "rhs-nagios-3-for-rhel-6-server-rpms"),
        },
    },
}


def rhn_migrate_classic_to_rhsm():
    # we have a special method here rather than calling the binary directly
    # so that we can add some mappings to the channel-cert-mapping.

    sys.path.append("/usr/share/rhsm")
    from subscription_manager.migrate import migrate

    orig_rccm = migrate.MigrationEngine.read_channel_cert_mapping

    def my_rccm(self, mappingfile):
        d = orig_rccm(self, mappingfile)
        d["rhel-x86_64-server-6.4.z"] = \
            "Server-Server-x86_64-6f455e15aed9-69.pem"
        d["rhel-x86_64-server-sfs-6.4.z"] = \
            "RHS-OnPremise-x86_64-5ae5a63102f0-186.pem"
        d["rhel-x86_64-server-6-rhs-2.1"] = \
            "RHS-OnPremise-x86_64-5ae5a63102f0-186.pem"
        d["rhel-x86_64-server-6-rhs-2.1-debuginfo"] = \
            "RHS-OnPremise-x86_64-5ae5a63102f0-186.pem"
        d["rhel-x86_64-server-6-rhs-rhsc-2.1"] = \
            "RHS-OnPremise-x86_64-5ae5a63102f0-186.pem"
        d["rhel-x86_64-server-6-rhs-rhsc-2.1-debuginfo"] = \
            "RHS-OnPremise-x86_64-5ae5a63102f0-186.pem"
        return d

    migrate.MigrationEngine.read_channel_cert_mapping = my_rccm
    migrate.MigrationEngine().main(args=["--force"])


def rhn_to_rhsm():
    status = get_status()
    if status.get("RHSM", None) == "subscribed":
        error("Already subscribed to RHSM")
    if status.get("RHN", None) != "subscribed":
        error("Can only migrate to RHSM if already registered to RHN")
    if not status.get("rhn-to-rhsm-map", {}):
        error("Can only migrate to RHSM if enabled channels are known")

    ver = set(v for v, p in status.get("product-version", []))
    if len(ver) != 1 or list(ver)[0] != "2.1":
        error("Can only migrate RHS 2.1 to RHSM")

    vs = get_pkg_vers("subscription-manager-migration",
                      "subscription-manager-migration-data")
    if len(vs) != 2:
        error(
            "Please run: yum install subscription-manager-migration subscription-manager-migration-data")

    channels = status["channels"]
    map = status["rhn-to-rhsm-map"]

    enable = set()
    for ch in channels:
        if ch not in map:
            print "Unknown channel: %s" % (ch)
    for ch in channels:
        if ch in map:
            name, repo = map[ch]
            enable.add(repo)

    print "Starting migration to RHSM..."

    rhn_migrate_classic_to_rhsm()

    print "Running \"subscription-manager repos\" to update repo list"

    auto_repos = set(get_rhsm_repos())

    if not auto_repos:
        print "Warning: not repos were automatically enabled, something likely went wrong."
        print "To correct your system's subscription, use:"
        print "    subscription-manager attach --auto"
        print "    subscription-manager repos"
        print "    yum-config-manager --disable '*'"
        print "    yum-config-manager --enable %s" % (" ".join(sorted(enable)))
        sys.exit(1)

    disable = auto_repos - enable

    if disable:
        print "Disabling repos:"
        for r in sorted(disable):
            print "  %s" % (r)
        run("yum-config-manager --disable %s" % (" ".join(sorted(disable))))
    if enable:
        print "Enabling repos:"
        for r in sorted(enable):
            print "  %s" % (r)
        run("yum-config-manager --enable %s" % (" ".join(sorted(enable))))

    print "Migration to subscription manager completed."


def upgrade(version=None):
    if version is None:
        version = max(rhs_product_info.keys())
    if version not in rhs_product_info:
        error("Target version %s not known (choices are: %s)"
              % (version, ", ".join(sorted(rhs_product_info.keys()))))

    status = get_status()
    if "product-version" not in status:
        error("Unknown RHS product/version")

    if "channels" in status:
        error(
            "Cannot upgrade system registed to RHN -- try --rhn-to-rhsm option")
    if "repos" not in status:
        error("Cannot upgrade system unless registed to RHSM")

    vs = sorted(set(pv[0] for pv in status["product-version"]))
    ps = sorted(set(pv[1] for pv in status["product-version"]))
    if len(vs) > 1:
        error("Multiple versions of RHS installed")
    cur_ver = vs[0]
    if cur_ver not in rhs_product_info:
        error("Unknown version of RHS installed: %s" % (cur_ver))

    unknown_p = [p for p in ps if p not in rhs_product_info[cur_ver]]
    if unknown_p:
        error("Unknown RHS %s product: %s" % (cur_ver,
                                              ", ".join(p for p in unknown_p)))

    noupgrade_p = [p for p in ps if p not in rhs_product_info[version]]
    if noupgrade_p:
        error("Unknown RHS %s product: %s"
              % (version, ", ".join(p for p in noupgrade_p)))

    fv = status.get("found-variants", {})

    del_repos = set()
    add_repos = set()
    miss_repos = set()

    for p in ps:
        vs = set(fv.get((cur_ver, p), []))
        del_repos.update((rhs_product_info[cur_ver][p][r][1] for r in vs))

        new_vs = set(rhs_product_info[cur_ver][p].keys())
        miss_vs = vs - new_vs
        force_vs = set(v for v in new_vs if v[0] != "_")
        get_vs = (new_vs | (force_vs - miss_vs))

        miss_repos.update((rhs_product_info[cur_ver][p][r][1] for r in miss_vs))

        add_repos.update((rhs_product_info[version][p][r][1] for r in get_vs))

    if status["product-version"][0][1] != "Console":
        print "Disable repos:\n%s" % (
            "\n".join(" %s" % (r) for r in sorted(del_repos)))
        run("yum-config-manager --disable %s" % (" ".join(sorted(del_repos))))
        print
    print "Enable repos:\n%s" % (
        "\n".join(" %s" % (r) for r in sorted(add_repos)))
    run("yum-config-manager --enable %s" % (" ".join(sorted(add_repos))))

    if miss_repos:
        print
        print "Warning: No replacements for repos:\n%s" % (
            "\n".join(" %s" % (r) for r in sorted(miss_repos)))


def get_rhsm_repos():
    rc, out = run("subscription-manager repos --list")
    if rc != 0:
        return []
    repos = []
    enabled = False
    for ln in out.split("\n"):
        if ln == "":
            if enabled and repoid is not None:
                repos.append(repoid)
            repoid = None
            enabled = False
        elif ln.startswith("Repo ID:"):
            _, repoid = ln.split(":", 1)
            repoid = repoid.strip()
        elif ln.startswith("Enabled:"):
            if "1" in ln:
                enabled = True
    if enabled and repoid is not None:
        repos.append(repoid)
    return sorted(repos)


def set_product_version(status):
    chans = set(status.get("channels", []))
    repos = set(status.get("repos", []))

    assert None not in chans and None not in repos

    found = {}
    channelmap = {}
    missing, suggested = set(), set()
    for ver in rhs_product_info:
        for prod in rhs_product_info[ver]:
            has_def_chan = False
            found[(ver, prod)] = []
            miss, sugg = set(), set()
            for repo in rhs_product_info[ver][prod]:
                c, r = rhs_product_info[ver][prod][repo]
                if c in chans or r in repos:
                    found[(ver, prod)].append(repo)
                    if repo[0] == "*":
                        has_def_chan = True
                    if c is not None:
                        channelmap[c] = (repo, r)
                else:
                    if repo.startswith("_"):
                        sugg.add((c, r))
                    else:
                        miss.add((c, r))
            if has_def_chan:
                missing.update(miss)
                suggested.update(sugg)
            else:
                del found[(ver, prod)]

    status["rhn-to-rhsm-map"] = channelmap

    if found:
        status["found-variants"] = found
        status["missing-chanrepo"] = missing
        status["suggested-chanrepo"] = suggested
        status["product-version"] = sorted(found.keys())
        return status

    pkgs = status.get("important-pkgs", [])
    if "appliance" in pkgs:
        if pkgs["appliance"][0].startswith("1.7"):
            status["product-version"] = [("2.0", "Server")]

    if "redhat-storage-server" in pkgs:
        v, r = pkgs["redhat-storage-server"]
        maj, min, _ = v.split(".", 2)
        status["product-version"] = [("%s.%s" % (maj, min), "Server")]

    assert (missing, suggested) == (set(), set())
    for v, p in status.get("product-version", []):
        if v not in rhs_product_info and "." in v:
            v, _ = v.split(".", 1)
        if v not in rhs_product_info or p not in rhs_product_info[v]:
            continue

        for repo in rhs_product_info[v][p]:
            c, r = rhs_product_info[v][p][repo]
            if c in chans or r in repos:
                pass
            else:
                if repo[0] == "_":
                    suggested.add((c, r))
                else:
                    missing.add((c, r))

    status["missing-chanrepo"] = missing
    status["suggested-chanrepo"] = suggested
    return status


def get_pkg_vers(*packages):
    ts = rpm.TransactionSet()
    mi = ts.dbMatch()
    res = {}
    for h in mi:
        if packages == [] or h['name'] in packages:
            res[h['name']] = (h['version'], h['release'])
    return res


def get_status():
    res = {}

    res["updates"] = []

    if os.path.exists("/etc/sysconfig/rhn/systemid"):
        rc, chans = run("rhn-channel -l")
        if rc == 0:
            res["RHN"] = "subscribed"
            res["channels"] = [x for x in chans.split("\n") if x]
            res["updates"].append("RHN")
        else:
            res["RHN"] = "error"

    if os.path.exists("/usr/sbin/subscription-manager"):
        rc, verinfo = run("subscription-manager version")
        if rc == 0:
            is_rhsm = False
            for x in verinfo.split("\n"):
                if not x.startswith("server type:"):
                    continue
                if "Subscription Management" in x:
                    is_rhsm = True
                break
            if is_rhsm:
                res["RHSM"] = "subscribed"
                res["repos"] = get_rhsm_repos()
                res["updates"].append("RHSM")

    if not res["updates"]:
        res["updates"].append("NOT SUBSCRIBED")

    res["important-pkgs"] = get_pkg_vers("redhat-storage-server", "appliance",
                                         "glusterfs", "ctdb", "ctdb2.5")

    set_product_version(res)

    return res


def status():
    status = get_status()

    if not status.get("product-version", None):
        print "RHS Version: unknown"
    elif len(status["product-version"]) == 1:
        print "RHS Version: RHS %s %s" % (status["product-version"][0])
    else:
        print "RHS Version:"
        for v, p in status["product-version"]:
            print "  %s %s" % (v, p)

    if "updates" in status:
        print "Updates: %s" % (" ".join(status["updates"]))

    if "RHN" in status:
        print "RHN: %s" % (status["RHN"])
    if "channels" in status:
        print "RHN Channels:"
        for ch in status["channels"]:
            print "  %s" % (ch)
        if status.get("missing-chanrepo", []):
            print "Missing RHN Channels:"
            for ch, repo in status["missing-chanrepo"]:
                if ch is None: continue
                print "  %s" % (ch)
        if status.get("suggested-chanrepo", []):
            print "Suggested RHN Channels:"
            for ch, repo in status["suggested-chanrepo"]:
                if ch is None: continue
                print "  %s" % (ch)

    if "RHSM" in status:
        print "RHSM: %s" % (status["RHSM"])
    if "repos" in status:
        print "RHSM Repos:"
        for r in status["repos"]:
            print "  %s" % (r)
        if status.get("missing-chanrepo", []):
            print "Missing RHSM Repos:"
            for ch, repo in status["missing-chanrepo"]:
                print "  %s" % (repo)
        if status.get("suggested-chanrepo", []):
            print "Suggested RHSM Repos:"
            for ch, repo in status["suggested-chanrepo"]:
                print "  %s" % (repo)

    if status.get("important-pkgs", []):
        print "Versions of selected packages:"
        for n, (v, r) in sorted(status["important-pkgs"].items()):
            print "  %s-%s-%s" % (n, v, r)


def main(args=None):
    if args is None:
        args = sys.argv[1:]

    parser = optparse.OptionParser()
    parser.add_option("--rhn-to-rhsm", action="store_true",
                      help="Upgrade from RHN to RHSM")
    parser.add_option("--upgrade", action="store_true",
                      help="Upgrade to new version of Red Hat Storage")
    parser.add_option("--status", action="store_true",
                      help="Output information about status of installation")
    parser.add_option("--version", metavar="VER",
                      help="Select version VER to upgrade to")
    (options, args) = parser.parse_args(args)

    actions = ["rhn_to_rhsm", "upgrade", "status"]
    action = [a for a in actions if getattr(options, a, False)]
    if action == []:
        parser.error("Must specify an action to take")
    elif len(action) > 1:
        parser.error("Can only specify one action to take at a time")
    action = action[0]

    if args != []:
        parser.error("This command does not accept positional arguments: %s" %
                     " ".join(args))

    if options.version is not None and action != "upgrade":
        parser.error("Can only specify version when upgrading")

    if action == "status":
        status()
    elif action == "upgrade":
        upgrade(options.version)
    elif action == "rhn_to_rhsm":
        rhn_to_rhsm()


if __name__ == "__main__":
    main()

