#!/usr/bin/python2.7 -E
#
# Copyright (c) 2016, Oracle and/or its affiliates. All rights reserved.
#

#
#
# This enumerator checks the system's cpu and platform banner against
# a set of minimum requirements. If the system does not meet that
# criteria, we print a message to stderr and return 1.
# If the system meets the criteria, we return 0 and print nothing.
#
# To invoke this utility, provide a check.version-type= argument, and one
# of check.exclude= or check.include= as well.
#
# For example, to exclude systems prior to SPARC-T4, you might invoke
# this utility with
#
# check.version-type=cpu check.exclude=SPARC-T1,SPARC-T2,SPARC-T3,SPARC64-V
#
# It's considerably easier to specify a cpu rather than a platform,
# since platform names may contain whitespace, and there are fewer
# variants of cpus (for SPARC, that is).
#
# To exclude x64 systems which do not have the SSE3 instruction set,
# you would invoke this utility with
#
# check.version-type=iset check.include=SSE3

from __future__ import print_function

import subprocess
import re
import sys


ENUMERATOR_ERROR = 242
VERSIONTYPE = "check.version-type"
ISEXCLUDE = "exclude"
ISINCLUDE = "include"
KSTAT = "/usr/bin/kstat"
ISAINFO = "/usr/bin/isainfo"
DEVPROP = "/usr/sbin/devprop"


def usage():
    print("cpu check.version-type=[plat|cpu|iset] "
          "check.include={csv list of arguments}",
          file=sys.stderr)
    print("cpu check.version-type=[plat|cpu|iset] "
          "check.exclude={csv list of arguments}",
          file=sys.stderr)


def do_exec(cmd=None, opts=None):
    """
    This function execs our external command, returns a buffer
    on success, and raises an exception on failure.
    """
    sys.stdout.flush()
    sys.stderr.flush()
    args = [cmd]
    args.extend(opts)
    try:
        proc = subprocess.Popen(args,
                                stdout=subprocess.PIPE,
                                stderr=subprocess.PIPE,
                                shell=False, close_fds=True)
        buf = proc.stdout.readlines()
        ret = proc.wait()
    except OSError as e:
        buf = (("Cannot exec {0}: {1}").format(" ".join(args), str(e)))
        ret = -1

    if ret < 0:
        # grumble
        raise OSError(msg="'{0}' returned ({1}): {2}\n".format(
            args, ret, buf), file=sys.stderr)
    return (ret, buf)


def check(exclude=True, s1=None, s2=None):
    """
    This function does the nitty gritty of the in/out check.
    Both s1 and s2 are sets, which makes this check function
    very easy to implement.

    Returns (True, None) or (False, "matching elements")
    """
    if exclude:
        matchset = s1 & s2
    else:
        matchset = s2 - s1
    if matchset:
        return (False, matchset)
    else:
        return (True, None)


def instructions(iset=None, exclude=True):
    """
    This function parses the output from getisainfo(1) to retrieve
    the list of capabilities present in this system's cpu.
    """

    opts = ["-b", "-v"]
    (status, buf) = do_exec(cmd=ISAINFO, opts=opts)

    # With that out of the way, let's process the buffer
    caps = set()
    for line in buf:
        if not line.startswith("64-bit"):
            caps.update(line.split())
    return check(exclude=exclude, s1=caps, s2=iset)


def platforms(plats=None, exclude=True):
    """
    This function parses the output from devprop(1m) to get
    the platform name for this system.
    """

    (status, buf) = do_exec(cmd=DEVPROP, opts=["compatible"])

    # With that out of the way, let's process the buffer
    banner = buf[0].split(" ")[0].split(".")[1].strip()

    # Since we want to allow platforms specified as base forms
    # (ie, SPARC-T5 rather than SPARC-T5-1,SPARC-T2-2,SPARC-T5-4 etc)
    # we do a bit of munging on the set before passing to the check
    # function above.
    compatible = set([banner])
    for el in plats:
        if banner.find(el) > -1:
            compatible = set([el])
    return check(exclude=exclude, s1=compatible, s2=plats)


def cpus(procs=None, exclude=True):
    """
    This function uses kstat(1) to retrieve the cpu implementation
    statistic. Since we do not allow different CPU types in a host,
    we only need to find the lowest-numbered CPU and use the
    implementation stat from it.
    """

    opts = ["-p", "-m", "cpu_info", "-s", "implementation"]
    (status, buf) = do_exec(cmd=KSTAT, opts=opts)

    # With that out of the way, let's process the buffer
    cpure = re.compile("""(.*implementation\s+)(\S+)( .*)$""")
    imp = cpure.match(buf[0])
    impl = set()
    impl.add(imp.group(2))
    return check(exclude=exclude, s1=impl, s2=procs)


if __name__ == "__main__":

    if "-h" in sys.argv:
        usage()
        sys.exit(0)

    if len(sys.argv) < 3 not in sys.argv:
        print("Missing arguments (check.version-type and "
              "check.[include|exclude])\n",
              file=sys.stderr)
        usage()
        sys.exit(ENUMERATOR_ERROR)

    vtype = str()
    inout = set()
    exclude = True
    for n in range(1, len(sys.argv)):
        key, value = sys.argv[n].split("=", 1)
        if key == "check.version-type":
            vtype = value
        elif key in ["check.include", "check.exclude"]:
            inout = set(value.split(","))
            if key == "check.include":
                exclude = False
        else:
            # fingers in ears time - ignore what we don't understand
            pass

    if vtype == "plat":
        (status, message) = platforms(plats=inout, exclude=exclude)
    elif vtype == "cpu":
        (status, message) = cpus(procs=inout, exclude=exclude)
    elif vtype == "iset":
        (status, message) = instructions(iset=inout, exclude=exclude)
    else:
        sys.exit(ENUMERATOR_ERROR)

    if message is not None:
        print("{0}".format(",".join(message)), file=sys.stderr)

    if not status:
        sys.exit(1)
    else:
        sys.exit(0)